import gradio as gr
from src.Generator import Generator
import cv2
import numpy as np
from src.Generator import Generator
import tensorflow as tf
def image_preprocessing(img):
    img=cv2.resize(img,(256,256))#resize image to 128*128
    img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)#BGR to RGB
    return img/255.0

def output_processing(img):
    kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
    im = cv2.filter2D(img, -1, kernel)
    im = cv2.blur(im,(3,3))
    return im
def generated_image(gen,data):
    if len(data.shape)==3:
        data=np.expand_dims(data,axis=0)
    predict=gen.predict(data)
    return predict


def comic_generator(image_path):
    image=cv2.imread(image_path)
    image=image_preprocessing(image)
    model=Generator(256,256,3)
    model.load_weights(r"./src/saved_model_v2/model.ckpt")

    prediction=generated_image(model,image)
    pred=output_processing(prediction[0])
    #ret_img=cv2.cvtColor(prediction[0],cv2.COLOR_BGR2RGB)
    return prediction[0]


#inputs = gr.inputs.Image(label="Input Image")
#outputs = gr.outputs.Image(label="Output Image")
title = "FACE2COMIC"
description = "Face to Comic Avatar Translation"
examples = [[r"./sampe_images/0.jpg"], [r"./sampe_images/1.jpg"],[r"./sampe_images/12.jpg"],[r"./sampe_images/13.jpg"], [r"./sampe_images/1001.jpg"],[r"./sampe_images/1009.jpg"],[r"./sampe_images/1037.jpg"],[r"./sampe_images/man_052.jpeg"],[r"./sampe_images/woman_0.58.jpeg"]]
gr.Interface(fn=comic_generator, 
             inputs=gr.Image(type="filepath",label="Input",scale=True, width=256, height=256), 
             outputs=gr.Image(type="numpy",label="Output",scale=True, width=256, height=256),
             title=title, 
             description=description, 
             examples=examples,
            # max_output_size=(256,256)
             ).launch(share=False)