DOCSEARCHbyAlgolia
Submit Docs

OSTrack

repository·main·Indexed 20 days ago

https://github.com/botaoye/ostrack

A high-performance, one-stream visual object tracking framework that utilizes self-attention for joint feature learning and relation modeling. It features an Early Candidate Elimination (ECE) module for efficient inference and fast training. The framework supports multiple head types (CORNER and CENTER) and provides tools for training and evaluating on datasets including GOT-10K, LaSOT, TrackingNet, and COCO.

Tokens
12.4K
Snippets
40
Records
46
Agent score
71%

What's inside OSTrack

  1. Train OSTrack models

    main

    To train OSTrack, you must first download the pre-trained MAE ViT-Base weights and place them in the $PROJECT_ROOT$/pretrained_models directory.

    Use tracking/train.py to start training. You can specify the model configuration using the --config flag (configs are located in experiments/ostrack).

    Arguments:

    • --script: Set to ostrack.
    • --config: The desired model configuration name.
    • --save_dir: Directory to save training outputs.
    • --mode: Training mode (e.g., multiple).
    • --nproc_per_node: Number of GPUs to use.
    • --use_wandb: Set to 1 to use Weights & Biases for logging, or 0 to disable.
    python tracking/train.py --script ostrack --config vitb_256_mae_ce_32x4_ep300 --save_dir ./output --mode multiple --nproc_per_node 4 --use_wandb 1

  2. Visualize the candidate elimination process with Visdom

    main

    OSTrack uses Visdom for visualization, specifically to debug the Early Candidate Elimination (ECE) module.

    1. Start the Visdom server: visdom.
    2. Run inference with the --debug 1 flag.
    3. Access the visualization at http://localhost:8097 (adjust IP/port if running on a remote server).
    # 1. Start visdom server
visdom

# 2. Run inference with debug mode enabled
python tracking/test.py ostrack vitb_384_mae_ce_32x4_ep300 --dataset vot22 --threads 1 --num_gpus 1 --debug 1

  3. Set project paths for training and testing

    main

    OSTrack requires local path configurations for workspaces, data, and output directories. You can initialize these automatically using a script, or manually edit the configuration files.

    1. Automatic Setup: Run tracking/create_default_local_file.py with your desired directories.
    2. Manual Setup: Edit the following files to modify paths:
      • lib/train/admin/local.py: Paths related to training.
      • lib/test/evaluation/local.py: Paths related to testing/evaluation.
    python tracking/create_default_local_file.py --workspace_dir . --data_dir ./data --save_dir ./output

  4. Evaluate OSTrack models

    main

    To evaluate models, download the weights from the provided Google Drive and place them in $PROJECT_ROOT$/output/checkpoints/train/ostrack.

    Important: You must update lib/test/evaluation/local.py with the actual paths to your benchmark data.

    Testing Commands by Dataset

    LaSOT (and other offline benchmarks):

    python tracking/test.py ostrack vitb_384_mae_ce_32x4_ep300 --dataset lasot --threads 16 --num_gpus 4
python tracking/analysis_results.py

    GOT10K-test:

    python tracking/test.py ostrack vitb_384_mae_ce_32x4_got10k_ep100 --dataset got10k_test --threads 16 --num_gpus 4
python lib/test/utils/transform_got10k.py --tracker_name ostrack --cfg_name vitb_384_mae_ce_32x4_got10k_ep100

    TrackingNet:

    python tracking/test.py ostrack vitb_384_mae_ce_32x4_ep300 --dataset trackingnet --threads 16 --num_gpus 4
python lib/test/utils/transform_trackingnet.py --tracker_name ostrack --cfg_name vitb_384_mae_ce_32x4_ep300

  5. Profile model FLOPs and speed

    main

    Use tracking/profile_model.py to measure the computational complexity and speed of different model configurations.

    # Profiling vitb_256_mae_ce_32x4_ep300
python tracking/profile_model.py --script ostrack --config vitb_256_mae_ce_32x4_ep300

# Profiling vitb_384_mae_ce_32x4_ep300
python tracking/profile_model.py --script ostrack --config vitb_384_mae_ce_32x4_ep300

  6. Prepare tracking datasets

    main

    Datasets must be organized under a ./data directory at the project root. The expected structure varies by dataset:

    • LaSOT: ./data/lasot/<category_name>
    • GOT-10K: ./data/got10k/ containing test, train, and val subdirectories.
    • COCO: ./data/coco/ containing annotations and images.
    • TrackingNet: ./data/trackingnet/ containing TRAIN_X and TEST subdirectories.

  7. Install the OSTrack environment

    main

    You can install the environment using one of three options depending on your CUDA version and preference:

    Option 1: Anaconda with CUDA 10.2 Create a new conda environment and run the provided installation script.

    Option 2: Anaconda with CUDA 11.3 Use the provided environment YAML file to create the environment.

    Option 3: Docker Use the provided Dockerfile for a containerized setup.

    # Option 1: CUDA 10.2
