MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / infer

Method infer

samples/python/detectron2/infer.py:103–156  ·  view source on GitHub ↗

Execute inference on a batch of images. The images should already be batched and preprocessed, as prepared by the ImageBatcher class. Memory copying to and from the GPU device will be performed here. :param batch: A numpy array holding the image batch. :param scales:

(self, batch, scales=None, nms_threshold=None)

Source from the content-addressed store, hash-verified

101 return specs
102
103 def infer(self, batch, scales=None, nms_threshold=None):
104 """
105 Execute inference on a batch of images. The images should already be batched and preprocessed, as prepared by
106 the ImageBatcher class. Memory copying to and from the GPU device will be performed here.
107 :param batch: A numpy array holding the image batch.
108 :param scales: The image resize scales for each image in this batch. Default: No scale postprocessing applied.
109 :return: A nested list for each image in the batch and each detection in the list.
110 """
111
112 # Prepare the output data.
113 outputs = []
114 for shape, dtype in self.output_spec():
115 outputs.append(np.zeros(shape, dtype))
116
117 # Process I/O and execute the network.
118 common.memcpy_host_to_device(self.inputs[0]['allocation'], np.ascontiguousarray(batch))
119
120 self.context.execute_v2(self.allocations)
121 for o in range(len(outputs)):
122 common.memcpy_device_to_host(outputs[o], self.outputs[o]['allocation'])
123
124 # Process the results.
125 nums = outputs[0]
126 boxes = outputs[1]
127 scores = outputs[2]
128 pred_classes = outputs[3]
129 masks = outputs[4]
130
131 detections = []
132 for i in range(self.batch_size):
133 detections.append([])
134 for n in range(int(nums[i])):
135 # Select a mask.
136 mask = masks[i][n]
137
138 # Calculate scaling values for bboxes.
139 scale = self.inputs[0]['shape'][2]
140 scale /= scales[i]
141 scale_y = scale
142 scale_x = scale
143
144 if nms_threshold and scores[i][n] < nms_threshold:
145 continue
146 # Append to detections
147 detections[i].append({
148 'ymin': boxes[i][n][0] * scale_y,
149 'xmin': boxes[i][n][1] * scale_x,
150 'ymax': boxes[i][n][2] * scale_y,
151 'xmax': boxes[i][n][3] * scale_x,
152 'score': scores[i][n],
153 'class': int(pred_classes[i][n]),
154 'mask': mask,
155 })
156 return detections
157
158
159def main(args):

Callers 2

mainFunction · 0.95
mainFunction · 0.95

Calls 2

output_specMethod · 0.95
appendMethod · 0.45

Tested by

no test coverage detected