File size: 983 Bytes
7c42e35 4a09b4b 7c42e35 1ecc26f 4a09b4b 7c42e35 1ecc26f 7c42e35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | from fastai.vision.all import *
import gradio as gr
import cloudpickle
import numpy as np
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Cargamos el learner
def load(map_location='cpu', pickle_module=pickle, **pickle_load_args):
with open ('modelomemes.pkl',mode='rb') as file:
learn=cloudpickle.load(file).to(device)
load()
# 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) |