| 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") |
|
|
|
|
| |
| 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() |
|
|
| |
| 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) |
| |
| |
| gr.Interface(fn=predict, inputs=gr.inputs.Image(), outputs=gr.outputs.Image()).launch(share=False) |
