CoTracker Documentation

repository·main·Indexed 26 days ago

https://github.com/facebookresearch/co-tracker

A transformer-based model for high-performance point tracking in videos. CoTracker supports tracking any pixel or quasi-dense sets of pixels via offline batch processing or memory-efficient online streaming. The library provides multiple model versions (v1, v2, and CoTracker3) accessible via PyTorch Hub, along with utilities for visualization, training on Kubric, and fine-tuning with pseudo labels.

Tokens: 3.4K
Snippets: 10
Records: 16
Agent score: 90%

What's inside CoTracker

Install CoTracker from source
main
To run local demos or perform evaluation/training, install the development version from the repository. Ensure PyTorch and TorchVision (ideally with CUDA support) are already installed.
```
git clone https://github.com/facebookresearch/co-tracker
cd co-tracker
pip install -e .
pip install matplotlib flow_vis tqdm tensorboard
```
https://github.com/facebookresearch/co-tracker/blob/main/README.md

Use CoTracker3 in Offline Mode via PyTorch Hub

main

Load the CoTracker3 offline model using torch.hub. This mode is suitable for processing entire videos at once. The model expects a video tensor of shape (B, T, C, H, W) and returns predicted tracks and visibility.

Note: You must install imageio[ffmpeg] to handle video loading.

import torch
import imageio.v3 as iio

# Download the video
url = 'https://github.com/facebookresearch/co-tracker/raw/refs/heads/main/assets/apple.mp4'
frames = iio.imread(url, plugin="FFMPEG")

device = 'cuda'
grid_size = 10
# Prepare video tensor: B T C H W
video = torch.tensor(frames).permute(0, 3, 1, 2)[None].float().to(device)

# Run Offline CoTracker:
cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker3_offline").to(device)
pred_tracks, pred_visibility = cotracker(video, grid_size=grid_size) # B T N 2,  B T N 1

Download CoTracker3 checkpoints
main
If installing from source, you may need to manually download model checkpoints into a checkpoints directory.
https://github.com/facebookresearch/co-tracker/blob/main/README.md

Fine-tune CoTracker with pseudo labels

main

Fine-tuning requires a custom dataset of real videos. Your dataset class should implement video loading and store it in a CoTrackerData class. Use train_on_real_data.py to generate pseudo labels during training. You must start with an existing Kubric-trained model checkpoint.

# Example: Fine-tuning the online model
python ./train_on_real_data.py --batch_size 1 --num_steps 15000 \
 --ckpt_path ./ --model_name cotracker_three --save_freq 200 --sequence_len 64 \
 --eval_datasets tapvid_stacking tapvid_davis_first --traj_per_sample 384 \
 --save_every_n_epoch 15 --evaluate_every_n_epoch 15 --model_stride 4 --dataset_root ${path_to_your_dataset} --num_nodes 4 --real_data_splits 0 \
 --num_virtual_tracks 64 --mixed_precision --random_frame_rate \
 --restore_ckpt ./checkpoints/baseline_online.pth \
 --lr 0.00005 --real_data_filter_sift --validate_at_start \
 --sliding_window_len 16 --limit_samples 15000

Evaluate CoTracker models on TAP-Vid
main
To reproduce paper results, download the TAP-Vid and Dynamic Replica datasets. Install the required dependencies using pip install hydra-core==1.1.0 mediapy.
Use the evaluate.py script to run evaluations. Note that evaluations are run jointly on all target points for speed; to reproduce exact paper numbers, add the single_point=True flag.
https://github.com/facebookresearch/co-tracker/blob/main/README.md

Use CoTracker3 in Online Mode via PyTorch Hub

main

Load the CoTracker3 online model using torch.hub. This mode is more memory-efficient and allows for processing longer videos or streams by processing video chunks.

To use online mode:

Initialize the model with the first video chunk using is_first_step=True.
Process subsequent chunks by passing video_chunk to the model.

import torch

device = 'cuda'
grid_size = 10
# video shape: B T C H W

cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker3_online").to(device)

