Spaces:
Sleeping
Sleeping
File size: 2,057 Bytes
ec47fb5 7e014cd ec47fb5 7e014cd ec47fb5 905f281 ec47fb5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import gradio as gr
import torch
from Unet import UNet
import torchvision
from torchvision.transforms import functional as f
import os
from timeit import default_timer as timer
device = 'cpu'
model = UNet(device=device, in_channels=3, num_classes=3)
model.load_state_dict(torch.load("./data/models/Unet_v1.pth", map_location=torch.device('cpu')))
image_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=(128, 128)),
torchvision.transforms.ToTensor()
])
def predict(img):
start_time = timer()
img_transformed = image_transforms(img).to(device)
model.eval()
with torch.inference_mode():
y_logits = model(img_transformed.unsqueeze(dim=0)).squeeze(dim=0)
predicted_label = torch.argmax(y_logits, dim=0).to('cpu')
for i in range(3):
for j in range(128):
for z in range(128):
img_transformed[i][j][z] = predicted_label[j][z]
img_transformed = f.to_pil_image(img_transformed)
return img_transformed, round((timer() - start_time), 3)
title = "Animal Segmentation"
description = "An UNet* feature extractor computer vision model to segment animal in an image.\nModel works more precisely on an image that only contains just one animal."
article = "U-Net: Convolutional Networks for Biomedical Image Segmentation (https://arxiv.org/abs/1505.04597)"
example_list = [["examples/" + example] for example in os.listdir("examples")]
demo = gr.Interface(fn=predict, # mapping function from input to output
inputs=gr.Image(type="pil"), # what are the inputs?
outputs=[gr.Image(label="Segmentation"), # what are the outputs?
gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
# Create examples list from "examples/" directory
examples=example_list,
title=title,
description=description,
article=article)
demo.launch()
|