PyTorch Implementation of LocoTrack

Preparing the Environment

git clone https://github.com/google-research/kubric.git

conda create -n locotrack-pytorch python=3.11
conda activate locotrack-pytorch

pip install torch torchvision torchaudio lightning==2.3.3 tensorflow_datasets tensorflow matplotlib mediapy tensorflow_graphics einshape wandb

LocoTrack Evaluation

1. Download Pre-trained Weights

To evaluate LocoTrack on the benchmarks, first download the pre-trained weights.

Model	Pre-trained Weights
LocoTrack-S	Link
LocoTrack-B	Link

2. Adjust the Config File

In config/default.ini (or any other config file), add the path to the evaluation datasets to [TRAINING]-val_dataset_path. Additionally, adjust the model size for evaluation in [MODEL]-model_kwargs-model_size.

3. Run Evaluation

To evaluate the LocoTrack model, use the experiment.py script with the following command-line arguments:

python experiment.py --config config/default.ini --mode eval_{dataset_to_eval_1}_..._{dataset_to_eval_N}[_q_first] --ckpt_path /path/to/checkpoint --save_path ./path_to_save_checkpoints/

--config: Specifies the path to the configuration file. Default is config/default.ini.
--mode: Specifies the mode to run the script. Use eval to perform evaluation. You can also include additional options for query first mode (q_first), and the name of the evaluation datasets. For example:
- Evaluation of the DAVIS dataset: eval_davis
- Evaluation of DAVIS and RoboTAP in query first mode: eval_davis_robotap_q_first
--ckpt_path: Specifies the path to the checkpoint file. If not provided, the script will use the default checkpoint.
--save_path: Specifies the path to save logs.

Replace /path/to/checkpoint with the actual path to your checkpoint file. This command will run the evaluation process and save the results in the specified save_path.

LocoTrack Training

Training Dataset Preparation

Download the panning-MOVi-E dataset used for training (approximately 273GB) from Huggingface using the following script. Git LFS should be installed to download the dataset. To install Git LFS, please refer to this link. Additionally, downloading instructions for the Huggingface dataset are available at this link.

git clone git@hf.co:datasets/hamacojr/LocoTrack-panning-MOVi-E

Training Script

Add the path to the downloaded panning-MOVi-E to the [TRAINING]-kubric_dir entry in config/default.ini (or any other config file). Optionally, for efficient training, change [TRAINING]-precision in the config file to bf16-mixed to use bfloat16. Then, run the training with the following script:

python experiment.py --config config/default.ini --mode train_davis --save_path ./path_to_save_checkpoints/