Update app.py

app.py

@@ -1,64 +1,218 @@
 import os
-import uuid
 import cv2
-import numpy as np
+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
-from modelscope.outputs import OutputKeys
-from gradio_imageslider import ImageSlider
 
-# Initialize the image colorization pipeline with a specified model revision
+# Load the colorization model at startup
 img_colorization = pipeline(
-    task=Tasks.image_colorization,
-    model='iic/cv_ddcolor_image-colorization',
-    model_revision='v1.02'  # Specify the model revision to avoid warnings
+    Tasks.image_colorization,
+    model="iic/cv_ddcolor_image-colorization"
 )
 
-def colorize_image(input_image):
-    """
-    Colorizes a grayscale image using the DDColor model.
+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)
 
-    Args:
-        input_image (numpy.ndarray): Grayscale input image.
+    temp_dir = tempfile.mkdtemp()
+    out_path = os.path.join(temp_dir, "colorized.png")
+    cv2.imwrite(out_path, result)
+    return out_path
 
-    Returns:
-        tuple: Original and colorized images.
-    """
-    # Convert RGB to BGR for OpenCV compatibility
-    bgr_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
+def enhance_image(
+    img_path: str,
+    brightness: float = 1.0,
+    contrast: float = 1.0,
+    edge_enhance: bool = False
+) -> str:
+    image = Image.open(img_path)
 
-    # Perform colorization
-    output = img_colorization(bgr_image)
-    colorized_bgr = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+    # Brightness adjustment
+    image = ImageEnhance.Brightness(image).enhance(brightness)
+    # Contrast adjustment
+    image = ImageEnhance.Contrast(image).enhance(contrast)
+    # Optional edge enhancement
+    if edge_enhance:
+        image = image.filter(ImageFilter.EDGE_ENHANCE)
 
-    # Convert BGR back to RGB for display
-    colorized_rgb = cv2.cvtColor(colorized_bgr, cv2.COLOR_BGR2RGB)
+    temp_dir = tempfile.mkdtemp()
+    enhanced_path = os.path.join(temp_dir, "enhanced.png")
+    image.save(enhanced_path)
+    return enhanced_path
 
-    return input_image, colorized_rgb
+def process_image(
+    img_path: str,
+    brightness: float,
+    contrast: float,
+    edge_enhance: bool,
+    output_format: str
+):
+    # 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)
+    # 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())
 
-# Define the Gradio interface
-title = "DDColor: Photo-Realistic Image Colorization"
-description = (
-    "Upload a black & white image to colorize it using the DDColor model. "
-    "This application utilizes a dual-decoder architecture to produce vivid and natural colorizations."
-)
+    # Return side-by-side (original, enhanced) and the downloadable file
+    return ([img_path, enhanced_path], output_path)
 
-# Example images (ensure these paths are correct or adjust accordingly)
-examples = [['./input.jpg']]
-
-# Create the Gradio interface
-demo = gr.Interface(
-    fn=colorize_image,
-    inputs=gr.Image(label="Upload Grayscale Image"),
-    outputs=ImageSlider(position=0.5, label='Colorization Comparison'),
-    examples=examples,
-    title=title,
-    description=description,
-    allow_flagging="never"
-)
+# CSS to give a modern, centered layout with a colored header and clean panels
+custom_css = """
+/* Overall background */
+body { 
+    background-color: #f0f2f5; 
+}
+
+/* Center the Gradio container and give it a max width */
+.gradio-container {
+    max-width: 900px !important;
+    margin: auto !important;
+}
+
+/* Header styling */
+#header {
+    background-color: #4CAF50;
+    padding: 20px;
+    border-radius: 8px;
+    text-align: center;
+    margin-bottom: 20px;
+}
+#header h2 {
+    color: white;
+    margin: 0;
+    font-size: 2rem;
+}
+#header p {
+    color: white;
+    margin: 5px 0 0 0;
+    font-size: 1rem;
+}
+
+/* White panel for controls */
+#control-panel {
+    background-color: white;
+    padding: 20px;
+    border-radius: 8px;
+    box-shadow: 0px 2px 8px rgba(0,0,0,0.1);
+    margin-bottom: 20px;
+}
+
+/* Style the “Colorize” button */
+#submit-btn {
+    background-color: #4CAF50 !important;
+    color: white !important;
+    border-radius: 8px !important;
+    font-weight: bold;
+    padding: 10px 20px !important;
+    margin-top: 10px !important;
+}
+
+/* Add some spacing around sliders and checkbox */
+#control-panel .gr-row {
+    gap: 15px;
+}
+.gr-slider, .gr-checkbox, .gr-dropdown {
+    margin-top: 10px;
+}
+
+/* Gallery panel styling */
+#comparison_gallery {
+    background-color: white;
+    padding: 10px;
+    border-radius: 8px;
+    box-shadow: 0px 2px 8px rgba(0,0,0,0.1);
+}
+
+/* Download button spacing */
+#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 section
+    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>
+        """
+    )
+
+    # Main control panel: white box with rounded corners
+    with gr.Column(elem_id="control-panel"):
+        with gr.Row():
+            # Left column: inputs and controls
+            with gr.Column():
+                input_image = gr.Image(
+                    type="filepath",
+                    label="Upload B&W Image",
+                    tool="editor",          # Enables zoom/pan
+                    interactive=True
+                )
+                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", 
+                    elem_id="submit-btn"
+                )
+
+            # Right column: results gallery & download
+            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"
+                )
+
+    submit_btn.click(
+        fn=process_image,
+        inputs=[
+            input_image,
+            brightness_slider,
+            contrast_slider,
+            edge_enhance_checkbox,
+            output_format_dropdown
+        ],
+        outputs=[comparison_gallery, download_btn]
+    )
 
+# “Production” launch: bind to 0.0.0.0 and use PORT env var if provided
 if __name__ == "__main__":
-    # Launch the Gradio app with production settings
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
-    
+    port = int(os.environ.get("PORT", 7860))
+    app.queue().launch(server_name="0.0.0.0", server_port=port)