image_processor.py

import argparse
import os
import shutil
from rembg import remove
from PIL import Image
import io


def add_background(image, background, default_color="#FFFFFF"):
    """
    Adds a background to an image, with a fallback to a default color if the specified background is not available.

    Args:
    - image (PIL.Image.Image): Image with a transparent background.
    - background (str or PIL.Image.Image): Background color (as a hex code) or a PIL Image to be used as background.
    - default_color (str): Fallback color if the specified background is not valid. Defaults to white.

    Returns:
    - PIL.Image.Image: The image with the new background.
    """
    foreground = image.convert("RGBA")

    if isinstance(background, str) and (background.startswith("#") or background.isalpha()):
        # Background is a color
        try:
            Image.new("RGBA", (1, 1), background)  # Test if valid color
            background_layer = Image.new("RGBA", foreground.size, background)
        except ValueError:
            print(
                f"Invalid color '{background}'. Using default color '{default_color}'.")
            background_layer = Image.new(
                "RGBA", foreground.size, default_color)
    elif isinstance(background, Image.Image):
        # Background is an image
        bg_img = background.convert("RGBA")
        background_layer = bg_img.resize(foreground.size)
    else:
        # Fallback to default color
        background_layer = Image.new("RGBA", foreground.size, default_color)

    final_img = Image.alpha_composite(
        background_layer, foreground).convert("RGB")

    return final_img


def cropnontrans(image, padding=0):
    """
    crops a nontransparent image

    Args:
    - image (PIL.Image.Image): Image to be cropped.

    Returns:
    - PIL.Image.Image: The autocropped image.
    """
    # first, 
    # Convert the PIL Image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()
     # Use rembg to remove the background
    result_bytes = remove(img_byte_arr)

    # Convert the result bytes back to a PIL Image
    transparent_image = Image.open(io.BytesIO(result_bytes))
    bbox = transparent_image.getbbox()
    # add padding area to the original bbox
    if bbox:
        bbox = (bbox[0]-padding, bbox[1]-padding, bbox[2]+padding, bbox[3]+padding)
        #delimig the bbox to the image size
        bbox = (max(0, bbox[0]), max(0, bbox[1]), min(transparent_image.width, bbox[2]), min(transparent_image.height, bbox[3]))
        print(f"Bounding box: {bbox}")
        return image.crop(bbox)
    return image



def autocrop_image(image):
    """
    Autocrops an image, focusing on the non-transparent pixels.

    Args:
    - image (PIL.Image.Image): Image to be autocropped.

    Returns:
    - PIL.Image.Image: The autocropped image.
    """
    bbox = image.getbbox()
    print(f"Bounding box: {bbox}")
    if bbox:
        return image.crop(bbox)
    return image


def remove_bg_func(image):
    """
    Removes the background from an image using the rembg library.

    Args:
    - image (PIL.Image.Image): Image object from which to remove the background.

    Returns:
    - PIL.Image.Image: New image object with the background removed.
    """
    # Convert the PIL Image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()

    # Use rembg to remove the background
    result_bytes = remove(img_byte_arr)

    # Convert the result bytes back to a PIL Image
    result_image = Image.open(io.BytesIO(result_bytes))

    return result_image


# Update the process_image function to include format selection
def process_image(img, crop=False, remove_bg=False, resize=None, padding=0, background=None, output_format='webp'):
    """
    Processes a single image with the specified options and format.
    
    Args:
    - img: Input image
    - crop: Whether to autocrop
    - remove_bg: Whether to remove background
    - resize: Tuple of (width, height) for resizing
    - padding: Padding to add
    - background: Background color or image
    - output_format: 'webp', 'png', or 'png-transparent'
    """
    # Existing processing code remains the same
    if remove_bg:
        img = remove_bg_func(img)
    
    if crop and remove_bg:
        img = autocrop_image(img)
    if crop and not remove_bg:
        img = cropnontrans(img, padding)
        
    if resize:
        img = resize_and_pad_image(img, resize, padding)
        
    if background and output_format != 'png-transparent':
        img = add_background(img, background)
    
    return img



