Update image_processor.py

image_processor.py

@@ -1,244 +1,106 @@
-import argparse
-import os
-import shutil
-from rembg import remove
+import gradio as gr
 from PIL import Image
-
-
-def add_background(image_path, background, output_path, default_color="#FFFFFF"):
-    """
-    Adds a background to an image, with a fallback to a default color if the specified color is not available.
-
-    Args:
-    - image_path (str): Path to the input image with transparent background.
-    - background (str): Background color (as a name or hex code) or path to a background image file.
-    - output_path (str): Path where the image with the new background should be saved.
-    - default_color (str): Fallback color if the specified background color is not valid. Defaults to white.
-    """
-    with Image.open(image_path).convert("RGBA") as foreground:
-        try:
-            # Attempt to create a background layer with the specified color or image
-            if background.startswith("#") or background.isalpha():
-                # Check if the color name is valid by creating a small test image
-                Image.new("RGBA", (1, 1), background)
-                background_layer = Image.new(
-                    "RGBA", foreground.size, background)
-            else:
-                # If background is an image file
-                with Image.open(background).convert("RGBA") as bg_img:
-                    bg_img = bg_img.resize(foreground.size)
-                    background_layer = bg_img
-        except ValueError:
-            # If the color is invalid, use the default color
-            print(
-                f"Invalid color '{background}'. Using default color '{default_color}'.")
-            background_layer = Image.new(
-                "RGBA", foreground.size, default_color)
-
-        # Composite the foreground over the background
-        with Image.alpha_composite(background_layer, foreground) as final_img:
-            # Convert to RGB to save in formats other than PNG
-            final_img = final_img.convert("RGB")
-            final_img.save(output_path)
-
-
-def autocrop_image(image_path, output_path):
-    """
-    Autocrops an image, focusing on the non-transparent pixels and saves as PNG.
-
-    Args:
-    - image_path (str): Path to the input image.
-    - output_path (str): Path where the cropped image should be saved.
-    """
-    with Image.open(image_path).convert("RGBA") as image:
-        bbox = image.getbbox()
-        if bbox:
-            cropped_image = image.crop(bbox)
-            cropped_image.save(output_path, format='PNG')
-        else:
-            image.save(output_path, format='PNG')
-
-
-def process_image(input_path, output_path, crop=False, remove_bg=False, resize=None, padding=0, background=None):
-    """
-    Processes a single image based on the provided options and saves it.
-
-    Args:
-    - input_path (str): Path to the input image.
-    - output_path (str): Path where the processed image should be saved.
-    - crop (bool): Whether to autocrop the image.
-    - remove_bg (bool): Whether to remove the background of the image.
-    - resize (tuple): Optional dimensions (width, height) to resize the image.
-    - padding (int): Number of padding pixels to add around the image.
-    - background (str): Optional background color (hex code or name) or path to an image file to set as the background.
-    """
-    need_processing = crop or resize or remove_bg or background
-    temp_path = output_path + ".tmp.png"
-
-    if remove_bg:
-        with open(input_path, 'rb') as input_file:
-            image_data = input_file.read()
-        image_data = remove(image_data)
-        with open(temp_path, 'wb') as temp_file:
-            temp_file.write(image_data)
-    else:
-        # Copy original image to temp_path if no background removal
-        shutil.copy(input_path, temp_path)
-
-    if crop:
-        autocrop_image(temp_path, temp_path)
-
-    if resize:
-        #        adjusted_resize = (resize[0],
-        #                           resize[1]) if padding else resize
-        resize_and_pad_image(temp_path, temp_path, resize, padding)
-
-    if background:
-        add_background(temp_path, background, temp_path)
-
-    # Finalize the process: move from temp_path to output_path
-    os.rename(temp_path, output_path)
-
-
-def resize_and_pad_image(image_path, output_path, dimensions, padding=0):
-    """
-    Resizes an image to fit within specified dimensions (AxB) and adds padding to make it exactly AxB,
-    ensuring the image content is centered within these dimensions.
-
-    Args:
-    - image_path (str): Path to the input image.
-    - output_path (str): Path where the resized and padded image should be saved.
-    - dimensions (tuple): Target dimensions (width, height) in pixels, before adding padding.
-    - padding (int): Number of padding pixels to add around the image.
-    """
-    target_width, target_height = dimensions
-    content_width, content_height = target_width - \
-        2*padding, target_height - 2*padding
-
-    with Image.open(image_path) as img:
-        # Resize the image, preserving the aspect ratio
-        img.thumbnail((content_width, content_height),
-                      Image.Resampling.LANCZOS)
-
-        # Create a new image with the target dimensions and a transparent background
-        # new image shall include padding spacig
-
-        new_img = Image.new("RGBA", dimensions, (255, 255, 255, 0))
-
-        # Calculate the position to paste the resized image to center it
-        paste_position = ((target_width - img.width) // 2,
-                          (target_height - img.height) // 2)
-
-        # Paste the resized image onto the new image, centered
-        new_img.paste(img, paste_position, img if img.mode == 'RGBA' else None)
-
-        # Save the output
-        new_img.save(output_path, format='PNG')
-
-
-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.
+import tempfile
+import os
+from image_processor import process_image
+
+
+def apply_standard_settings(setting):
+    """Returns the parameters for the selected standard setting."""
+    settings_dict = {
+        "S light": (True, True, "240x240", 48, "whitesmoke"),
+        "M light": (True, True, "480x480", 96, "whitesmoke"),
+        "L light": (True, True, "960x960", 128, "whitesmoke"),
+        "S dark": (True, True, "240x240", 48, "#2A373D"),
+        "M dark": (True, True, "480x480", 96, "#2A373D"),
+        "L dark": (True, True, "960x960", 128, "#2A373D"),
+    }
+    # Default to no special settings
+    return settings_dict.get(setting, (None, None, None, None, None))
+
+
+def settings_description(crop, remove_bg, resize, padding, background):
+    """Generate an HTML text description of the current settings in a smaller font and list format."""
+    description = f"""
+    <ul style="font-size:small;">
+        <li>Crop: {crop}</li>
+        <li>Remove Background: {remove_bg}</li>
+        <li>Resize: {resize if resize else 'No resize'}</li>
+        <li>Padding: {padding}</li>
+        <li>Background: {background}</li>
+    </ul>
     """
-    base, _ = os.path.splitext(os.path.basename(input_path))
-    suffix = ""
+    return description
 
-    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"
+def gradio_interface(image, standard_settings, crop=False, remove_bg=False, resize=None, padding=0, background="white"):
+    # Apply standard settings if selected and not "None"
+    if image is None:
+        # Load the standard image from the specified path if no image is uploaded
+        standard_image_path = './data/examples/supermario.png'
+        image = Image.open(standard_image_path)
 
+    if standard_settings and standard_settings != "None":
+        crop, remove_bg, resize, padding, background = apply_standard_settings(
+            standard_settings)
 
-def generate_output_filename2(input_path, remove_bg=False, crop=False, resize=None, padding=0, background=None):
-    """
-    Generates an output filename based on the input path and processing options applied.
-    Takes into account the effect of padding on the final image size for naming.
-    """
-    base, _ = os.path.splitext(os.path.basename(input_path))
-    suffix = ""
+    # Generate settings description
+    applied_settings = settings_description(
+        crop, remove_bg, resize, padding, background)
 
-    if remove_bg:
-        suffix += "_b"
-    if crop:
-        suffix += "_c"
+    # Convert resize string to tuple (if provided)
+    resize_dimensions = None
     if resize:
-        # Adjust the resize dimensions to reflect the final image size after padding
-        if padding > 0:
-            # Adjust dimensions to reflect content size before padding if that's the intent
-            # Otherwise, add padding to the dimensions to reflect final size including padding
-            adjusted_width = resize[0] + 2*padding
-            adjusted_height = resize[1] + 2*padding
-            suffix += f"_{adjusted_width}x{adjusted_height}"
-        else:
-            width, height = resize
-            suffix += f"_{width}x{height}"
-    if background:
-        suffix += "_bg"
-
-    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.
-
-    Args:
-    - input_dir (str): Directory containing the images to be processed.
-    - output_dir (str): Directory where processed images will be saved.
-    - crop (bool): Whether to crop the images.
-    - remove_bg (bool): Whether to remove the background of the images.
-    - resize (tuple): Optional dimensions (width, height) to resize the image.
-    - padding (int): Number of padding pixels to add around the image.
-    - background (str): Optional background color (hex code or name) or path to an image file to set as the background.
-    """
-    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 images are not in the input directory, print a message and return
-    if not inputs:
-        print("No images found in the input directory.")
-        return
-
-    for i, input_path in enumerate(inputs, start=1):
-        filename = os.path.basename(input_path)
-        output_filename = generate_output_filename(
-            input_path, remove_bg=remove_bg, crop=crop, resize=resize, background=background)
-        output_path = os.path.join(output_dir, output_filename)
-        print(f"Processing image {i}/{len(inputs)}...{filename}")
-
-        # Update the call to process_image with all parameters including background
-        process_image(input_path, output_path, crop=crop, remove_bg=remove_bg,
-                      resize=resize, padding=padding, background=background)
-
-        shutil.move(input_path, os.path.join(processed_input_dir, filename))
-
-    print("All images have been processed.")
+        try:
+            width, height = map(int, resize.split('x'))
+            resize_dimensions = (width, height)
+        except ValueError:
+            return "Invalid format for resize dimensions. Please use 'WxH'.", "original", applied_settings
+    # Process the image directly
+    processed_image = process_image(
+        image, crop, remove_bg, resize_dimensions, padding, background)
+
+    # Generate settings description
+    applied_settings = settings_description(
+        crop, remove_bg, resize, padding, background)
+
+    return processed_image, applied_settings
+
+
+example_images = [
+    [os.path.join("data", "examples", "supermario.png"),
+     "S light", True, True, "480x420", 10, "whitesmoke"],
+    [os.path.join("data", "examples",
+                  "depositphotos_520707962-stock-photo-fujifilm-s10-body-black-fujifilm.jpg"), "None", True, True, "480x320", 48, "blue"],
+    [os.path.join("data", "examples", "batman_b_c_320x280_bg.png"),
+     "None", True, True, "360x360", 48, "yellow"],
+
+
+]
+
+# Define the Gradio interface
+interface = gr.Interface(fn=gradio_interface,
+                         inputs=[
+                             gr.components.Image(
+                                 type="pil", label="Input Image"),
+                             gr.components.Radio(choices=[
+                                                 "None", "S light", "M light", "L light", "S dark", "M dark", "L dark"], label="Settings"),
+                             gr.components.Checkbox(label="Crop"),
+                             gr.components.Checkbox(label="Remove Background"),
+                             gr.components.Textbox(
+                                 label="Resize (WxH)", placeholder="Example: 100x100"),
+                             gr.components.Slider(
+                                 minimum=0, maximum=200, label="Padding"),
+                             gr.components.Textbox(
+                                 label="Background", placeholder="Color name or hex code")
+                         ],
+                         outputs=[
+                             gr.components.Image(type="pil"),
+                             gr.components.HTML(label="Applied Settings")
+                         ],
+                         examples=example_images,
+                         title="IMAGER ___ Image Processor",
+                         description="Upload an image and select processing options or choose a standard setting. Supports crop, autoremove background, resize, add padding, and set the background color.",)
+
+if __name__ == "__main__":
+    interface.launch()