fix when loading image GT columns exclusively

TrackingLaboratory · Mar 5, 2024 · 99c03b7 · 99c03b7
1 parent 4225408
commit 99c03b7
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tracklab/datastruct/tracker_state.py b/tracklab/datastruct/tracker_state.py
@@ -114,9 +114,12 @@ def load_groundtruth(self, load_columns):
                 load_columns = {k: list(set(self.load_from_groundtruth.get(k, [])) & set(v))
                                 for k, v in load_columns.items()
                                 }
-            self.detections_pred_gt = self.detections_gt.copy()[
-                self.detections_gt.columns.intersection(load_columns["detection"]+["image_id", "video_id"])
-            ]
+            if len(load_columns["detection"]) == 0:
+                self.detections_pred_gt = pd.DataFrame(columns=["video_id","image_id"])
+            else:
+                self.detections_pred_gt = self.detections_gt.copy()[
+                    self.detections_gt.columns.intersection(load_columns["detection"]+["image_id", "video_id"])
+                ]
             self.image_pred_gt = merge_dataframes(
                 self.image_metadatas.copy(), self.image_gt.copy()
             )[list(set(load_columns["image"]) | {"video_id", "file_path", "frame"})]