DOCSEARCHbyAlgolia
Submit Docs

TransTrack Documentation

repository·main·Indexed 20 days ago

https://github.com/peizesun/transtrack

TransTrack is a Transformer-based framework for Multiple Object Tracking (MOT) built on Deformable DETR and CenterTrack. It utilizes detection and tracking queries for efficient object tracking in video sequences. The framework includes tools for dataset conversion (CrowdHuman and MOT to COCO), a training pipeline involving pre-training and fine-tuning, and a CLI via main.py for managing hyperparameters, model architecture, and evaluation.

Tokens
2.2K
Snippets
8
Records
8
Agent score
22%

What's inside TransTrack

  1. Prepare CrowdHuman and MOT datasets

    main

    Before training, you must organize your datasets and convert annotations to the COCO format.

    1. Organize Directory Structure

    mkdir crowdhuman
cp -r /path_to_crowdhuman_dataset/CrowdHuman_train crowdhuman/CrowdHuman_train
cp -r /path_to_crowdhuman_dataset/CrowdHuman_val crowdhuman/CrowdHuman_val

mkdir mot
cp -r /path_to_mot_dataset/train mot/train
cp -r /path_to_mot_dataset/test mot/test

    2. Convert Annotations to COCO

    Use the provided tools to convert the raw dataset annotations into the required COCO format:

    • CrowdHuman: python3 track_tools/convert_crowdhuman_to_coco.py
    • MOT: python3 track_tools/convert_mot_to_coco.py
    # Convert CrowdHuman
python3 track_tools/convert_crowdhuman_to_coco.py

# Convert MOT
python3 track_tools/convert_mot_to_coco.py

  2. Train on a mixture of CrowdHuman and MOT17

    main

    To achieve higher performance on the test set, you can train on a mixture of CrowdHuman and MOT17 datasets.

    1. Mix Datasets

    Create a mix directory and link the datasets together:

    mkdir -p mix/annotations
cp mot/annotations/val_half.json mix/annotations/val_half.json
cp mot/annotations/test.json mix/annotations/test.json
cd mix
ln -s ../mot/train mot_train
ln -s ../crowdhuman/CrowdHuman_train crowdhuman_train
cd ..
python3 track_tools/mix_data.py

    2. Train TransTrack

    Run the training script for the mixed dataset:

    sh track_exps/crowdhuman_plus_mot_train.sh
    # Mix data
mkdir -p mix/annotations
cp mot/annotations/val_half.json mix/annotations/val_half.json
cp mot/annotations/test.json mix/annotations/test.json
cd mix
ln -s ../mot/train mot_train
ln -s ../crowdhuman/CrowdHuman_train crowdhuman_train
cd ..
python3 track_tools/mix_data.py

# Train
sh track_exps/crowdhuman_plus_mot_train.sh

  3. Train and Evaluate TransTrack

    main

    Follow these steps to perform the full training pipeline, from pre-training on CrowdHuman to final evaluation.

    Training Pipeline

    1. Pre-train on CrowdHuman:
      sh track_exps/crowdhuman_train.sh
python3 track_tools/crowdhuman_model_to_mot.py
    2. Fine-tune TransTrack (e.g., on mot17_half):
      sh track_exps/crowdhuman_mot_trainhalf.sh

    Evaluation and Visualization

    1. Evaluate performance:
      sh track_exps/mot_val.sh
sh track_exps/mota.sh
    2. Visualize results (convert text output to video):
      python3 track_tools/txt2video.py

    Performance Tuning Tip

    If the MOTA (Multiple Object Tracking Accuracy) is not meeting expectations, try adjusting the --track_thresh parameter. The default is 0.4 (or 0.5 for CrowdHuman).

    # Pre-train
sh track_exps/crowdhuman_train.sh
python3 track_tools/crowdhuman_model_to_mot.py

# Fine-tune
sh track_exps/crowdhuman_mot_trainhalf.sh

# Evaluate
sh track_exps/mot_val.sh
sh track_exps/mota.sh

# Visualize
python3 track_tools/txt2video.py

  4. Install TransTrack

    main

    TransTrack is built on top of Deformable DETR and CenterTrack. Follow these steps to set up the environment and build the necessary libraries.

    Requirements

    • OS: Linux
    • CUDA: >= 9.2
    • GCC: >= 5.4
    • Python: >= 3.7
    • PyTorch: ≥ 1.5 and a matching torchvision installation.
    • OpenCV: Optional (required for demo and visualization).

    Installation Steps

    1. Clone the repository and build the CUDA operators:
    git clone https://github.com/PeizeSun/TransTrack.git
