ssd_inference.py
3.95 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
# import the necessary packages
from utils_tools.data_utils import get_dataloader
from utils_tools.data_utils import normalize
from utils_tools import config
from torchvision.datasets import ImageFolder
from torch.utils.data import Subset
from sklearn.model_selection import train_test_split
from torchvision.transforms import Compose
from torchvision.transforms import ToTensor
from torchvision.transforms import Resize
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import random
import torch
import cv2
import os
# initialize test transform pipeline
testTransform = Compose([
Resize((config.IMAGE_SIZE, config.IMAGE_SIZE)), ToTensor()])
# create the test dataset and initialize the test data loader
testDataset = ImageFolder(config.TEST_PATH, testTransform)
testLoader = get_dataloader(testDataset, config.PRED_BATCH_SIZE)
# initialize iterable variable
sweeper = iter(testLoader)
# list to store permuted images
imageInput = []
# grab a batch of test data
print("[INFO] getting the test data...")
batch = next(sweeper)
(images, _) = (batch[0], batch[1])
# switch off autograd
with torch.no_grad():
# send the images to the device
images = images.to(config.DEVICE)
# loop over all the batch
for index in range(0, config.PRED_BATCH_SIZE):
# grab the image, de-normalize it, scale the raw pixel
# intensities to the range [0, 255], and change the channel
# ordering from channels first tp channels last
image = images[index]
image = image.permute((1, 2, 0))
imageInput.append(image.cpu().detach().numpy())
# call the required entry points
ssdModel = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd')
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd_processing_utils')
# flash model to the device and set it to eval mode
ssdModel.to(config.DEVICE)
ssdModel.eval()
# new list for processed input
processedInput = []
# loop over images and preprocess them
for image in imageInput:
image = normalize(image)
processedInput.append(image)
# convert the preprocessed images into tensors
inputTensor = utils.prepare_tensor(processedInput)
# turn off auto-grad
print("[INFO] getting detections from the test data...")
with torch.no_grad():
# feed images to model
detections = ssdModel(inputTensor)
# decode the results and filter them using the threshold
resultsPerInput = utils.decode_results(detections)
bestResults = [utils.pick_best(results,
config.THRESHOLD) for results in resultsPerInput]
# get coco labels
classesToLabels = utils.get_coco_object_dictionary()
# loop over the image batch
for image_idx in range(len(bestResults)):
(fig, ax) = plt.subplots(1)
# denormalize the image and plot the image
image = processedInput[image_idx] / 2 + 0.5
ax.imshow(image)
# grab bbox, class, and confidence values
(bboxes, classes, confidences) = bestResults[image_idx]
# loop over the detected bounding boxes
for idx in range(len(bboxes)):
# scale values up according to image size
(left, bot, right, top) = bboxes[idx] * 300
# draw the bounding box on the image
(x, y, w, h) = [val for val in [left, bot, right - left,
top - bot]]
rect = patches.Rectangle((x, y), w, h, linewidth=1,
edgecolor="r", facecolor="none")
ax.add_patch(rect)
ax.text(x, y,
"{} {:.0f}%".format(classesToLabels[classes[idx] - 1],
confidences[idx] * 100),
bbox=dict(facecolor="white", alpha=0.5))
# check to see if the output directory already exists, if not
# make the output directory
if not os.path.exists(config.SSD_OUTPUT):
os.makedirs(config.SSD_OUTPUT)
# save plots to output directory
print("[INFO] saving the inference...")
outputFileName = os.path.join(config.SSD_OUTPUT, "output.png")
plt.savefig(outputFileName)