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)