app.py

import gradio as gr
from PIL import Image, ImageEnhance, ImageFilter, ImageOps
import numpy as np
import random
import cv2
import os

# --- Core Transformation Logic ---

def apply_transformations(
    image,
    # New parameter for resizing
    resize_percentage,
    # ---
    apply_mosaic_trigger,
    crop_box,
    scale_factor,
    rotation_angle,
    h_flip,
    v_flip,
    shear_x,
    shear_y,
    brightness,
    contrast,
    saturation,
    hue,
    gamma,
    grayscale,
    invert,
    blur_radius,
    sharpen_factor,
    noise_intensity,
    cutout_n_holes,
    cutout_ratio
):
    """
    Applies a series of transformations to an input image.
    """
    # BUG FIX: More robust input handling for gr.ImageEditor or standard np.array
    if image is None:
        return None, False  # Return nothing if no image is present

    # Get the background image from the ImageEditor's dictionary output, or use the array directly
    image_data = image.get("background") if isinstance(image, dict) else image
    
    if image_data is None:
        return None, False # Return nothing if image_data is still None

    img = Image.fromarray(image_data).convert("RGB")

    # --- FEATURE: Resize functionality added ---
    # Applied first to improve performance on subsequent operations.
    if resize_percentage != 100:
        new_width = int(img.width * resize_percentage / 100.0)
        new_height = int(img.height * resize_percentage / 100.0)
        # Ensure dimensions are at least 1x1
        if new_width > 0 and new_height > 0:
            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

    # Mosaic is a special one-shot trigger
    if apply_mosaic_trigger:
        w, h = img.size
        # Ensure dimensions are even for clean quadrants
        w, h = w - (w % 2), h - (h % 2)
        if w > 0 and h > 0:
            img = img.crop((0, 0, w, h))
            cx, cy = w // 2, h // 2
            crops = [img.crop((0, 0, cx, cy)), img.crop((cx, 0, w, cy)), img.crop((0, cy, w, h)), img.crop((cx, cy, w, h))]
            random.shuffle(crops)
            mosaic_img = Image.new('RGB', (w, h))
            mosaic_img.paste(crops[0], (0, 0))
            mosaic_img.paste(crops[1], (cx, 0))
            mosaic_img.paste(crops[2], (0, cy))
            mosaic_img.paste(crops[3], (cx, cy))
            img = mosaic_img

    if crop_box is not None:
        try:
            # Ensure crop coordinates are within image bounds
            x1, y1, x2, y2 = map(int, crop_box)
            if x1 < x2 and y1 < y2:
                img = img.crop((x1, y1, x2, y2))
        except (ValueError, TypeError):
            pass # Ignore if crop box is invalid

    if scale_factor != 1.0: img = img.resize((int(img.width * scale_factor), int(img.height * scale_factor)), Image.Resampling.LANCZOS)
    if shear_x != 0 or shear_y != 0: img = img.transform(img.size, Image.Transform.AFFINE, (1, shear_x, 0, shear_y, 1, 0), Image.Resampling.BICUBIC)
    if rotation_angle != 0: img = img.rotate(rotation_angle, expand=True, fillcolor=(128, 128, 128))
    if h_flip: img = ImageOps.mirror(img)
    if v_flip: img = ImageOps.flip(img)
    if brightness != 1.0: img = ImageEnhance.Brightness(img).enhance(brightness)
    if contrast != 1.0: img = ImageEnhance.Contrast(img).enhance(contrast)
    if saturation != 1.0: img = ImageEnhance.Color(img).enhance(saturation)

    if hue != 0:
        hsv_img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2HSV)
        hsv_img[:, :, 0] = (hsv_img[:, :, 0].astype(int) + hue) % 180
        img = Image.fromarray(cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB))

    if gamma != 1.0 and gamma > 0:
        inv_gamma = 1.0 / gamma
        table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
        # For RGB images, PIL's point operation applies the LUT to each channel
        img = Image.eval(img, lambda p: table[p])


    if sharpen_factor > 0:
        for _ in range(int(sharpen_factor)): img = img.filter(ImageFilter.SHARPEN)
    if blur_radius > 0: img = img.filter(ImageFilter.GaussianBlur(radius=blur_radius))

    if cutout_n_holes > 0 and cutout_ratio > 0:
        np_img = np.array(img).copy() # Use copy to avoid overwriting original array if needed elsewhere
        h, w, _ = np_img.shape
        hole_w, hole_h = int(w * cutout_ratio), int(h * cutout_ratio)
        if hole_w > 0 and hole_h > 0:
            for _ in range(cutout_n_holes):
                if h > hole_h and w > hole_w:
                    y1, x1 = random.randint(0, h - hole_h), random.randint(0, w - hole_w)
                    # BUG FIX: Fill with a consistent gray color instead of a calculated mean
                    np_img[y1:y1+hole_h, x1:x1+hole_w] = 128
            img = Image.fromarray(np_img)

    if noise_intensity > 0:
        np_img = np.array(img)
        noise = np.random.normal(0, noise_intensity, np_img.shape)
        img = Image.fromarray(np.clip(np_img + noise, 0, 255).astype(np.uint8))

    if grayscale: img = ImageOps.grayscale(img)
    # BUG FIX: Simplified invert logic since image is already converted to RGB at the start
    if invert: img = ImageOps.invert(img)

    # The second return value resets the mosaic_trigger state to False after one use
    return img, False

