import matplotlib.pyplot as plt
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 according to the transforms.json as in
# https://help.hasty.ai/model-playground/image-transformations
# Convert to torch tensor
x = torch.from_numpy(image).to(device)
# 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
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)
for mask, bbox, label in zip(reversed(pred_masks),
bbox = list(map(int, bbox))
class_name = class_mapping[class_idx]
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 4)
cv2.FONT_HERSHEY_SIMPLEX,
mask = cv2.resize(mask.squeeze().numpy(), dsize=(w, h),
interpolation=cv2.INTER_LINEAR)
all_masks[mask > 0.5] = instance_idx
# Display predicted masks, boxes and classes on your image
plt.imshow(all_masks, alpha=0.5)