| 75 | return output |
| 76 | |
| 77 | def process(self, batch, scales=None, nms_threshold=None): |
| 78 | # Infer network |
| 79 | output = self.infer(batch) |
| 80 | |
| 81 | # Extract the results depending on what kind of saved model this is |
| 82 | boxes = None |
| 83 | scores = None |
| 84 | classes = None |
| 85 | if len(self.outputs) == 1: |
| 86 | # Detected as AutoML Saved Model |
| 87 | assert len(self.outputs[0]['shape']) == 3 and self.outputs[0]['shape'][2] == 7 |
| 88 | results = output[self.outputs[0]['name']].numpy() |
| 89 | boxes = results[:, :, 1:5] |
| 90 | scores = results[:, :, 5] |
| 91 | classes = results[:, :, 6].astype(np.int32) |
| 92 | elif len(self.outputs) >= 4: |
| 93 | # Detected as TFOD Saved Model |
| 94 | assert output['num_detections'] |
| 95 | num = int(output['num_detections'].numpy().flatten()[0]) |
| 96 | boxes = output['detection_boxes'].numpy()[:, 0:num, :] |
| 97 | scores = output['detection_scores'].numpy()[:, 0:num] |
| 98 | classes = output['detection_classes'].numpy()[:, 0:num] |
| 99 | |
| 100 | # Process the results |
| 101 | detections = [[]] |
| 102 | normalized = (np.max(boxes) < 2.0) |
| 103 | for n in range(scores.shape[1]): |
| 104 | if scores[0][n] == 0.0: |
| 105 | break |
| 106 | scale = self.inputs[0]['shape'][2] if normalized else 1.0 |
| 107 | if scales: |
| 108 | scale /= scales[0] |
| 109 | if nms_threshold and scores[0][n] < nms_threshold: |
| 110 | continue |
| 111 | detections[0].append({ |
| 112 | 'ymin': boxes[0][n][0] * scale, |
| 113 | 'xmin': boxes[0][n][1] * scale, |
| 114 | 'ymax': boxes[0][n][2] * scale, |
| 115 | 'xmax': boxes[0][n][3] * scale, |
| 116 | 'score': scores[0][n], |
| 117 | 'class': int(classes[0][n]) - 1, |
| 118 | }) |
| 119 | return detections |
| 120 | |
| 121 | |
| 122 | def main(args): |