Dev-Hattori/Human_Segmentation

1import gradio as gr2import torch3import numpy as np4import cv25 6model = torch.load('best_model.pt')7 8def predict(img: torch.Tensor):9    image = cv2.resize(img, (320, 320), interpolation = cv2.INTER_CUBIC)10    image = np.transpose(image, (2, 0, 1))11    image = torch.Tensor(image).unsqueeze(0)12    13    pred = model(image)14    prob_mask = torch.sigmoid(pred)15    pred_mask = (prob_mask > 0.5)*1.016    return pred_mask.squeeze(0)17 18 19iface = gr.Interface(fn=predict, inputs="image", outputs="image")20iface.launch()