Running semantic segmentation

Good job designing the U-Net! You will find an already pre-trained model very similar to the one you have just built available to you. This model has been trained on a large set of images and contains a few small additions to the architecture, such as the batch norm layers.

You can instantiate the model as UNet(), which will provide the model with the pre-trained weights. You task is to use it to produce segmentation masks for the following image of a car.

car image

Image from PIL has already been imported for you.

This exercise is part of the course

Deep Learning for Images with PyTorch

View Course

Exercise instructions

Instantiate UNet() in a variable called model.
Load the image at car.jpg to a variable called image.
Produce segmentation masks by passing the image to the model and squeeze(0)-ing the output.

Hands-on interactive exercise

Have a go at this exercise by completing this sample code.

# Load model
model = ____
model.eval()

# Load and transform image
image = ____
transform = transforms.Compose([transforms.ToTensor()])
image_tensor = transform(image).unsqueeze(0)

# Predict segmentation mask
with torch.no_grad():
    prediction = ____

# Display mask
plt.imshow(prediction[1, :, :])
plt.show()

Edit and Run Code