wenmengzhou's picture
use rgb instead of bgr
b9bb4d7 verified
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 = str(uuid.uuid4()) + '.png'
cv2.imwrite(unique_imgfilename, result)
print('infer finished!')
return (image, unique_imgfilename)
title = "old_photo_restoration"
description = "upload old photo, ddcolor image colorization"
examples = [['./input.jpg'],]
demo = gr.Interface(fn=color,inputs="image",outputs=ImageSlider(position=0.5,label='Colored image with slider-view'),examples=examples,title=title,description=description)
if __name__ == "__main__":
demo.launch(share=False)