| 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) |