diff --git a/preprocessing/sports/space_data/ultimate/ultimate_space_class.py b/preprocessing/sports/space_data/ultimate/ultimate_space_class.py index fcd7a8c..5f8b031 100644 --- a/preprocessing/sports/space_data/ultimate/ultimate_space_class.py +++ b/preprocessing/sports/space_data/ultimate/ultimate_space_class.py @@ -95,3 +95,84 @@ def preprocessing(self): ) return event_data_dict, home_tracking_dict, away_tracking_dict + + def detect_initiations( + self, + velocity_threshold=3.0, + acceleration_threshold=4.0, + distance_threshold=5.0, + player_threshold=2, + ): + """ + Detect play initiations in tracking data. + + This method processes raw tracking data (2_1.csv or 1_1_1.csv format), + calculates required motion features, detects play initiations, and optionally + extracts individual play segments. + + Args: + velocity_threshold (float): Threshold for velocity (default: 3.0). + acceleration_threshold (float): Threshold for acceleration (default: 4.0). + distance_threshold (float): Threshold for distance (default: 5.0). + player_threshold (int): Threshold for player proximity (default: 2). + output_detected_dir (str): Optional directory to save detected play data. + output_extracted_dir (str): Optional directory to save extracted play segments. + + Returns: + tuple: (detected_plays_dict, extracted_plays_dict) + - detected_plays_dict: Dictionary mapping match_id to detected play DataFrame + - extracted_plays_dict: Dictionary mapping (match_id, offense_id, play_num) to extracted play DataFrame + """ + tracking_files = self.get_files() + if self.testing_mode: + tracking_files = tracking_files[:2] + print("Running in testing mode. Limited files will be processed.") + + from .ultimate_space_detect_initiation import detect_play, extract_play + from .ultimate_space_preprocessing import create_intermediate_file + + detected_plays_dict = {} + for tracking_path_i in tqdm( + tracking_files, total=len(tracking_files), desc="Processing tracking files" + ): + match_i = os.path.splitext( + os.path.splitext(os.path.basename(tracking_path_i))[0] + )[0] + + # Read raw tracking data + match_tracking_df = pd.read_csv(tracking_path_i) + + # Prepare data with required columns + intermediate_df = create_intermediate_file(match_tracking_df) + + # Detect play initiations + detected_df = detect_play( + intermediate_df, + velocity_threshold=velocity_threshold, + acceleration_threshold=acceleration_threshold, + distance_threshold=distance_threshold, + player_threshold=player_threshold, + ) + + # Extract individual play segments + + plays = extract_play(detected_df) + + for play_key, play_df in plays.items(): + detected_plays_dict[len(detected_plays_dict)] = ( + f"{match_i}-id{play_key[1]}-play{play_key[2]}" + ) + + if self.out_path: + os.makedirs(self.out_path + "/initiation/plays", exist_ok=True) + + for play_key, df in plays.items(): + df.to_csv( + os.path.join( + self.out_path, + "initiation/plays", + f"{match_i}-id{play_key[1]}-play{play_key[2]}.csv", + ), + index=False, + ) + return detected_plays_dict diff --git a/pyproject.toml b/pyproject.toml index e99530b..09a8699 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" [project] name = "openstarlab_preprocessing" -version = "0.1.54" +version = "0.1.55" description = "openstarlab preprocessing package" readme = "README.md" requires-python = ">=3.8"