import gradio as gr
import os
import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import PIL
import numpy as np
import uuid
from gradio_imageslider import ImageSlider

img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
img_path = 'input.png'
##result = img_colorization(img_path)
##cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
def color(image):
    output = img_colorization(image[...,::-1])
    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
    # result = result[...,::-1]
    # Generate a unique filename using UUID
    unique_imgfilename = 'colorize_'+str(uuid.uuid4()) + '.png'
    cv2.imwrite(unique_imgfilename, result)
    print('infer finished!')
    return (image, unique_imgfilename) , unique_imgfilename
    

title = "Colorize Black and White Photos"
description = "Use AI image coloring algorithm to easily colorize your old black and white photos instead of traditional color filters."
slider = ImageSlider(position=0.5,label='Original vs AI Processed')
download_button = gr.File(label="Download")
examples = [['./input.jpg'],]

demo = gr.Interface(fn=color,inputs="image",outputs=[slider,download_button],examples=examples,title=title,description='')

if __name__ == "__main__":
    demo.launch(show_error=True)