# --- UI Helper Functions ---
def process_selection(evt: gr.SelectData):
    return (evt.index[0], evt.index[1], evt.index[2], evt.index[3])

def update_slider(min_val, max_val, current_val):
    if min_val > max_val: min_val = max_val
    new_val = max(min_val, min(max_val, current_val))
    return gr.update(minimum=min_val, maximum=max_val, value=new_val)

def reset_all_controls():
    # Added reset value for the new resize slider (100)
    return (
        100, False, None, 1.0, 0, False, False, 0.0, 0.0, 1.0, 1.0, 1.0, 0, 1.0, False, False, 0.0, 0, 0, 0, 0.0,
        0.1, 3.0, -180, 180, -0.5, 0.5, -0.5, 0.5, 0.0, 3.0, 0.0, 3.0, 0.0, 3.0,
        -90, 90, 0.2, 2.2, 0.0, 15.0, 0, 50, 0, 50, 0.0, 0.5,
        None, None
    )

def on_upload():
    # Slicing is adjusted to correctly reset components without touching the image inputs
    return reset_all_controls()[:-2] + (None,)

# --- Gradio UI Layout ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Advanced Image Augmentation Tool (Manual Control)")
    gr.Markdown("Set your parameters on the left, then click **Apply Transformations** to see the result.")

    crop_box_state = gr.State(None)
    mosaic_trigger = gr.State(False)

    with gr.Row(variant="panel"):
        with gr.Column(scale=1, min_width=400):
            gr.Markdown("### Control Panel")
            with gr.Accordion("Geometric Transformations", open=True):
                # NEW: Resize Slider
                resize_percent_slider = gr.Slider(1, 200, 100, step=1, label="Resize (%)")
                scale_slider = gr.Slider(0.1, 3.0, 1.0, step=0.05, label="Scale")
                rotation_slider = gr.Slider(-180, 180, 0, step=1, label="Rotation Angle")
                with gr.Row():
                    # BUG FIX: Buttons now set the angle directly for predictable behavior
                    rotate_90_btn = gr.Button("Rotate 90°")
                    rotate_180_btn = gr.Button("Rotate 180°")
                    rotate_270_btn = gr.Button("Rotate 270°")
                shear_x_slider, shear_y_slider = gr.Slider(-0.5, 0.5, 0.0, step=0.01, label="Shear X"), gr.Slider(-0.5, 0.5, 0.0, step=0.01, label="Shear Y")
                h_flip_check, v_flip_check = gr.Checkbox(label="Horizontal Flip"), gr.Checkbox(label="Vertical Flip")

            with gr.Accordion("Color & Tone Adjustments", open=True):
                brightness_slider, contrast_slider, saturation_slider = gr.Slider(0.0, 3.0, 1.0, step=0.05, label="Brightness"), gr.Slider(0.0, 3.0, 1.0, step=0.05, label="Contrast"), gr.Slider(0.0, 3.0, 1.0, step=0.05, label="Saturation")
                hue_slider, gamma_slider = gr.Slider(-90, 90, 0, step=1, label="Hue"), gr.Slider(0.2, 2.2, 1.0, step=0.05, label="Exposure (Gamma)")
                grayscale_check, invert_check = gr.Checkbox(label="Grayscale"), gr.Checkbox(label="Invert Colors")

            with gr.Accordion("Filters & Distortions", open=True):
                blur_slider, sharpen_slider, noise_slider = gr.Slider(0.0, 15.0, 0.0, step=0.1, label="Blur Radius"), gr.Slider(0, 5, 0, step=1, label="Sharpen Intensity"), gr.Slider(0, 50, 0, step=1, label="Add Noise")

            with gr.Accordion("Enhanced Augmentations", open=False):
                cutout_n_slider, cutout_ratio_slider = gr.Slider(0, 50, 0, step=1, label="Number of Holes"), gr.Slider(0.0, 0.5, 0.0, step=0.01, label="Hole Size Ratio")
                mosaic_btn = gr.Button("Apply Mosaic (then click Apply Transformations)")

            with gr.Accordion("Parameter Range Control (Advanced)", open=False):
                # This section remains the same, no new controls needed here
                with gr.Tabs():
                    with gr.TabItem("Geometric"):
                        scale_min, scale_max, rotation_min, rotation_max = gr.Number(0.1, label="Scale Min"), gr.Number(3.0, label="Scale Max"), gr.Number(-180, label="Rotation Min"), gr.Number(180, label="Rotation Max")
                        shear_x_min, shear_x_max, shear_y_min, shear_y_max = gr.Number(-0.5, label="Shear X Min"), gr.Number(0.5, label="Shear X Max"), gr.Number(-0.5, label="Shear Y Min"), gr.Number(0.5, label="Shear Y Max")
                    with gr.TabItem("Color"):
                        brightness_min, brightness_max, contrast_min, contrast_max = gr.Number(0.0, label="Brightness Min"), gr.Number(3.0, label="Brightness Max"), gr.Number(0.0, label="Contrast Min"), gr.Number(3.0, label="Contrast Max")
                        saturation_min, saturation_max, hue_min, hue_max = gr.Number(0.0, label="Saturation Min"), gr.Number(3.0, label="Saturation Max"), gr.Number(-90, label="Hue Min"), gr.Number(90, label="Hue Max")
                        gamma_min, gamma_max = gr.Number(0.2, label="Exposure Min"), gr.Number(2.2, label="Exposure Max")
                    with gr.TabItem("Filters/Other"):
                        blur_min, blur_max, noise_min, noise_max = gr.Number(0.0, label="Blur Min"), gr.Number(15.0, label="Blur Max"), gr.Number(0, label="Noise Min"), gr.Number(50, label="Noise Max")
                        cutout_n_min, cutout_n_max, cutout_ratio_min, cutout_ratio_max = gr.Number(0, label="Holes Min"), gr.Number(50, label="Holes Max"), gr.Number(0.0, label="Ratio Min"), gr.Number(0.5, label="Ratio Max")

            reset_btn = gr.Button("Reset All Settings", variant="stop", size="lg")

        with gr.Column(scale=3):
            apply_btn = gr.Button("Apply Transformations", variant="primary")
            image_input = gr.ImageEditor(type="numpy", label="Original Image (Select an area to crop)", interactive=True)
            image_output = gr.Image(label="Transformed Image", interactive=False)
            
            # Use a helper function to find example images if they exist
            def find_examples():
                example_files = ["cat.jpg", "cheetah.jpg", "lion.jpg"]
                existing_examples = []
                # __file__ might not work in all environments (like notebooks), so be defensive
                try:
                    script_dir = os.path.dirname(__file__)
                    for f in example_files:
                        path = os.path.join(script_dir, f)
                        if os.path.exists(path):
                            existing_examples.append(path)
                except NameError:
                    # __file__ is not defined, so skip examples
                    pass
                return existing_examples
                
            example_paths = find_examples()
            if example_paths:
                gr.Examples(examples=example_paths, inputs=image_input, label="Example Images")

    # --- Event Listeners ---
    
    # Updated list of inputs for the main function
    all_inputs_for_transform = [image_input, resize_percent_slider, mosaic_trigger, crop_box_state, scale_slider, rotation_slider, h_flip_check, v_flip_check, shear_x_slider, shear_y_slider, brightness_slider, contrast_slider, saturation_slider, hue_slider, gamma_slider, grayscale_check, invert_check, blur_slider, sharpen_slider, noise_slider, cutout_n_slider, cutout_ratio_slider]
    
    # Updated list of all components that can be reset
    all_resettable_components = [resize_percent_slider, mosaic_trigger, crop_box_state, scale_slider, rotation_slider, h_flip_check, v_flip_check, shear_x_slider, shear_y_slider, brightness_slider, contrast_slider, saturation_slider, hue_slider, gamma_slider, grayscale_check, invert_check, blur_slider, sharpen_slider, noise_slider, cutout_n_slider, cutout_ratio_slider, scale_min, scale_max, rotation_min, rotation_max, shear_x_min, shear_x_max, shear_y_min, shear_y_max, brightness_min, brightness_max, contrast_min, contrast_max, saturation_min, saturation_max, hue_min, hue_max, gamma_min, gamma_max, blur_min, blur_max, noise_min, noise_max, cutout_n_min, cutout_n_max, cutout_ratio_min, cutout_ratio_max, image_input, image_output]
    
    # A subset used for upload/clear events
    partial_resettable_components = all_resettable_components[:-2] + [image_output]

    apply_btn.click(fn=apply_transformations, inputs=all_inputs_for_transform, outputs=[image_output, mosaic_trigger], api_name="predict")

    image_input.select(fn=process_selection, inputs=None, outputs=[crop_box_state], show_progress="hidden")
    image_input.upload(fn=on_upload, inputs=None, outputs=partial_resettable_components)
    image_input.clear(fn=reset_all_controls, inputs=None, outputs=all_resettable_components)
    reset_btn.click(fn=reset_all_controls, inputs=None, outputs=all_resettable_components)

    # BUG FIX: Buttons now set a specific value, not increment
    rotate_90_btn.click(lambda: 90, inputs=None, outputs=[rotation_slider])
    rotate_180_btn.click(lambda: 180, inputs=None, outputs=[rotation_slider])
    rotate_270_btn.click(lambda: -90, inputs=None, outputs=[rotation_slider])
    mosaic_btn.click(lambda: True, None, mosaic_trigger)

    range_map = {(scale_min, scale_max): scale_slider, (rotation_min, rotation_max): rotation_slider, (shear_x_min, shear_x_max): shear_x_slider, (shear_y_min, shear_y_max): shear_y_slider, (brightness_min, brightness_max): brightness_slider, (contrast_min, contrast_max): contrast_slider, (saturation_min, saturation_max): saturation_slider, (hue_min, hue_max): hue_slider, (gamma_min, gamma_max): gamma_slider, (blur_min, blur_max): blur_slider, (noise_min, noise_max): noise_slider, (cutout_n_min, cutout_n_max): cutout_n_slider, (cutout_ratio_min, cutout_ratio_max): cutout_ratio_slider}
    for (min_comp, max_comp), slider in range_map.items():
        min_comp.change(fn=update_slider, inputs=[min_comp, max_comp, slider], outputs=[slider])
        max_comp.change(fn=update_slider, inputs=[min_comp, max_comp, slider], outputs=[slider])

demo.launch(debug=True)