Create app.py

app.py

@@ -0,0 +1,79 @@
+import os
+import cv2
+import tempfile
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from pathlib import Path
+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
+img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
+
+def apply_clahe(image):
+    # Convert to LAB color space
+    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
+    l, a, b = cv2.split(lab)
+    # Apply CLAHE to the L-channel
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+    cl = clahe.apply(l)
+    # Merge the CLAHE enhanced L-channel back with A and B channels
+    limg = cv2.merge((cl,a,b))
+    # Convert back to RGB color space
+    final = cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)
+    return final
+
+def enhance_image_pil(image_pil, brightness=1.0, contrast=1.0):
+    # Apply brightness and contrast enhancements
+    enhancer_brightness = ImageEnhance.Brightness(image_pil)
+    image_pil = enhancer_brightness.enhance(brightness)
+    enhancer_contrast = ImageEnhance.Contrast(image_pil)
+    image_pil = enhancer_contrast.enhance(contrast)
+    # Apply edge enhancement
+    image_pil = image_pil.filter(ImageFilter.EDGE_ENHANCE)
+    return image_pil
+
+def process_image(img_path, brightness, contrast):
+    # Read the input image
+    image = cv2.imread(str(img_path))
+    # Colorize the image
+    output = img_colorization(image[..., ::-1])
+    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+    # Apply CLAHE
+    result_clahe = apply_clahe(result)
+    # Convert to PIL Image for further enhancements
+    result_pil = Image.fromarray(result_clahe)
+    # Apply brightness, contrast, and edge enhancements
+    enhanced_pil = enhance_image_pil(result_pil, brightness, contrast)
+    # Save the enhanced image to a temporary file
+    temp_dir = tempfile.mkdtemp()
+    enhanced_path = os.path.join(temp_dir, 'enhanced.png')
+    enhanced_pil.save(enhanced_path)
+    return [img_path, enhanced_path], enhanced_path
+
+title = "🌈 Color Restorization Model"
+description = "Upload a black & white photo to restore it in color using a deep learning model."
+
+with gr.Blocks(title=title) as demo:
+    gr.Markdown(f"## {title}")
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="filepath", label="Upload B&W 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")
+            submit_btn = gr.Button("Colorize")
+        with gr.Column():
+            comparison = gr.Gallery(label="Original vs Colorized").style(grid=[2], height="auto")
+            download_btn = gr.File(label="Download Colorized Image")
+
+    submit_btn.click(
+        fn=process_image,
+        inputs=[input_image, brightness_slider, contrast_slider],
+        outputs=[comparison, download_btn]
+    )
+
+demo.launch(enable_queue=True)