Create app.py

app.py

@@ -0,0 +1,88 @@
+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
+
+# Load the model into memory to make running multiple predictions efficient
+img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
+
+def colorize_image(img_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, brightness=1.0, contrast=1.0):
+    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)
+    temp_dir = tempfile.mkdtemp()
+    enhanced_path = os.path.join(temp_dir, 'enhanced.png')
+    image.save(enhanced_path)
+    return enhanced_path
+
+def process_single_image(img_path, brightness, contrast):
+    colorized_path = colorize_image(img_path)
+    enhanced_path = enhance_image(colorized_path, brightness, contrast)
+    return [img_path, enhanced_path], enhanced_path
+
+def process_batch_images(img_paths, brightness, contrast):
+    results = []
+    for img_path in img_paths:
+        colorized_path = colorize_image(img_path)
+        enhanced_path = enhance_image(colorized_path, brightness, contrast)
+        results.append((enhanced_path, os.path.basename(img_path)))
+    return results
+
+title = "🌈 Color Restorization Model"
+description = "Upload black & white photos to restore them in color using a deep learning model."
+
+with gr.Blocks(title=title) as demo:
+    gr.Markdown(f"## {title}")
+    gr.Markdown(description)
+
+    with gr.Tab("Single Image"):
+        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_single_image,
+            inputs=[input_image, brightness_slider, contrast_slider],
+            outputs=[comparison, download_btn]
+        )
+
+    with gr.Tab("Batch Processing"):
+        with gr.Row():
+            with gr.Column():
+                batch_input = gr.File(file_types=["image"], file_count="multiple", label="Upload Multiple B&W Images")
+                batch_brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
+                batch_contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
+                batch_submit_btn = gr.Button("Colorize Batch")
+            with gr.Column():
+                batch_output = gr.Gallery(label="Batch Colorized Images").style(grid=[3], height="auto")
+
+        batch_submit_btn.click(
+            fn=process_batch_images,
+            inputs=[batch_input, batch_brightness_slider, batch_contrast_slider],
+            outputs=batch_output
+        )
+
+demo.launch(enable_queue=True)