Merge remote-tracking branch 'origin/main'

TrackingLaboratory · Feb 19, 2024 · 5544a37 · 5544a37
2 parents ce43860 + a96d7dd
commit 5544a37
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/tracklab/datastruct/tracker_state.py b/tracklab/datastruct/tracker_state.py
@@ -296,6 +296,7 @@ def load(self):
         video_detections = pd.DataFrame()
         video_image_preds = self.image_metadatas[self.image_metadatas.video_id == self.video_id]
         if self.load_from_groundtruth:
+            video_detections = self.detections_pred[self.detections_pred.video_id == self.video_id]
             video_image_preds = self.image_pred[self.image_pred.video_id == self.video_id]
         if self.load_file is not None:
             if f"{self.video_id}.pkl" in self.zf["load"].namelist():

diff --git a/tracklab/wrappers/datasets/soccernet/soccernet_game_state.py b/tracklab/wrappers/datasets/soccernet/soccernet_game_state.py
@@ -206,7 +206,7 @@ def video_dir_to_dfs(args):
     video_level_categories = []
     video_folder_path = os.path.join(dataset_path, video_folder)
     if os.path.isdir(video_folder_path):
-        if split == "challenge":
+        if not (Path(video_folder_path) / "Labels-GameState.json").exists():
             img_folder_path = os.path.join(video_folder_path, 'img1')
             video_id = str(int(video_folder.split('-')[-1]))
             video_metadata = {
@@ -320,7 +320,7 @@ def load_set(dataset_path, nvid=-1, vids_filter_set=None):
             annotations_pitch_camera_list.append(result["annotations_pitch_camera"])
             categories_list += result["video_level_categories"]
 
-    if split == "challenge":
+    if len(categories_list) == 0:
         video_metadata = pd.DataFrame(video_metadatas_list)
         image_metadata = pd.concat(image_metadata_list, ignore_index=True)
         detections = None