NEW
All annotation is now free in Hasty.

Instance Segmentation

Sample inference script for torchscript exported instance segmentation model

Mask RCNN code example

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('model.pt').to(device)

image_path = '/path/to/your/image'
image = Image.open(image_path)
\# 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).float()
    y = model(x)
    \# Some optional postprocessing, you can change the 0.5 iou
    \# overlap as needed
    to_keep = torchvision.ops.nms(y['pred_boxes'], y['scores'], 0.5)
    y['pred_boxes'] = y['pred_boxes'][to_keep]
    y['pred_classes'] = y['pred_classes'][to_keep]
    y['pred_masks'] = y['pred_masks'][to_keep]

    \# Draw you box predictions:
    all_masks = np.zeros((h, w), dtype=np.int8)
    instance_idx = 1
    for mask, bbox, label in zip(reversed(y['pred_masks']),
                                 y['pred_boxes'],
                                 y['pred_classes']):
        bbox = list(map(int, bbox))
        x1, y1, x2, y2 = bbox
        class_idx = label.item()
        class_name = class_mapping[class_idx]
        cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 4)
        cv2.putText(
            image,
            class_name,
            (x1, y1),
            cv2.FONT_HERSHEY_SIMPLEX,
            4,
            (255, 0, 0)
        )
        all_masks[mask == 1] = instance_idx
        instance_idx += 1
\# Display predicted masks, boxes and classes on your image
plt.imshow(image)
plt.imshow(all_masks, alpha=0.5)
plt.show()

FBNet v3 code example

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('model.pt').to(device)

image_path = '/path/to/your/image'
image = Image.open(image_path)
\# 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).float()
    pred_boxes, pred_classes, pred_masks, scores, _ = model(x)
    \# Some optional postprocessing, you can change the 0.5 iou
    \# overlap as needed
    to_keep = torchvision.ops.nms(pred_boxes, scores, 0.5)
    pred_boxes = pred_boxes[to_keep]
    pred_classes = pred_classes[to_keep]
    pred_masks = pred_masks[to_keep]

    \# Draw you box predictions:
    all_masks = np.zeros((h, w), dtype=np.int8)
    instance_idx = 1
    for mask, bbox, label in zip(reversed(pred_masks),
                                 pred_boxes,
                                 pred_classes):
        bbox = list(map(int, bbox))
        x1, y1, x2, y2 = bbox
        class_idx = label.item()
        class_name = class_mapping[class_idx]
        cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 4)
        cv2.putText(
            image,
            class_name,
            (x1, y1),
            cv2.FONT_HERSHEY_SIMPLEX,
            4,
            (255, 0, 0)
        )
        mask = cv2.resize(mask.squeeze().numpy(), dsize=(w, h),
                          interpolation=cv2.INTER_LINEAR)
        all_masks[mask > 0.5] = instance_idx
        instance_idx += 1
\# Display predicted masks, boxes and classes on your image
plt.imshow(image)
plt.imshow(all_masks, alpha=0.5)
plt.show()
Example output from the instance segmentation sample inference script
Last updated on Jun 01, 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%.