from fastai.vision.all import *
import gradio as gr
import cloudpickle
import numpy as np
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")


# Cargamos el learner
def load(f, map_location='cpu', pickle_module=pickle, **pickle_load_args):
    with open ('modelomemes.pkl',mode='rb') as file:
        learn=cloudpickle.load(file)
    



# 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)