|
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 |
|
import requests |
|
from io import BytesIO |
|
from PIL import Image |
|
|
|
|
|
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization') |
|
|
|
|
|
def load_image_from_url(url): |
|
response = requests.get(url) |
|
img = Image.open(BytesIO(response.content)) |
|
return img |
|
|
|
def inference(img, img_url=None): |
|
if img_url: |
|
img = load_image_from_url(img_url) |
|
img = np.array(img) |
|
output = img_colorization(img[..., ::-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) |
|
upload_url = "https://api.postimages.org/upload" |
|
files = {'file': open(out_path, 'rb')} |
|
response = requests.post(upload_url, files=files) |
|
files.close() |
|
image_url = response.json()['url'] |
|
|
|
return Path(out_path), image_url |
|
|
|
title = "Color Restorization Model" |
|
interface = gr.Interface( |
|
inference, |
|
inputs=[ |
|
gr.inputs.Image(type="pil", label="Input Image"), |
|
gr.inputs.Textbox(placeholder="Enter Image URL (optional)", label="Image URL (optional)") |
|
], |
|
outputs=[ |
|
gr.outputs.Image(type="pil", label="Output Image"), |
|
gr.outputs.Textbox(label="Download Link") |
|
], |
|
title=title |
|
) |
|
|
|
interface.launch(enable_queue=True) |
