Semantic Segmentation

Sample inference script for torchscript exported semantic segmentation model

import torch
import numpy as np
from PIL import Image
import torchvision
import json
import matplotlib.pyplot as plt
import cv2

with open('class_mapping.json') as data:
    mappings = json.load(data)

class_mapping = {item['model_idx']: item['class_name'] for item in mappings}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = torch.jit.load('').to(device)

image_path = '/path/to/your/image'
image =
# Transform your image if the config.yaml shows
# you used any image transforms for validation data
image = np.array(image)
h, w = image.shape[:2]
# Convert to torch tensor
x = torch.from_numpy(image).to(device)
with torch.no_grad():
    # Convert to channels first, convert to float datatype
    x = x.permute(2, 0, 1).unsqueeze(dim=0).float()
    y = model(x)
    mask = torch.argmax(y, dim=1).squeeze()

# Overlay predicted mask on image and display
plt.imshow(mask, alpha=0.5)

The script above should produce outputs that look like this:

Example output from the semseg inference script, yellow highlights the present class.
Last updated on Jul 29, 2022

Removing the risk from vision AI.

Only 13% of vision AI projects make it to production, with Hasty we boost that number to 100%.