Fix bug causing labels to be keypoints instead of bboxes in bbox mAP

luxonis · Jun 27, 2024 · ce5ed26 · ce5ed26
1 parent f3d4280
commit ce5ed26
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/luxonis_train/attached_modules/metrics/mean_average_precision.py b/luxonis_train/attached_modules/metrics/mean_average_precision.py
@@ -38,7 +38,9 @@ def update(
     def prepare(
         self, outputs: Packet[Tensor], labels: Labels
     ) -> tuple[list[dict[str, Tensor]], list[dict[str, Tensor]]]:
-        label = labels[self.task][0]
+        label = labels["boundingbox"][
+            0
+        ]  # TODO: Think of a better way to deal with multi-task heads
         output_nms = self.get_input_tensors(outputs)
 
         image_size = self.node.original_in_shape[1:]