Object Detection
Inference sample for exported torchscript object detector model.
Mask R-CNN
FBNetV3
sample.py
1
import torch
2
import numpy as np
3
from PIL import Image
4
import torchvision
5
import json
6
import matplotlib.pyplot as plt
7
import cv2
8
9
with open('class_mapping.json') as data:
10
mappings = json.load(data)
11
12
class_mapping = {item['model_idx']: item['class_name'] for item in mappings}
13
14
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
15
16
model = torch.jit.load('model.pt').to(device)
17
18
image_path = '/path/to/your/image'
19
image = Image.open(image_path)
20
# Transform your image according to the transforms.json as in
21
# https://help.hasty.ai/model-playground/image-transformations
22
image = np.array(image)
23
# Convert to torch tensor
24
x = torch.from_numpy(image).to(device)
25
with torch.no_grad():
26
# Convert to channels first, convert to float datatype
27
x = x.permute(2, 0, 1).float()
28
y = model(x)
29
30
# Some optional postprocessing, you can change the 0.5 iou
31
# overlap as needed
32
to_keep = torchvision.ops.nms(y['pred_boxes'], y['scores'], 0.5)
33
y['pred_boxes'] = y['pred_boxes'][to_keep]
34
y['pred_classes'] = y['pred_classes'][to_keep]
35
36
# Draw you box predictions:
37
for bbox, label in zip(y['pred_boxes'], y['pred_classes']):
38
bbox = list(map(int, bbox))
39
x1, y1, x2, y2 = bbox
40
class_idx = label.item()
41
class_name = class_mapping[class_idx]
42
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 4)
43
cv2.putText(
44
image,
45
class_name,
46
(x1, y1),
47
cv2.FONT_HERSHEY_SIMPLEX,
48
4,
49
(255, 0, 0)
50
)
51
# Display predicted boxes and classes on your image
52
plt.imshow(image)
53
plt.show()
Copied!
sample.py
1
import torch
2
import numpy as np
3
from PIL import Image
4
import torchvision
5
import json
6
import matplotlib.pyplot as plt
7
import cv2
8
9
with open('class_mapping.json') as data:
10
mappings = json.load(data)
11
12
class_mapping = {item['model_idx']: item['class_name'] for item in mappings}
13
14
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
15
16
model = torch.jit.load('model.pt').to(device)
17
18
image_path = '/path/to/your/image'
19
image = Image.open(image_path)
20
# Transform your image according to the transforms.json as in
21
# https://help.hasty.ai/model-playground/image-transformations
22
image = np.array(image)
23
# Convert to torch tensor
24
x = torch.from_numpy(image).to(device)
25
with torch.no_grad():
26
# Convert to channels first, convert to float datatype
27
x = x.permute(2, 0, 1).float()
28
pred_boxes, pred_classes, scores, _ = model(x)
29
30
# Some optional postprocessing, you can change the 0.5 iou
31
# overlap as needed
32
to_keep = torchvision.ops.nms(pred_boxes, scores, 0.5)
33
pred_boxes = pred_boxes[to_keep]
34
pred_classes = pred_classes[to_keep]
35
36
# Draw you box predictions:
37
for bbox, label in zip(pred_boxes, pred_classes):
38
bbox = list(map(int, bbox))
39
x1, y1, x2, y2 = bbox
40
class_idx = label.item()
41
class_name = class_mapping[class_idx]
42
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 4)
43
cv2.putText(
44
image,
45
class_name,
46
(x1, y1),
47
cv2.FONT_HERSHEY_SIMPLEX,
48
4,
49
(255, 0, 0)
50
)
51
# Display predicted boxes and classes on your image
52
plt.imshow(image)
53
plt.show()
Copied!
The script above should produce outputs that look like this:
Copy link