Performs post-processing on the model's output to extract bounding boxes, scores, and class IDs. Args: input_image (numpy.ndarray): The input image. output (numpy.ndarray): The output of the model. Returns: numpy.ndarray: The input image
(self, output, shape_raw, cat_id=[1])
| 109 | return image_data, np.array([img_height, img_width]) |
| 110 | |
| 111 | def postprocess(self, output, shape_raw, cat_id=[1]): |
| 112 | """ |
| 113 | Performs post-processing on the model's output to extract bounding boxes, scores, and class IDs. |
| 114 | |
| 115 | Args: |
| 116 | input_image (numpy.ndarray): The input image. |
| 117 | output (numpy.ndarray): The output of the model. |
| 118 | |
| 119 | Returns: |
| 120 | numpy.ndarray: The input image with detections drawn on it. |
| 121 | """ |
| 122 | # Transpose and squeeze the output to match the expected shape |
| 123 | |
| 124 | outputs = np.squeeze(output) |
| 125 | if len(outputs.shape) == 1: |
| 126 | outputs = outputs[None] |
| 127 | if output.shape[-1] != 6 and output.shape[1] == 84: |
| 128 | outputs = np.transpose(outputs) |
| 129 | |
| 130 | # Get the number of rows in the outputs array |
| 131 | rows = outputs.shape[0] |
| 132 | |
| 133 | # Calculate the scaling factors for the bounding box coordinates |
| 134 | x_factor = shape_raw[1] / self.input_width |
| 135 | y_factor = shape_raw[0] / self.input_height |
| 136 | |
| 137 | # Lists to store the bounding boxes, scores, and class IDs of the detections |
| 138 | boxes = [] |
| 139 | scores = [] |
| 140 | class_ids = [] |
| 141 | |
| 142 | if outputs.shape[-1] == 6: |
| 143 | max_scores = outputs[:, 4] |
| 144 | classid = outputs[:, -1] |
| 145 | |
| 146 | threshold_conf_masks = max_scores >= self.threshold_conf |
| 147 | classid_masks = classid[threshold_conf_masks] != 3.14159 |
| 148 | |
| 149 | max_scores = max_scores[threshold_conf_masks][classid_masks] |
| 150 | classid = classid[threshold_conf_masks][classid_masks] |
| 151 | |
| 152 | boxes = outputs[:, :4][threshold_conf_masks][classid_masks] |
| 153 | boxes[:, [0, 2]] *= x_factor |
| 154 | boxes[:, [1, 3]] *= y_factor |
| 155 | boxes[:, 2] = boxes[:, 2] - boxes[:, 0] |
| 156 | boxes[:, 3] = boxes[:, 3] - boxes[:, 1] |
| 157 | boxes = boxes.astype(np.int32) |
| 158 | |
| 159 | else: |
| 160 | classes_scores = outputs[:, 4:] |
| 161 | max_scores = np.amax(classes_scores, -1) |
| 162 | threshold_conf_masks = max_scores >= self.threshold_conf |
| 163 | |
| 164 | classid = np.argmax(classes_scores[threshold_conf_masks], -1) |
| 165 | |
| 166 | classid_masks = classid != 3.14159 |
| 167 | |
| 168 | classes_scores = classes_scores[threshold_conf_masks][classid_masks] |
no test coverage detected