from PIL import Image
import gradio as gr
import subprocess
import shutil, os
from gradio_imageslider import ImageSlider
import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

img_colorization = pipeline(Tasks.image_colorization, model='damo/cv_ddcolor_image-colorization')
result = img_colorization('https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg')
cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
def generate(image):
    image_in = cv2.imread(image)
    img_colorization = pipeline(Tasks.image_colorization, model='damo/cv_ddcolor_image-colorization')
    result = img_colorization(image_in)
    cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
    image_in_pil = Image.fromarray(cv2.cvtColor(image_in, cv2.COLOR_BGR2RGB))
    image_out_pil = Image.fromarray(cv2.cvtColor(image_out, cv2.COLOR_BGR2RGB))
    return (image_in_pil, image_out_pil)

with gr.Blocks() as demo:
    with gr.Row():
      with gr.Column():
          image = gr.Image(type='filepath')
          button = gr.Button()
      output_image = ImageSlider(show_label=False, type="filepath", interactive=False)
    button.click(fn=generate, inputs=[image], outputs=[output_image])

demo.queue().launch(inline=False, share=True, debug=True)