Version_02.py

# -*- coding: utf-8 -*-
"""Seam Carving.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1qV8wT2ZnlDw0sbDnRNcuOlCcR7-5lSgx
"""

import sys
import time
from typing import List, Tuple

import numpy as np
from numba import jit
from PIL import Image

# Configuration constants
DESIRED_SCALE_FACTOR = 0.5  # Final dimension will be 50% of original
INPUT_IMAGE_PATH = "original.jpg"
OUTPUT_IMAGE_PATH = "resized.jpg"
SEAM_VISUALIZATION_PATH = "seams_visualization.jpg"

#########################
# ENERGY CALCULATION
#########################
def convert_to_grayscale(image_array: np.ndarray) -> np.ndarray:
    """
    Convert an RGB image to grayscale using:
         Gray = 0.299*R + 0.587*G + 0.114*B
    """
    return (0.299 * image_array[:, :, 0] +
            0.587 * image_array[:, :, 1] +
            0.114 * image_array[:, :, 2]).astype(np.float64)

@jit(nopython=True)
def compute_energy_jit(grayscale: np.ndarray) -> np.ndarray:
    """
    Compute the energy map of a grayscale image using:
         e₁ = |∂ₓI| + |∂ᵧI|
    With JIT compilation for performance.
    """
    h, w = grayscale.shape
    energy = np.zeros((h, w), dtype=np.float64)
    # Set borders
    energy[0, :] = 1000
    energy[h-1, :] = 1000
    energy[:, 0] = 1000
    energy[:, w-1] = 1000

    # Compute energy for interior pixels
    for i in range(1, h-1):
        for j in range(1, w-1):
            dx = abs(grayscale[i, j+1] - grayscale[i, j-1])
            dy = abs(grayscale[i+1, j] - grayscale[i-1, j])
            energy[i, j] = dx + dy

    return energy

def compute_energy(grayscale: np.ndarray) -> np.ndarray:
    """
    Wrapper for JIT-compiled energy computation function
    """
    return compute_energy_jit(grayscale)

#########################
# VERTICAL SEAM CARVING
#########################
@jit(nopython=True)
def find_vertical_seam_jit(energy: np.ndarray) -> np.ndarray:
    """
    Find the vertical seam with the minimum cumulative energy using JIT.
    Returns a 1D array of length h; seam[i] is the column to remove at row i.
    """
    h, w = energy.shape
    M = np.copy(energy)
    backtrack = np.zeros((h, w), dtype=np.int32)

    # Build cumulative energy map row by row
    for i in range(1, h):
        for j in range(w):
            # Handle left edge
            if j == 0:
                if M[i-1, j] <= M[i-1, j+1]:
                    min_energy = M[i-1, j]
                    backtrack[i, j] = j
                else:
                    min_energy = M[i-1, j+1]
                    backtrack[i, j] = j+1
            # Handle right edge
            elif j == w - 1:
                if M[i-1, j-1] <= M[i-1, j]:
                    min_energy = M[i-1, j-1]
                    backtrack[i, j] = j-1
                else:
                    min_energy = M[i-1, j]
                    backtrack[i, j] = j
            # Handle middle columns
            else:
                if M[i-1, j-1] <= M[i-1, j] and M[i-1, j-1] <= M[i-1, j+1]:
                    min_energy = M[i-1, j-1]
                    backtrack[i, j] = j-1
                elif M[i-1, j] <= M[i-1, j-1] and M[i-1, j] <= M[i-1, j+1]:
                    min_energy = M[i-1, j]
                    backtrack[i, j] = j
                else:
                    min_energy = M[i-1, j+1]
                    backtrack[i, j] = j+1

            M[i, j] += min_energy

    # Backtrack to find the seam
    seam = np.zeros(h, dtype=np.int32)
    min_val = M[h-1, 0]
    seam[h-1] = 0

    for j in range(1, w):
        if M[h-1, j] < min_val:
            min_val = M[h-1, j]
            seam[h-1] = j

    for i in range(h-2, -1, -1):
        seam[i] = backtrack[i+1, seam[i+1]]

    return seam

def find_vertical_seam(energy: np.ndarray) -> np.ndarray:
    """
    Wrapper for JIT-compiled seam finding function
    """
    return find_vertical_seam_jit(energy)