# Initialize online processing
cotracker(video_chunk=video, is_first_step=True, grid_size=grid_size)

# Process the video in chunks
for ind in range(0, video.shape[1] - cotracker.step, cotracker.step):
    pred_tracks, pred_visibility = cotracker(
        video_chunk=video[:, ind : ind + cotracker.step * 2]
    )  # B T N 2,  B T N 1

Train CoTracker baseline on Kubric

main

To train the baseline model, you must first generate annotations for the Google Kubric MOVI-f dataset.

Install training dependencies:

pip install pip==24.0
pip install pytorch_lightning==1.6.0 tensorboard opencv-python

Launch training using train_on_kubric.py. Ensure you modify --dataset_root and --ckpt_path before execution. For multi-node training (e.g., 4 nodes), include the --num_nodes 4 flag.

# Example: Launch training of the online model on Kubric
python train_on_kubric.py --batch_size 1 --num_steps 50000 \
 --ckpt_path ./ --model_name cotracker_three --save_freq 200 --sequence_len 64 \
  --eval_datasets tapvid_davis_first tapvid_stacking --traj_per_sample 384 \
  --sliding_window_len 16 --train_datasets kubric --save_every_n_epoch 5 \
  --evaluate_every_n_epoch 5 --model_stride 4 --dataset_root ${path_to_your_dataset} \
   --num_nodes 4 --num_virtual_tracks 64 --mixed_precision --corr_radius 3 \ 
   --wdecay 0.0005 --linear_layer_for_vis_conf --validate_at_start --add_huber_loss

Build CoTracker documentation
main
To generate documentation from Sphinx, install sphinx and sphinxcontrib-bibtex, then run the make command from the docs directory.
```
pip install sphinx
pip install sphinxcontrib-bibtex
make -C docs html
```
https://github.com/facebookresearch/co-tracker/blob/main/README.md
Visualize predicted tracks with Visualizer
main
After installing the CoTracker package, you can use the Visualizer utility to create videos of the predicted tracks.
```
from cotracker.utils.visualizer import Visualizer

vis = Visualizer(save_dir="./saved_videos", pad_value=120, linewidth=3)
vis.visualize(video, pred_tracks, pred_visibility)
```
https://github.com/facebookresearch/co-tracker/blob/main/README.md

Use CoTracker v2 via PyTorch Hub

main

CoTracker v2 can be loaded via torch.hub. It supports both Offline and Online modes.

Offline Mode: Processes the entire video at once. Online Mode: Uses a different API designed for chunked/streaming processing.

import torch
import imageio.v3 as iio

# Setup
device = 'cuda'
grid_size = 10
url = 'https://github.com/facebookresearch/co-tracker/blob/main/assets/apple.mp4'
frames = iio.imread(url, plugin="FFMPEG")
video = torch.tensor(frames).permute(0, 3, 1, 2)[None].float().to(device)  # B T C H W

# Offline Mode
cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker2").to(device)
pred_tracks, pred_visibility = cotracker(video, grid_size=grid_size) # B T N 2, B T N 1

# Online Mode
cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker2_online").to(device)
cotracker(video_chunk=video, is_first_step=True, grid_size=grid_size)  

for ind in range(0, video.shape[1] - cotracker.step, cotracker.step):
    pred_tracks, pred_visibility = cotracker(
        video_chunk=video[:, ind : ind + cotracker.step * 2]
    )

Use CoTracker v1 via PyTorch Hub
main
CoTracker v1 is available directly via PyTorch Hub using the v1.0 tag.
```
import torch
cotracker = torch.hub.load("facebookresearch/co-tracker:v1.0", "cotracker_w8")
```
https://github.com/facebookresearch/co-tracker/blob/main/README.md
Use CoTracker model utilities
main
The cotracker.models.core.model_utils module provides utility functions for working with CoTracker models. These utilities are intended to assist with model-related operations during inference or setup.
https://github.com/facebookresearch/co-tracker/blob/main/docs/source/apis/models.rst