def resize_and_pad_image(image, dimensions, padding=0):
    """
    Resizes an image to fit the specified dimensions and adds padding.

    Args:
    - image (PIL.Image.Image): Image object to be resized and padded.
    - dimensions (tuple): Target dimensions (width, height).
    - padding (int): Padding to add around the resized image.

    Returns:
    - PIL.Image.Image: Resized and padded image object.
    """
    target_width, target_height = dimensions
    content_width, content_height = target_width - \
        2*padding, target_height - 2*padding

    # Determine new size, preserving aspect ratio
    img_ratio = image.width / image.height
    target_ratio = content_width / content_height

    if target_ratio > img_ratio:
        new_height = content_height
        new_width = int(new_height * img_ratio)
    else:
        new_width = content_width
        new_height = int(new_width / img_ratio)

    # Resize the image
    resized_img = image.resize(
        (new_width, new_height), Image.Resampling.LANCZOS)

    # Create a new image with the target dimensions and a transparent background
    new_img = Image.new(
        "RGBA", (target_width, target_height), (255, 255, 255, 0))

    # Calculate the position to paste the resized image to center it
    paste_position = ((target_width - new_width) // 2,
                      (target_height - new_height) // 2)

    # Paste the resized image onto the new image, centered
    new_img.paste(resized_img, paste_position,
                  resized_img if resized_img.mode == 'RGBA' else None)

    return new_img


def generate_output_filename(input_path, remove_bg=False, crop=False, resize=None, background=None):
    """
    Generates an output filename based on the input path and processing options applied.
    Appends specific suffixes based on the operations: '_b' for background removal, '_c' for crop,
    and '_bg' if a background is added. It ensures the file extension is '.png'.

    Args:
    - input_path (str): Path to the input image.
    - remove_bg (bool): Indicates if background removal was applied.
    - crop (bool): Indicates if autocrop was applied.
    - resize (tuple): Optional dimensions (width, height) for resizing the image.
    - background (str): Indicates if a background was added (None if not used).

    Returns:
    - (str): Modified filename with appropriate suffix and '.png' extension.
    """
    base, _ = os.path.splitext(os.path.basename(input_path))
    suffix = ""

    if remove_bg:
        suffix += "_b"
    if crop:
        suffix += "_c"
    if resize:
        width, height = resize
        suffix += f"_{width}x{height}"
    if background:
        suffix += "_bg"  # Append "_bg" if the background option was used

    # Ensure the file saves as PNG, accommodating for transparency or added backgrounds
    return f"{base}{suffix}.png"


# The main and process_images functions remain the same, but ensure to update them to handle the new PNG output correctly.

# Update the process_images and main functions to include the new autocrop functionality
# Ensure to pass the crop argument to process_image and adjust the output filename generation accordingly


def process_images(input_dir="./input", output_dir="./output", crop=False, remove_bg=False, resize=None, padding=0, background=None):
    """
    Processes images in the specified directory based on the provided options.
    """
    processed_input_dir = os.path.join(input_dir, "processed")
    os.makedirs(processed_input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    inputs = [os.path.join(input_dir, f) for f in os.listdir(
        input_dir) if os.path.isfile(os.path.join(input_dir, f))]

    if not inputs:
        print("No images found in the input directory.")
        return

    for i, input_path in enumerate(inputs, start=1):
        try:
            with Image.open(input_path) as img:
                # Define filename here, before it's used
                filename = os.path.basename(input_path)

                # Process the image
                processed_img = process_image(
                    img, crop=crop, remove_bg=remove_bg, resize=resize, padding=padding, background=background)

                # Generate output filename based on processing parameters
                output_filename = generate_output_filename(
                    filename, remove_bg=remove_bg, crop=crop, resize=resize, background=background)
                output_path = os.path.join(output_dir, output_filename)

                # Save the processed image to the output directory
                processed_img.save(output_path)

                print(
                    f"Processed image {i}/{len(inputs)}: {filename} -> {output_filename}")

                # Optionally move the processed input image to a "processed" subdirectory
                shutil.move(input_path, os.path.join(
                    processed_input_dir, filename))
        except Exception as e:
            print(f"Error processing image {input_path}: {e}")

    print("All images have been processed.")



def save_image_with_format(image, output_path, format='webp', quality=90, custom_filename=None):
    """
    Saves the image in the specified format with appropriate settings.
    
    Args:
    - image (PIL.Image.Image): The image to save
    - output_path (str): Base path for the output file (without extension)
    - format (str): 'webp', 'png', 'png-transparent', or 'jpg'
    - quality (int): Quality setting for compression (1-100)
    - custom_filename (str): Optional custom filename for the output
    """
    # Get image dimensions for filename
    width, height = image.size
    
    # Generate filename with schema: originalname_size.type
    if custom_filename:
        base_dir = os.path.dirname(output_path)
        filename = f"{custom_filename}_{width}x{height}"
        final_path = os.path.join(base_dir, filename)
    else:
        final_path = output_path

    if format == 'webp':
        final_path = f"{final_path}.webp"
        image.save(final_path, 'webp', quality=quality)
    elif format == 'png-transparent':
        final_path = f"{final_path}.png"
        image.save(final_path, 'PNG', optimize=True)
    elif format == 'png':
        final_path = f"{final_path}.png"
        if image.mode in ('RGBA', 'LA'):
            background = Image.new('RGB', image.size, 'white')
            background.paste(image, mask=image.split()[-1])
            background.save(final_path, 'PNG', optimize=True)
        else:
            image.save(final_path, 'PNG', optimize=True)
    elif format == 'jpg':
        final_path = f"{final_path}.jpg"
        if image.mode in ('RGBA', 'LA'):
            background = Image.new('RGB', image.size, 'white')
            background.paste(image, mask=image.split()[-1])
            background.save(final_path, 'JPEG', quality=quality, optimize=True)
        else:
            image.convert('RGB').save(final_path, 'JPEG', quality=quality, optimize=True)
    else:
        raise ValueError(f"Unsupported format: {format}")
    
    return final_path