Spaces:
Sleeping
Sleeping
| 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) |