diff --git a/tracklab/datastruct/tracker_state.py b/tracklab/datastruct/tracker_state.py index 844eae03..d54e40f3 100644 --- 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 index 2a28f72e..bf7da9c8 100644 --- 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