conda create -n ostrack python=3.8
conda activate ostrack
bash install.sh

# Option 2: CUDA 11.3
conda env create -f ostrack_cuda113_env.yaml

  8. Implement custom data processing with BaseProcessing

    main

    The BaseProcessing class is the foundation for data augmentation and preprocessing pipelines in OSTrack. It is used to transform data returned by a dataset before it enters the neural network (e.g., cropping search regions, applying augmentations).

    When initializing a subclass of BaseProcessing, you can provide specific transformations for different data components:

    • transform: The default transformation applied to images if specific transforms are not provided.
    • template_transform: Transformations applied specifically to template images.
    • search_transform: Transformations applied specifically to search images.
    • joint_transform: Transformations applied to both template and search images simultaneously (e.g., converting both to grayscale).

    Subclasses must implement the __call__(self, data: TensorDict) method to define the actual processing logic.

    from lib.train.data.processing import BaseProcessing
import torchvision.transforms as transforms

class MyCustomProcessing(BaseProcessing):
    def __call__(self, data):
        # Implement custom logic here
        return data

processor = MyCustomProcessing(
    transform=transforms.ToTensor(),
    template_transform=transforms.ColorJitter(),
    joint_transform=transforms.Grayscale()
)

  9. Configure OSTrack head types: CORNER vs CENTER

    main

    The OSTrack model supports two distinct head types, which change how bounding boxes are predicted and what data is returned. This is controlled by the head_type parameter during initialization (or via cfg.MODEL.HEAD.TYPE in build_ostrack).

    FeatureCORNER HeadCENTER Head
    Prediction LogicUses box_head to predict corners, then converts to cxcywh via box_xyxy_to_cxxywh.Uses box_head to directly predict center, bbox, size, and offset maps.
    Output Keyspred_boxes, score_mappred_boxes, score_map, size_map, offset_map
    Coordinate Format[cx, cy, w, h]Varies based on box_head implementation (typically bbox output)

  10. Implement a custom trainer by inheriting from BaseTrainer

    main

    To implement custom training logic, create a new class that inherits from BaseTrainer. You must override the train_epoch method to define the actual training loop for a single epoch. The BaseTrainer handles high-level orchestration including multi-epoch loops, checkpoint saving/loading, and error recovery.

    Required Implementation

    • train_epoch(): This method is called by the base class in every epoch. It should contain the logic for iterating over the data loaders, performing forward/backward passes, and updating weights.
    class MyCustomTrainer(BaseTrainer):
    def train_epoch(self):
        # Implement your custom training logic here
        for batch in self.loaders[0]:
            # ... training steps ...
            pass

  11. Load trackers and datasets for evaluation

    main

    To perform analysis, you must first load your tracker results and the corresponding dataset using the lib.test utilities.

    1. Get Dataset: Use get_dataset(dataset_name) to load the dataset object.
    2. Get Tracker List: Use trackerlist(...) to retrieve specific tracker runs. You need to provide the name (e.g., 'ostrack'), the parameter_name (the specific model configuration), and the dataset_name.

    Example workflow:

    from lib.test.analysis.plot_results import print_results
from lib.test.evaluation import get_dataset, trackerlist

dataset_name = 'lasot'
dataset = get_dataset(dataset_name)

trackers = []
trackers.extend(trackerlist(name='ostrack', parameter_name='vitb_256_mae_ce_32x4_ep300', 
                            dataset_name=dataset_name, run_ids=None, display_name='OSTrack256'))

print_results(trackers, dataset, dataset_name)
    from lib.test.analysis.plot_results import print_results
from lib.test.evaluation import get_dataset, trackerlist

dataset_name = 'lasot'

trackers = []
trackers.extend(trackerlist(name='ostrack', parameter_name='vitb_256_mae_ce_32x4_ep300', dataset_name=dataset_name,
                            run_ids=None, display_name='OSTrack256'))

dataset = get_dataset(dataset_name)
print_results(trackers, dataset, dataset_name, merge_results=True, plot_types=('success', 'prec', 'norm_prec'))

  12. Understand LaSOT and TrackingNet data specification files

    main

    The following files are used for LaSOT and TrackingNet data specifications:

    • lasot_train_split.txt: Defines the complete LaSOT training set.
    • trackingnet_classmap.txt: Provides a mapping from sequence names to their corresponding target classes for the TrackingNet dataset.