# PyTorch Implementation of LocoTrack

## Preparing the Environment

```bash
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](https://huggingface.co/datasets/hamacojr/LocoTrack-pytorch-weights/resolve/main/locotrack_small.ckpt) |
| LocoTrack-B | [Link](https://huggingface.co/datasets/hamacojr/LocoTrack-pytorch-weights/resolve/main/locotrack_base.ckpt)  |

### 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:

```bash
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](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage?platform=linux). Additionally, downloading instructions for the Huggingface dataset are available at this [link](https://huggingface.co/docs/hub/en/datasets-downloading).

```bash
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:

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