from fastai.vision.all import *
import gradio as gr


# Cargamos el learner
learn = load_learner('modelomemes.pkl')


# Definimos una función que se encarga de llevar a cabo las predicciones
def predict(img):
    imgLAB = cv2.cvtColor(img.astype('uint8'), cv2.COLOR_BGR2LAB)
    img_pred,a,b = learn.predict(imgLAB[:,:,0])
    arrL = np.array(img_pred)[0,:,:]
    arrA = np.array(img_pred)[1,:,:]
    arrB = np.array(img_pred)[2,:,:]
    imgP = np.stack((arrB,arrA,arrL),axis=2)
    imgColorRGB = cv2.cvtColor(imgP.astype('uint8'), cv2.COLOR_LAB2BGR)
    return(imgColorRGB)
    
# Creamos la interfaz y la lanzamos. 
gr.Interface(fn=predict, inputs=gr.inputs.Image(), outputs=gr.outputs.Image()).launch(share=False)