Spaces:
Running
Running
import gradio as gr | |
from PIL import Image | |
import numpy as np | |
import segmentation_models_pytorch as smp | |
import torch | |
from torchvision import transforms as T | |
from tensorflow.keras.models import load_model | |
model = smp.MAnet( | |
encoder_name="efficientnet-b7", | |
encoder_weights="imagenet", | |
in_channels=3, | |
classes=1, | |
activation='sigmoid',) | |
model.load_state_dict(torch.load("weights.pt", map_location=torch.device('cpu'))) | |
model.eval() | |
def segment(image): | |
image = T.functional.to_tensor(image) | |
prediction = model(image[None, ...]) | |
prediction = np.squeeze(prediction.detach().numpy()) | |
return Image.fromarray(prediction) | |
iface = gr.Interface(fn=segment, inputs="image", outputs="image").launch() |