Update app.py
Browse files
app.py
CHANGED
@@ -4,67 +4,38 @@ import tempfile
|
|
4 |
from modelscope.outputs import OutputKeys
|
5 |
from modelscope.pipelines import pipeline
|
6 |
from modelscope.utils.constant import Tasks
|
7 |
-
|
8 |
import gradio as gr
|
|
|
9 |
|
10 |
-
# Load model
|
11 |
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
|
12 |
|
13 |
def inference(img):
|
14 |
-
"""Process input image and return colorized output"""
|
15 |
image = cv2.imread(str(img))
|
16 |
output = img_colorization(image[..., ::-1])
|
17 |
result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
|
18 |
-
|
19 |
-
# Save result to temporary directory
|
20 |
temp_dir = tempfile.mkdtemp()
|
21 |
-
out_path = os.path.join(temp_dir, '
|
22 |
cv2.imwrite(out_path, result)
|
23 |
return Path(out_path)
|
24 |
|
25 |
-
#
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
|
|
|
|
31 |
with gr.Row():
|
32 |
with gr.Column():
|
33 |
-
|
34 |
-
|
35 |
-
type="filepath",
|
36 |
-
elem_id="input-image"
|
37 |
-
)
|
38 |
-
submit_btn = gr.Button("🎨 Colorize Image", variant="primary")
|
39 |
-
|
40 |
with gr.Column():
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
# Examples section
|
45 |
-
gr.Examples(
|
46 |
-
examples=[
|
47 |
-
["examples/vintage.jpg"],
|
48 |
-
["examples/portrait.png"],
|
49 |
-
["examples/architecture.jpeg"]
|
50 |
-
],
|
51 |
-
inputs=input_img,
|
52 |
-
outputs=output_img,
|
53 |
-
fn=inference,
|
54 |
-
cache_examples=True
|
55 |
-
)
|
56 |
-
|
57 |
-
# Event handlers
|
58 |
-
submit_btn.click(
|
59 |
-
fn=inference,
|
60 |
-
inputs=[input_img],
|
61 |
-
outputs=[output_img]
|
62 |
-
)
|
63 |
-
output_img.change(
|
64 |
-
fn=lambda img: gr.File.update(value=img) if img else None,
|
65 |
-
inputs=[output_img],
|
66 |
-
outputs=[download_btn]
|
67 |
-
)
|
68 |
|
69 |
-
|
70 |
-
demo.launch(enable_queue=True)
|
|
|
4 |
from modelscope.outputs import OutputKeys
|
5 |
from modelscope.pipelines import pipeline
|
6 |
from modelscope.utils.constant import Tasks
|
7 |
+
from pathlib import Path
|
8 |
import gradio as gr
|
9 |
+
import numpy as np
|
10 |
|
11 |
+
# Load the model into memory to make running multiple predictions efficient
|
12 |
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
|
13 |
|
14 |
def inference(img):
|
|
|
15 |
image = cv2.imread(str(img))
|
16 |
output = img_colorization(image[..., ::-1])
|
17 |
result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
|
18 |
+
|
|
|
19 |
temp_dir = tempfile.mkdtemp()
|
20 |
+
out_path = os.path.join(temp_dir, 'old-to-color.png')
|
21 |
cv2.imwrite(out_path, result)
|
22 |
return Path(out_path)
|
23 |
|
24 |
+
# Modernized UI using Gradio 3.9 components
|
25 |
+
title = "🌈 Color Restorization Model"
|
26 |
+
description = "Upload a black & white photo to restore it in color using a deep learning model."
|
27 |
+
|
28 |
+
with gr.Blocks(title=title) as demo:
|
29 |
+
gr.Markdown(f"## {title}")
|
30 |
+
gr.Markdown(description)
|
31 |
+
|
32 |
with gr.Row():
|
33 |
with gr.Column():
|
34 |
+
input_image = gr.Image(type="filepath", label="Upload B&W Image")
|
35 |
+
submit_btn = gr.Button("Colorize")
|
|
|
|
|
|
|
|
|
|
|
36 |
with gr.Column():
|
37 |
+
output_image = gr.Image(type="pil", label="Colorized Output")
|
38 |
+
|
39 |
+
submit_btn.click(fn=inference, inputs=input_image, outputs=output_image)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
40 |
|
41 |
+
demo.launch(enable_queue=True)
|
|