Create app.py

app.py

@@ -0,0 +1,168 @@
+import os
+import warnings
+import logging
+
+# ----------------------------
+# 1. Warning & logging setup
+# ----------------------------
+# Suppress FutureWarning from timm internals
+warnings.filterwarnings(
+    "ignore",
+    category=FutureWarning,
+    module="timm.models.layers"
+)
+# Suppress UserWarning from modelscope (e.g. missing preprocessor config)
+warnings.filterwarnings(
+    "ignore",
+    category=UserWarning,
+    module="modelscope"
+)
+# Only show ERROR+ logs from modelscope
+logging.getLogger("modelscope").setLevel(logging.ERROR)
+
+# ----------------------------
+# 2. Standard imports
+# ----------------------------
+import cv2
+import tempfile
+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
+
+# ----------------------------
+# 3. Load your colorization model
+# ----------------------------
+img_colorization = pipeline(
+    Tasks.image_colorization,
+    model="iic/cv_ddcolor_image-colorization",
+    model_revision="v1.02",          # explicitly specify revision
+)
+
+# ----------------------------
+# 4. Image processing fns
+# ----------------------------
+def colorize_image(img_path: str) -> str:
+    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:
+    image = Image.open(img_path)
+    image = ImageEnhance.Brightness(image).enhance(brightness)
+    image = ImageEnhance.Contrast(image).enhance(contrast)
+    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
+):
+    # Colorize → Enhance → Re‑save in chosen format
+    colorized_path = colorize_image(img_path)
+    enhanced_path = enhance_image(colorized_path, brightness, contrast, edge_enhance)
+
+    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 side-by-side gallery and downloadable file
+    return ([img_path, enhanced_path], output_path)
+
+# ----------------------------
+# 5. Gradio UI + custom CSS
+# ----------------------------
+custom_css = """
+body { background-color: #f0f2f5; }
+.gradio-container { max-width: 900px !important; margin: auto !important; }
+#header { background-color: #4CAF50; padding: 20px; border-radius: 8px;
+         text-align: center; margin-bottom: 20px; }
+#header h2, #header p { color: white; margin: 0; }
+#header p { margin-top: 5px; font-size: 1rem; }
+#control-panel { background: white; padding: 20px; border-radius: 8px;
+                 box-shadow: 0 2px 8px rgba(0,0,0,0.1); margin-bottom: 20px; }
+#submit-btn { background-color: #4CAF50 !important; color: white !important;
+              border-radius: 8px !important; font-weight: bold;
+              padding: 10px 20px !important; margin-top: 10px !important; }
+#control-panel .gr-row { gap: 15px; }
+.gr-slider, .gr-checkbox, .gr-dropdown { margin-top: 10px; }
+#comparison_gallery { background: white; padding: 10px;
+                      border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); }
+#download-btn { margin-top: 15px !important; }
+"""
+
+TITLE = "🌈 Color Restorization Model"
+DESCRIPTION = "Bring your old black & white photos back to life—upload, adjust, and download in vivid color."
+
+with gr.Blocks(title=TITLE, css=custom_css) as app:
+    # Header
+    gr.HTML(
+        """
+        <div id="header">
+          <h2>🌈 Color Restorization Model</h2>
+          <p>Bring your old black & white photos back to life—upload, adjust, and download in vivid color.</p>
+        </div>
+        """
+    )
+
+    # Controls & results
+    with gr.Column(elem_id="control-panel"):
+        with gr.Row():
+            # Inputs
+            with gr.Column():
+                input_image = gr.Image(type="filepath", label="Upload B&W Image", interactive=True)
+                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(["PNG", "JPEG", "TIFF"], value="PNG", label="Output Format")
+                submit_btn = gr.Button("Colorize", elem_id="submit-btn")
+
+            # Outputs
+            with gr.Column():
+                comparison_gallery = gr.Gallery(
+                    label="Original vs. Colorized",
+                    columns=2,
+                    elem_id="comparison_gallery",
+                    height="auto"
+                )
+                download_btn = gr.File(label="Download Colorized Image", elem_id="download-btn")
+
+    # Wire up
+    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
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    app.queue().launch(server_name="0.0.0.0", server_port=port)