File size: 1,380 Bytes
ca7e923
 
 
021b9b8
ca7e923
 
5d612e7
021b9b8
5d612e7
ca7e923
5d612e7
021b9b8
ca7e923
 
021b9b8
 
 
5d612e7
021b9b8
5d612e7
021b9b8
 
ca7e923
5d612e7
 
 
 
 
 
 
 
ca7e923
021b9b8
5d612e7
 
021b9b8
5d612e7
 
 
ca7e923
5d612e7
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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

# Load the model into memory to make running multiple predictions efficient
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)

# Modernized UI using Gradio 3.9 components
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)