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)
| 127 | return [o['host_allocation'] for o in self.outputs] |
| 128 | |
| 129 | def process(self, batch, scales=None, nms_threshold=None): |
| 130 | """ |
| 131 | Execute inference on a batch of images. The images should already be batched and preprocessed, as prepared by |
| 132 | the ImageBatcher class. Memory copying to and from the GPU device will be performed here. |
| 133 | :param batch: A numpy array holding the image batch. |
| 134 | :param scales: The image resize scales for each image in this batch. Default: No scale postprocessing applied. |
| 135 | :return: A nested list for each image in the batch and each detection in the list. |
| 136 | """ |
| 137 | # Run inference |
| 138 | outputs = self.infer(batch) |
| 139 | |
| 140 | # Process the results |
| 141 | nums = outputs[0] |
| 142 | boxes = outputs[1] |
| 143 | scores = outputs[2] |
| 144 | classes = outputs[3] |
| 145 | detections = [] |
| 146 | normalized = (np.max(boxes) < 2.0) |
| 147 | for i in range(self.batch_size): |
| 148 | detections.append([]) |
| 149 | for n in range(int(nums[i])): |
| 150 | scale = self.inputs[0]['shape'][2] if normalized else 1.0 |
| 151 | if scales and i < len(scales): |
| 152 | scale /= scales[i] |
| 153 | if nms_threshold and scores[i][n] < nms_threshold: |
| 154 | continue |
| 155 | detections[i].append( |
| 156 | { |
| 157 | "ymin": boxes[i][n][0] * scale, |
| 158 | "xmin": boxes[i][n][1] * scale, |
| 159 | "ymax": boxes[i][n][2] * scale, |
| 160 | "xmax": boxes[i][n][3] * scale, |
| 161 | "score": scores[i][n], |
| 162 | "class": int(classes[i][n]), |
| 163 | } |
| 164 | ) |
| 165 | return detections |
| 166 | |
| 167 | |
| 168 | def main(args): |