File size: 4,660 Bytes
e1832f4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#!/usr/bin/env python3
"""
This script runs a hyperparameter tuning process for a multi-object tracking (MOT) tracker using Ray Tune.
It loads the tracker configuration from a YAML file, sets up the search space for hyperparameters, and evaluates
the tracker to optimize selected metrics (e.g., MOTA, HOTA, IDF1).
"""

import os
from pathlib import Path
import yaml

# Check required packages
from boxmot.utils.checks import RequirementsChecker
checker = RequirementsChecker()
checker.check_packages(('ray[tune]',))  # Install ray[tune] if not already present

import ray
from ray import tune
from ray.air import RunConfig

from boxmot.utils.checks import RequirementsChecker
from boxmot.utils import EXAMPLES, TRACKER_CONFIGS, ROOT, NUM_THREADS
from tracking.val import (
    run_generate_dets_embs,
    run_generate_mot_results,
    run_trackeval,
    parse_opt as parse_optt,
    download_mot_eval_tools
)


class Tracker:
    """
    Encapsulates the evaluation of a tracking configuration.
    """
    def __init__(self, opt):
        self.opt = opt

    def objective_function(self, config: dict) -> dict:
        """
        Evaluates a given tracker configuration.

        Args:
            config (dict): A dictionary of tracker hyperparameters.

        Returns:
            dict: Combined evaluation metrics extracted from run_trackeval.
        """
        # Ensure evaluation tools are available
        download_mot_eval_tools(self.opt.val_tools_path)
        # Generate MOT-compliant results with the specified tracker parameters
        run_generate_mot_results(self.opt, config)
        # Retrieve evaluation metrics (e.g., MOTA, HOTA, IDF1)
        results = run_trackeval(self.opt)
        # Extract only the desired objective results
        combined_results = {key: results.get(key) for key in self.opt.objectives}
        return combined_results


def load_yaml_config(tracking_method: str) -> dict:
    """
    Loads the YAML configuration file for the given tracking method.

    Args:
        tracking_method (str): Name of the tracking method.

    Returns:
        dict: Configuration parameters loaded from the YAML file.
    """
    config_path = TRACKER_CONFIGS / f"{tracking_method}.yaml"
    with open(config_path, 'r') as file:
        config = yaml.safe_load(file)
    return config


def yaml_to_search_space(config: dict) -> dict:
    """
    Converts a YAML configuration dictionary to a Ray Tune search space.

    Args:
        config (dict): YAML configuration parameters.

    Returns:
        dict: A dictionary representing the search space for hyperparameters.
    """
    search_space = {}
    for param, details in config.items():
        search_type = details.get('type')
        if search_type == 'uniform':
            search_space[param] = tune.uniform(*details['range'])
        elif search_type == 'randint':
            search_space[param] = tune.randint(*details['range'])
        elif search_type == 'qrandint':
            search_space[param] = tune.qrandint(*details['range'])
        elif search_type == 'choice':
            search_space[param] = tune.choice(details['options'])
        elif search_type == 'grid_search':
            search_space[param] = tune.grid_search(details['values'])
        elif search_type == 'loguniform':
            search_space[param] = tune.loguniform(*details['range'])
    return search_space


def main():
    # Parse options and set necessary paths
    opt = parse_optt()
    opt.val_tools_path = EXAMPLES / 'val_utils'
    opt.source = Path(opt.source).resolve()
    opt.yolo_model = [Path(y).resolve() for y in opt.yolo_model]
    opt.reid_model = [Path(r).resolve() for r in opt.reid_model]

    # Load YAML configuration and convert it to a Ray Tune search space
    yaml_config = load_yaml_config(opt.tracking_method)
    search_space = yaml_to_search_space(yaml_config)

    # Create a Tracker instance
    tracker = Tracker(opt)

    # Generate detection and embedding files required for evaluation
    run_generate_dets_embs(opt)

    # Define a wrapper for the objective function for Ray Tune
    def tune_wrapper(config):
        return tracker.objective_function(config)

    results_dir = os.path.abspath("ray/")

    # Set up and run the hyperparameter tuning using Ray Tune
    tuner = tune.Tuner(
        tune.with_resources(tune_wrapper, {"cpu": NUM_THREADS, "gpu": 0}),
        param_space=search_space,
        tune_config=tune.TuneConfig(num_samples=opt.n_trials),
        run_config=RunConfig(storage_path=results_dir)
    )
    tuner.fit()

    # Print the tuning results
    print(tuner.get_results())


if __name__ == "__main__":
    main()