Spaces:
Runtime error
Runtime error
File size: 3,359 Bytes
ba20db4 c836824 ba20db4 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import torch
import numpy as np
import gradio as gr
from PIL import Image
from unet import UNet
from torchvision import transforms
from torchvision.transforms.functional import to_tensor, to_pil_image
import matplotlib.pyplot as plt
device = "cuda:0" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
# Load the trained model
model_path = 'cityscapes_dataUNet.pth'
num_classes = 10
model = UNet(num_classes=num_classes)
model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
model.to(device)
model.eval()
# Define the prediction function that takes an input image and returns the segmented image
def predict_segmentation(image):
print(device)
# Convert the input image to a PyTorch tensor and normalize it
image = Image.fromarray(image, 'RGB')
# image = transforms.functional.resize(image, (256, 256))
image = to_tensor(image).unsqueeze(0)
image = transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))(image)
image=image.to(device)
print("input shape",image.shape) # input shape torch.Size([1, 3, 256, 256])
print("input dtype",image.dtype) # input dtype torch.float32
# Make a prediction using the model
with torch.no_grad():
print(image.shape, image.dtype) # torch.Size([1, 3, 256, 256]) torch.float32
output= model(image)
# print(output.shape,output.dtype) # torch.Size([1, 10, 256, 256]) torch.float32
predicted_class = torch.argmax(output, dim=1).squeeze(0)
predicted_class = predicted_class.cpu().detach().numpy().astype(np.uint8)
print(predicted_class.dtype , predicted_class.shape) # int64 (256, 256)
# Visualize the predicted segmentation mask
plt.imshow(predicted_class)
plt.show()
# Apply the inverse transform to convert the normalized image back to RGB
# predicted_class = inverse_transform(torch.from_numpy(predicted_class))
print("predicted class ",predicted_class)
predicted_class = to_pil_image(predicted_class)
# Return the predicted segmentation
return predicted_class
# Define the Gradio interface
input_image = gr.inputs.Image()
output_image = gr.outputs.Image(type='numpy')
gr.Interface(fn=predict_segmentation, inputs=input_image, outputs=output_image,
title='UNet Image Segmentation',
description='Segment an image using a UNet model').launch() |