Spaces:
Running
Running
File size: 722 Bytes
44a9d05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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() |