Create app.py

app.py

@@ -0,0 +1,144 @@
+import os
+import cv2
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageEnhance, ImageFilter
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+
+# Load the colorization model into memory once at startup
+img_colorization = pipeline(
+    Tasks.image_colorization,
+    model="iic/cv_ddcolor_image-colorization"
+)
+
+def colorize_image(img_path: str) -> str:
+    """
+    Reads a B&W image from disk, runs the colorization model, 
+    writes the colorized result to a temp file, and returns its path.
+    """
+    image = cv2.imread(str(img_path))
+    output = img_colorization(image[..., ::-1])
+    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+
+    temp_dir = tempfile.mkdtemp()
+    out_path = os.path.join(temp_dir, "colorized.png")
+    cv2.imwrite(out_path, result)
+    return out_path
+
+def enhance_image(
+    img_path: str,
+    brightness: float = 1.0,
+    contrast: float = 1.0,
+    edge_enhance: bool = False
+) -> str:
+    """
+    Opens a colorized image from disk, applies brightness, contrast, 
+    and optional edge enhancement, saves to a temp file, and returns its path.
+    """
+    image = Image.open(img_path)
+
+    # Adjust brightness
+    image = ImageEnhance.Brightness(image).enhance(brightness)
+    # Adjust contrast
+    image = ImageEnhance.Contrast(image).enhance(contrast)
+    # Optionally apply an edge enhancement filter
+    if edge_enhance:
+        image = image.filter(ImageFilter.EDGE_ENHANCE)
+
+    temp_dir = tempfile.mkdtemp()
+    enhanced_path = os.path.join(temp_dir, "enhanced.png")
+    image.save(enhanced_path)
+    return enhanced_path
+
+def process_image(
+    img_path: str,
+    brightness: float,
+    contrast: float,
+    edge_enhance: bool,
+    output_format: str
+):
+    """
+    1) Colorizes the uploaded B&W image.
+    2) Applies the chosen brightness/contrast/edge-enhancement.
+    3) Re‐saves in the user’s chosen format (PNG/JPEG/TIFF).
+    Returns:
+      - A list [original_path, final_path] for side-by-side display.
+      - The final image’s file path for download.
+    """
+    # Step 1: colorize
+    colorized_path = colorize_image(img_path)
+    # Step 2: enhancement
+    enhanced_path = enhance_image(colorized_path, brightness, contrast, edge_enhance)
+    # Step 3: convert to chosen format
+    img = Image.open(enhanced_path)
+    temp_dir = tempfile.mkdtemp()
+    filename = f"colorized_image.{output_format.lower()}"
+    output_path = os.path.join(temp_dir, filename)
+    img.save(output_path, format=output_format.upper())
+
+    # Return ([original, enhanced], download_path)
+    return ([img_path, enhanced_path], output_path)
+
+# Title and description shown at the top of the interface
+TITLE = "🌈 Color Restorization Model"
+DESCRIPTION = "Upload a black & white photo to restore it in color using a deep learning model."
+
+# Build the Gradio Blocks interface
+with gr.Blocks(title=TITLE) as app:
+    gr.Markdown(f"## {TITLE}")
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(
+                type="filepath",
+                label="Upload B&W Image",
+                tool="editor"       # Enables zoom/pan on the uploaded image
+            )
+            brightness_slider = gr.Slider(
+                minimum=0.5, maximum=2.0, value=1.0,
+                label="Brightness"
+            )
+            contrast_slider = gr.Slider(
+                minimum=0.5, maximum=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_gallery = gr.Gallery(
+                label="Original vs Colorized",
+                columns=2,      # two images side by side
+                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_gallery, download_btn]
+    )
+
+# Launch in “production” mode: bind to 0.0.0.0 on configurable port
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    app.queue().launch(server_name="0.0.0.0", server_port=port)