|
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 |
|
|
|
|
|
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization') |
|
|
|
def inference(img): |
|
image = cv2.imread(str(img)) |
|
output = img_colorization(image[..., ::-1]) |
|
result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8) |
|
|
|
temp_dir = tempfile.mkdtemp() |
|
out_path = os.path.join(temp_dir, 'old-to-color.png') |
|
cv2.imwrite(out_path, result) |
|
return Path(out_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") |
|
submit_btn = gr.Button("Colorize") |
|
with gr.Column(): |
|
output_image = gr.Image(type="pil", label="Colorized Output") |
|
|
|
submit_btn.click(fn=inference, inputs=input_image, outputs=output_image) |
|
|
|
demo.launch(enable_queue=True) |
|
|