Update app.py

app.py

@@ -9,7 +9,7 @@ import gradio as gr
 import numpy as np
 from PIL import Image, ImageEnhance, ImageFilter
 
-# Load the model into memory to make running multiple predictions efficient
+# Load the model once at startup for efficiency
 img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
 
 def colorize_image(img_path):
@@ -23,10 +23,11 @@ def colorize_image(img_path):
 
 def enhance_image(img_path, brightness=1.0, contrast=1.0, edge_enhance=False):
     image = Image.open(img_path)
-    enhancer_brightness = ImageEnhance.Brightness(image)
-    image = enhancer_brightness.enhance(brightness)
-    enhancer_contrast = ImageEnhance.Contrast(image)
-    image = enhancer_contrast.enhance(contrast)
+    # Adjust brightness
+    image = ImageEnhance.Brightness(image).enhance(brightness)
+    # Adjust contrast
+    image = ImageEnhance.Contrast(image).enhance(contrast)
+    # Optionally apply edge enhancement
     if edge_enhance:
         image = image.filter(ImageFilter.EDGE_ENHANCE)
     temp_dir = tempfile.mkdtemp()
@@ -35,15 +36,18 @@ def enhance_image(img_path, brightness=1.0, contrast=1.0, edge_enhance=False):
     return enhanced_path
 
 def process_image(img_path, brightness, contrast, edge_enhance, output_format):
+    # Step 1: colorize
     colorized_path = colorize_image(img_path)
+    # Step 2: enhance (brightness / contrast / edge)
     enhanced_path = enhance_image(colorized_path, brightness, contrast, edge_enhance)
-    # Convert to desired format
-    image = Image.open(enhanced_path)
+    # Step 3: convert to chosen format
+    img = Image.open(enhanced_path)
     temp_dir = tempfile.mkdtemp()
-    output_filename = f'colorized_image.{output_format.lower()}'
-    output_path = os.path.join(temp_dir, output_filename)
-    image.save(output_path, format=output_format.upper())
-    return (img_path, output_path), output_path
+    filename = f'colorized_image.{output_format.lower()}'
+    output_path = os.path.join(temp_dir, filename)
+    img.save(output_path, format=output_format.upper())
+    # Return side-by-side (original, enhanced) plus the single downloadable file
+    return ([img_path, enhanced_path], output_path)
 
 title = "🌈 Color Restorization Model"
 description = "Upload a black & white photo to restore it in color using a deep learning model."
@@ -55,19 +59,27 @@ with gr.Blocks(title=title) as demo:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="filepath", label="Upload B&W Image", tool="editor")
-            brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness", info="Adjust the brightness of the colorized image.")
-            contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast", info="Adjust the contrast of the colorized image.")
-            edge_enhance_checkbox = gr.Checkbox(label="Apply Edge Enhancement", info="Enhance edges to improve image detail.")
-            output_format_dropdown = gr.Dropdown(choices=["PNG", "JPEG", "TIFF"], value="PNG", label="Output Format", info="Select the desired format for the output image.")
+            brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
+            contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
+            edge_enhance_checkbox = gr.Checkbox(label="Apply Edge Enhancement")
+            output_format_dropdown = gr.Dropdown(
+                choices=["PNG", "JPEG", "TIFF"],
+                value="PNG",
+                label="Output Format"
+            )
             submit_btn = gr.Button("Colorize")
         with gr.Column():
-            comparison = gr.ImageSlider(label="Original vs Colorized")
+            comparison_gallery = gr.Gallery(
+                label="Original vs Colorized", 
+                columns=2, 
+                height="auto"
+            )
             download_btn = gr.File(label="Download Colorized Image")
 
     submit_btn.click(
         fn=process_image,
         inputs=[input_image, brightness_slider, contrast_slider, edge_enhance_checkbox, output_format_dropdown],
-        outputs=[comparison, download_btn]
+        outputs=[comparison_gallery, download_btn]
     )
 
 demo.launch(enable_queue=True)