cd TransTrack
cd models/ops
python setup.py build install
cd ../..
    1. Install Python dependencies:
    pip install -r requirements.txt
    git clone https://github.com/PeizeSun/TransTrack.git
cd TransTrack
cd models/ops
python setup.py build install
cd ../..
pip install -r requirements.txt

  5. Resume training from a checkpoint

    main

    To resume training from a previous state, use the --resume argument pointing to a .pth file. The script attempts to restore:

    • Model weights (model key).
    • Optimizer state (optimizer key).
    • Learning rate scheduler state (lr_scheduler key).
    • The current epoch (epoch key).

    Note on Learning Rate: If you resume a checkpoint, the script includes a mechanism to allow --lr_drop to override the resumed scheduler's step size if args.override_resumed_lr_drop is triggered.

    python main.py --resume ./output/checkpoint.pth --output_dir ./output

  6. Train the model using main.py

    main

    To start a training session, run main.py without the --eval flag. The script will initialize the model, build datasets, set up the optimizer (AdamW by default, or SGD if --sgd is passed), and begin the training loop.

    Checkpoints are automatically saved to the --output_dir at specific intervals:

    1. Every 5 epochs.
    2. At the epoch specified by --lr_drop.

    If using distributed training, ensure you provide the appropriate distributed environment variables or flags as expected by utils.init_distributed_mode.

    python main.py --batch_size 2 --epochs 50 --output_dir ./training_run

  7. Evaluate the model using main.py

    main

    To evaluate a trained model, use the --eval flag. You can also provide a --resume path to load a specific checkpoint.

    When --eval is active, the script will:

    1. Load the model and postprocessors.
    2. Load the validation dataset.
    3. Run the evaluate function.
    4. Save evaluation results (e.g., eval.pth or latest.pth) to the --output_dir.
    python main.py --eval --resume ./output/checkpoint.pth --output_dir ./eval_results

  8. Configure Deformable DETR training and evaluation via CLI

    main

    The main.py script serves as the primary entrypoint for training and evaluating the TransTrack/Deformable DETR model. It uses argparse to manage a wide range of hyperparameters including learning rates, model architecture, and dataset paths.

    Key Argument Groups

    Optimization & Learning Rate

    • --lr: Base learning rate (default: 2e-4).
    • --lr_backbone: Learning rate for the backbone (default: 2e-5).
    • --lr_backbone_names: List of parameter names to apply --lr_backbone to (default: ['backbone.0']).
    • --lr_linear_proj_names: List of parameter names for linear projections (default: ['reference_points', 'sampling_offsets']).
    • --lr_linear_proj_mult: Multiplier for the linear projection learning rate (default: 0.1).
    • --epochs: Total training epochs (default: 50).
    • --lr_drop: Epoch at which to drop the learning rate (default: 40).
    • --weight_decay: Weight decay coefficient (default: 1e-4).
    • --sgd: If set, uses SGD optimizer instead of AdamW.

    Model Architecture

    • --backbone: Name of the convolutional backbone (e.g., resnet50).
    • --position_embedding: Type of positional embedding (sine or learned).
    • --num_queries: Number of query slots (default: 300).
    • --enc_layers / --dec_layers: Number of encoder/decoder layers in the transformer (default: 6).
    • --hidden_dim: Dimension of the transformer embeddings (default: 256).
    • --with_box_refine: Enables box refinement.
    • --two_stage: Enables two-stage mechanism.
    • --masks: Enables training the segmentation head.
    • --frozen_weights: Path to a pretrained model. If provided, only the mask head is trained (requires --masks).

    Dataset & Execution

    • --dataset_file: Dataset type (e.g., coco).
    • --coco_path: Path to COCO dataset (default: ./data/coco).
    • --output_dir: Directory to save checkpoints and logs.
    • --eval: Run evaluation mode instead of training.
    • --resume: Path to a checkpoint to resume training from.
    • --device: Device to use (cuda or cpu).
    • --distributed: Enables distributed training.
    # Example: Training with specific backbone and output directory
python main.py --backbone resnet50 --output_dir ./output --epochs 50

# Example: Running evaluation on a saved checkpoint
python main.py --eval --resume ./output/checkpoint.pth --output_dir ./eval_results