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

im_size = (1152, 768)
with open("learn.pkl", 'rb') as f:
    model = pickle.load(f)

def predict_mask(file):
    img = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, im_size)
    pred = model.predict(img)
    return pred[0]

iface = gr.Interface(fn=predict_mask, 
                     inputs=gr.inputs.Image(label="Upload Image"), 
                     outputs="image",
                     capture_session=True,
                     interpretation="default")

if __name__ == "__main__":
    iface.launch()