VerokeAI/Object_tracking_boxmot

1#!/usr/bin/env python32"""3This script runs a hyperparameter tuning process for a multi-object tracking (MOT) tracker using Ray Tune.4It loads the tracker configuration from a YAML file, sets up the search space for hyperparameters, and evaluates5the tracker to optimize selected metrics (e.g., MOTA, HOTA, IDF1).6"""7 8import os9from pathlib import Path10import yaml11 12# Check required packages13from boxmot.utils.checks import RequirementsChecker14checker = RequirementsChecker()15checker.check_packages(('ray[tune]',))  # Install ray[tune] if not already present16 17import ray18from ray import tune19from ray.air import RunConfig20 21from boxmot.utils.checks import RequirementsChecker22from boxmot.utils import EXAMPLES, TRACKER_CONFIGS, ROOT, NUM_THREADS23from tracking.val import (24    run_generate_dets_embs,25    run_generate_mot_results,26    run_trackeval,27    parse_opt as parse_optt,28    download_mot_eval_tools29)30 31 32class Tracker:33    """34    Encapsulates the evaluation of a tracking configuration.35    """36    def __init__(self, opt):37        self.opt = opt38 39    def objective_function(self, config: dict) -> dict:40        """41        Evaluates a given tracker configuration.42 43        Args:44            config (dict): A dictionary of tracker hyperparameters.45 46        Returns:47            dict: Combined evaluation metrics extracted from run_trackeval.48        """49        # Ensure evaluation tools are available50        download_mot_eval_tools(self.opt.val_tools_path)51        # Generate MOT-compliant results with the specified tracker parameters52        run_generate_mot_results(self.opt, config)53        # Retrieve evaluation metrics (e.g., MOTA, HOTA, IDF1)54        results = run_trackeval(self.opt)55        # Extract only the desired objective results56        combined_results = {key: results.get(key) for key in self.opt.objectives}57        return combined_results58 59 60def load_yaml_config(tracking_method: str) -> dict:61    """62    Loads the YAML configuration file for the given tracking method.63 64    Args:65        tracking_method (str): Name of the tracking method.66 67    Returns:68        dict: Configuration parameters loaded from the YAML file.69    """70    config_path = TRACKER_CONFIGS / f"{tracking_method}.yaml"71    with open(config_path, 'r') as file:72        config = yaml.safe_load(file)73    return config74 75 76def yaml_to_search_space(config: dict) -> dict:77    """78    Converts a YAML configuration dictionary to a Ray Tune search space.79 80    Args:81        config (dict): YAML configuration parameters.82 83    Returns:84        dict: A dictionary representing the search space for hyperparameters.85    """86    search_space = {}87    for param, details in config.items():88        search_type = details.get('type')89        if search_type == 'uniform':90            search_space[param] = tune.uniform(*details['range'])91        elif search_type == 'randint':92            search_space[param] = tune.randint(*details['range'])93        elif search_type == 'qrandint':94            search_space[param] = tune.qrandint(*details['range'])95        elif search_type == 'choice':96            search_space[param] = tune.choice(details['options'])97        elif search_type == 'grid_search':98            search_space[param] = tune.grid_search(details['values'])99        elif search_type == 'loguniform':100            search_space[param] = tune.loguniform(*details['range'])101    return search_space102 103 104def main():105    # Parse options and set necessary paths106    opt = parse_optt()107    opt.val_tools_path = EXAMPLES / 'val_utils'108    opt.source = Path(opt.source).resolve()109    opt.yolo_model = [Path(y).resolve() for y in opt.yolo_model]110    opt.reid_model = [Path(r).resolve() for r in opt.reid_model]111 112    # Load YAML configuration and convert it to a Ray Tune search space113    yaml_config = load_yaml_config(opt.tracking_method)114    search_space = yaml_to_search_space(yaml_config)115 116    # Create a Tracker instance117    tracker = Tracker(opt)118 119    # Generate detection and embedding files required for evaluation120    run_generate_dets_embs(opt)121 122    # Define a wrapper for the objective function for Ray Tune123    def tune_wrapper(config):124        return tracker.objective_function(config)125 126    results_dir = os.path.abspath("ray/")127 128    # Set up and run the hyperparameter tuning using Ray Tune129    tuner = tune.Tuner(130        tune.with_resources(tune_wrapper, {"cpu": NUM_THREADS, "gpu": 0}),131        param_space=search_space,132        tune_config=tune.TuneConfig(num_samples=opt.n_trials),133        run_config=RunConfig(storage_path=results_dir)134    )135    tuner.fit()136 137    # Print the tuning results138    print(tuner.get_results())139 140 141if __name__ == "__main__":142    main()143