from random import choices
import numpy as np
import gradio as gr
from glob import glob
from huggingface_hub import from_pretrained_keras
import cv2


model = from_pretrained_keras('Rietta/CycleGAN_DL', compile=False)

imagen = cv2.imread('Kirby.png', cv2.COLOR_BGR2RGB) 

def transform(img, direction):
    img = (img / 127.5) - 1
    if direction==0:
        pred = model.generator_wow.predict(img[None,:,:,:])[0]
    elif direction == 1:
        pred = model.generator_sims.predict(img[None,:,:,:])[0]
    else:
        pred = imagen
    pred = (pred-pred.min())/(pred.max()-pred.min())
    pred = (pred * 255).astype(np.uint8)
    return pred

#examples_gta = [[path, 'GTA->REAL'] for path in glob('Examples/gta*')]
#examples_real = [[path, 'REAL->GTA'] for path in glob('Examples/real*')]
#examples = [*examples_gta, *examples_real]

demo = gr.Interface(fn=transform, 
                    inputs=[gr.inputs.Image(shape=(256, 256), type='numpy'),
                            gr.inputs.Radio(choices=['Sims', 'Warcraft', "Kirby"],
                                            type='index')], 
                    outputs=gr.outputs.Image(type='numpy'))
                    #examples=examples)

if __name__ == '__main__':
    demo.launch()