| import os |
| import cv2 |
| import tempfile |
| from modelscope.outputs import OutputKeys |
| from modelscope.pipelines import pipeline |
| from modelscope.utils.constant import Tasks |
| import PIL |
| 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) |
|
|
|
|
| title = "Color Restorization Model" |
| gr.Interface( |
| inference, |
| [gr.inputs.Image(type="filepath", label="Input")], |
| gr.outputs.Image(type="pil", label="Output"), |
| title=title |
| ).launch(enable_queue=True) |
