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)