@jit(nopython=True)
def remove_vertical_seam_jit(image: np.ndarray, indices_map: np.ndarray, seam: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """
    Remove the vertical seam from the image and update the indices map with JIT.
    """
    h, w, c = image.shape
    new_image = np.zeros((h, w-1, c), dtype=np.uint8)
    new_indices = np.zeros((h, w-1), dtype=np.int32)

    for i in range(h):
        col = seam[i]
        # Copy pixels before the seam
        for j in range(col):
            for k in range(c):
                new_image[i, j, k] = image[i, j, k]
            new_indices[i, j] = indices_map[i, j]

        # Copy pixels after the seam
        for j in range(col, w-1):
            for k in range(c):
                new_image[i, j, k] = image[i, j+1, k]
            new_indices[i, j] = indices_map[i, j+1]

    return new_image, new_indices

def remove_vertical_seam(image: np.ndarray, indices_map: np.ndarray, seam: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """
    Wrapper for JIT-compiled seam removal function
    """
    return remove_vertical_seam_jit(image, indices_map, seam)

def carve_vertical(image: Image.Image, desired_width: int) -> Tuple[np.ndarray, List[List[Tuple[int, int]]]]:
    """
    Iteratively remove vertical seams until the image's width equals desired_width.
    Returns:
      - The carved image as an array.
      - A list of removed seam coordinates (each seam is a list of (row, original_column) tuples).
    """
    start_time = time.time()

    image_np = np.array(image)
    h, w, _ = image_np.shape
    if w <= desired_width:
        raise ValueError("Image width must be greater than desired width.")

    # Initialize indices map to track original positions
    indices_map = np.tile(np.arange(w), (h, 1)).astype(np.int32)
    seams_removed: List[List[Tuple[int, int]]] = []
    current_image = image_np
    current_w = w

    # Progress tracking
    total_seams = w - desired_width

    print(f"Starting vertical carving: removing {total_seams} seams...")

    for i in range(total_seams):
        if i % 10 == 0 or i == total_seams - 1:
            print(f"Processing vertical seam {i+1}/{total_seams} ({(i+1)/total_seams*100:.1f}%)")

        gray = convert_to_grayscale(current_image)
        energy = compute_energy(gray)
        seam = find_vertical_seam(energy)

        # Record seam coordinates in original image space using indices_map
        seam_coords = [(i, int(indices_map[i, seam[i]])) for i in range(h)]
        seams_removed.append(seam_coords)

        current_image, indices_map = remove_vertical_seam(current_image, indices_map, seam)
        current_w -= 1

    elapsed_time = time.time() - start_time
    print(f"Vertical carving completed in {elapsed_time:.2f} seconds")

    return current_image, seams_removed

#########################
# HORIZONTAL SEAM CARVING
#########################
def carve_horizontal(image: Image.Image, desired_height: int) -> Tuple[np.ndarray, List[List[Tuple[int, int]]]]:
    """
    Perform horizontal seam carving by transposing the image, carving vertically,
    then transposing back.
    Returns the carved image and list of removed seam coordinates (in original orientation).
    """
    start_time = time.time()

    np_img = np.array(image)
    original_h = np_img.shape[0]

    # Transpose so that rows become columns
    transposed = np.transpose(np_img, (1, 0, 2))
    transposed_image = Image.fromarray(transposed)
    desired_width = desired_height  # since transposed width equals original height

    # Perform vertical carving on the transposed image
    carved_transposed, seams_removed_transposed = carve_vertical(transposed_image, desired_width)

    # Transpose back to get the carved horizontal image
    carved_image = np.transpose(carved_transposed, (1, 0, 2))

    # Convert seam coordinates back to original orientation:
    # In the transposed image, a coordinate (row, col) corresponds to (col, row) in original
    transformed_seams = []
    for seam in seams_removed_transposed:
        transformed = [(col, row) for row, col in seam]
        transformed_seams.append(transformed)

    elapsed_time = time.time() - start_time
    print(f"Horizontal carving completed in {elapsed_time:.2f} seconds (including vertical carving time)")

    return carved_image, transformed_seams

#########################
# AUTOMATIC DUAL-DIRECTION CARVING
#########################
def carve_image(image: Image.Image, scale_factor: float) -> Tuple[Image.Image, Image.Image]:
    """
    Automatically reduce both width and height based on scale_factor.
    For example, a scale_factor of 0.5 reduces both dimensions to 50% of the original.

    Returns a tuple:
      (resized image, seam visualization image)
    The seam visualization image is the original image with all removed seam pixels marked in red.
    """
    overall_start_time = time.time()

    image_np = np.array(image)
    h, w, _ = image_np.shape
    desired_width = int(w * scale_factor)
    desired_height = int(h * scale_factor)

    print(f"Original dimensions: {w}x{h}")
    print(f"Target dimensions: {desired_width}x{desired_height}")

    # Calculate number of seams to remove in each dimension
    cols_to_remove = w - desired_width
    rows_to_remove = h - desired_height

    # Keep a copy of the original image for seam visualization
    original_image = image_np.copy()

    vertical_seams: List[List[Tuple[int, int]]] = []
    horizontal_seams: List[List[Tuple[int, int]]] = []

    # Remove horizontal seams first (rows)
    carved_image = image_np
    if rows_to_remove > 0:
        carved_image, horizontal_seams = carve_horizontal(Image.fromarray(carved_image), desired_height)

    # Remove vertical seams next (columns)
    if cols_to_remove > 0:
        carved_image, vertical_seams = carve_vertical(Image.fromarray(carved_image), desired_width)

    # Mark seams on the original image
    marking_start_time = time.time()
    print("Generating visualization image...")

    marked_image = original_image.copy()
    # Mark vertical seams
    for seam in vertical_seams:
        for (row, col) in seam:
            if 0 <= row < marked_image.shape[0] and 0 <= col < marked_image.shape[1]:
                marked_image[row, col] = [255, 0, 0]  # red

    # Mark horizontal seams
    for seam in horizontal_seams:
        for (row, col) in seam:
            if 0 <= row < marked_image.shape[0] and 0 <= col < marked_image.shape[1]:
                marked_image[row, col] = [255, 0, 0]  # red

    marking_time = time.time() - marking_start_time
    print(f"Visualization generation completed in {marking_time:.2f} seconds")

    overall_elapsed_time = time.time() - overall_start_time
    print(f"Total execution time: {overall_elapsed_time:.2f} seconds")
    print(f"Final dimensions: {carved_image.shape[1]}x{carved_image.shape[0]}")

    resized_img = Image.fromarray(carved_image)
    vis_img = Image.fromarray(marked_image)
    return resized_img, vis_img

#########################
# MAIN FUNCTION
#########################
def main():
    overall_start_time = time.time()

    # Command-line usage:
    # python seam_carving.py input.jpg resized.jpg seams_vis.jpg
    if len(sys.argv) >= 4:
        input_path = sys.argv[1]
        output_path = sys.argv[2]
        seams_vis_path = sys.argv[3]
    else:
        input_path = INPUT_IMAGE_PATH
        output_path = OUTPUT_IMAGE_PATH
        seams_vis_path = SEAM_VISUALIZATION_PATH

    print(f"Loading image from: {input_path}")
    # Load the image
    input_image = Image.open(input_path).convert("RGB")

    # Warn if dimensions exceed 800x800
    if input_image.width > 800 or input_image.height > 800:
        print("Warning: Image dimensions exceed 800x800. Processing may take longer.")

    # Check that dimensions allow reduction by the scale factor
    if input_image.width < int(input_image.width * DESIRED_SCALE_FACTOR) or input_image.height < int(input_image.height * DESIRED_SCALE_FACTOR):
        raise ValueError(f"Image dimensions must be sufficiently large for reduction by scale factor {DESIRED_SCALE_FACTOR}.")

    # Optionally adjust recursion limit
    sys.setrecursionlimit(max(input_image.height, input_image.width) + 10)

    # Save grayscale and energy visualization for reference
    input_np = np.array(input_image)
    gray_np = convert_to_grayscale(input_np)
    gray_img = Image.fromarray(gray_np.astype(np.uint8))
    gray_img.save("grayscale.jpg", quality=100)

    energy = compute_energy(gray_np)
    max_energy = np.max(energy)
    energy_vis = np.zeros((energy.shape[0], energy.shape[1], 3), dtype=np.uint8)
    for i in range(energy.shape[0]):
        for j in range(energy.shape[1]):
            val = int((energy[i, j] * 255) / max_energy) if max_energy > 0 else 0
            energy_vis[i, j] = (val, val, val)
    energy_img = Image.fromarray(energy_vis)
    energy_img.save("energy.jpg", quality=100)

    # Perform automatic seam carving in both directions
    resized_img, seams_vis_img = carve_image(input_image, DESIRED_SCALE_FACTOR)
    resized_img.save(output_path, quality=100)
    seams_vis_img.save(seams_vis_path, quality=100)

    print("Seam carving completed.")
    print("Resized image saved as:", output_path)
    print("Seam visualization saved as:", seams_vis_path)

    overall_elapsed_time = time.time() - overall_start_time
    print(f"Total execution time: {overall_elapsed_time:.2f} seconds")

if __name__ == "__main__":
    main()