# cvlab-kaist/locotrack architecture

> Architecture diagram of the GitHub repository cvlab-kaist/locotrack, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/cvlab-kaist/locotrack
- Repository: https://github.com/cvlab-kaist/locotrack
- Video tour (about a minute): https://gitdiagram.com/cvlab-kaist/locotrack/video

## Overview

LocoTrack is a point-tracking research product with JAX and PyTorch implementations, training and evaluation support, and an interactive video demo. The principal user workflow is to provide a video and query points, run the tracker, and view a rendered track video. The graph also captures the training-data and evaluation branches documented in the README. Some implementation excerpts are partial, so uncertain internal wiring is omitted.

Main entry points: [`demo/demo.py`](https://github.com/cvlab-kaist/locotrack/blob/main/demo/demo.py) (Gradio demo), [`locotrack_pytorch/models/locotrack_model.py`](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/models/locotrack_model.py) (PyTorch tracker), [`locotrack_pytorch/experiment.py`](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/experiment.py) (PyTorch training), [`locotrack/supervised_point_prediction.py`](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/supervised_point_prediction.py) (JAX prediction task)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_demo["Interactive demo"]
  node_video["Video and queries"]
  node_demo_app["Gradio demo<br/>[demo.py]"]
  node_tracked_video["Rendered tracks"]
end

subgraph group_tracking["Tracking models"]
  node_torch_model["PyTorch tracker<br/>[locotrack_model.py]"]
  node_torch_networks["Feature networks<br/>[nets.py]"]
  node_jax_model["JAX tracker<br/>[locotrack_model.py]"]
  node_jax_resnet["JAX ResNet<br/>[resnet.py]"]
end

subgraph group_learning["Training and evaluation"]
  node_torch_experiment["PyTorch training<br/>[experiment.py]"]
  node_jax_prediction["JAX prediction task"]
end

subgraph group_data["Dataset handling"]
  node_torch_eval_data["PyTorch evaluation data"]
  node_torch_kubric["PyTorch Kubric pipeline<br/>[kubric_data.py]"]
  node_jax_kubric["JAX Kubric pipeline<br/>[kubric_data.py]"]
  node_tapvid_prep["TAP-Vid preparation<br/>[generate_tapvid.py]"]
end

node_user(("User"))
node_tapvid_data["TAP-Vid data"]
node_kubric_data["Kubric training data"]
node_evaluation_data["Evaluation datasets"]

node_user -->|"submits video and points"| node_demo_app
node_demo_app -->|"reads input"| node_video
node_demo_app -->|"loads and runs"| node_torch_model
node_torch_model -->|"uses networks"| node_torch_networks
node_torch_model -->|"predicts tracks"| node_tracked_video
node_demo_app -->|"renders output"| node_tracked_video
node_torch_experiment -->|"trains or evaluates"| node_torch_model
node_torch_experiment -->|"loads training data"| node_torch_kubric
node_torch_experiment -->|"loads evaluation data"| node_torch_eval_data
node_torch_kubric -->|"reads"| node_kubric_data
node_torch_eval_data -->|"reads"| node_evaluation_data
node_jax_prediction -->|"runs prediction"| node_jax_model
node_jax_model -->|"uses backbone"| node_jax_resnet
node_jax_prediction -->|"loads evaluation data"| node_jax_kubric
node_jax_kubric -->|"loads"| node_kubric_data
node_tapvid_prep -->|"converts videos and tracks"| node_tapvid_data
node_evaluation_data -.->|"evaluates tracker"| node_jax_prediction

click node_demo_app "https://github.com/cvlab-kaist/locotrack/blob/main/demo/demo.py"
click node_torch_model "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/models/locotrack_model.py"
click node_torch_networks "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/models/nets.py"
click node_jax_model "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/locotrack_model.py"
click node_jax_resnet "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/models/resnet.py"
click node_torch_experiment "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/experiment.py"
click node_jax_prediction "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/supervised_point_prediction.py"
click node_torch_eval_data "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/data/evaluation_datasets.py"
click node_torch_kubric "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/data/kubric_data.py"
click node_jax_kubric "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/data/kubric_data.py"
click node_tapvid_prep "https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/data/generate_tapvid.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_video,node_demo_app,node_tracked_video,node_user toneBlue
class node_torch_model,node_torch_networks,node_jax_model,node_jax_resnet toneAmber
class node_torch_experiment,node_jax_prediction toneMint
class node_torch_eval_data,node_torch_kubric,node_jax_kubric,node_tapvid_prep toneRose
class node_tapvid_data,node_kubric_data,node_evaluation_data toneIndigo
```

## Components

### Interactive demo

- **Video and queries**
- [Gradio demo](https://github.com/cvlab-kaist/locotrack/blob/main/demo/demo.py): `demo/demo.py`
- **Rendered tracks**

### Tracking models

- [PyTorch tracker](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/models/locotrack_model.py): `locotrack_pytorch/models/locotrack_model.py`
- [Feature networks](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/models/nets.py): `locotrack_pytorch/models/nets.py`
- [JAX tracker](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/locotrack_model.py): `locotrack/locotrack_model.py`
- [JAX ResNet](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/models/resnet.py): `locotrack/models/resnet.py`

### Training and evaluation

- [PyTorch training](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/experiment.py): `locotrack_pytorch/experiment.py`
- [JAX prediction task](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/supervised_point_prediction.py): `locotrack/supervised_point_prediction.py`

### Dataset handling

- [PyTorch evaluation data](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/data/evaluation_datasets.py): `locotrack_pytorch/data/evaluation_datasets.py`
- [PyTorch Kubric pipeline](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack_pytorch/data/kubric_data.py): `locotrack_pytorch/data/kubric_data.py`
- [JAX Kubric pipeline](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/data/kubric_data.py): `locotrack/data/kubric_data.py`
- [TAP-Vid preparation](https://github.com/cvlab-kaist/locotrack/blob/main/locotrack/data/generate_tapvid.py): `locotrack/data/generate_tapvid.py`

### Other components

- **User**
- **TAP-Vid data**
- **Kubric training data**
- **Evaluation datasets**

## Connections

- User → Gradio demo: submits video and points
- Gradio demo → Video and queries: reads input
- Gradio demo → PyTorch tracker: loads and runs
- PyTorch tracker → Feature networks: uses networks
- PyTorch tracker → Rendered tracks: predicts tracks
- Gradio demo → Rendered tracks: renders output
- PyTorch training → PyTorch tracker: trains or evaluates
- PyTorch training → PyTorch Kubric pipeline: loads training data
- PyTorch training → PyTorch evaluation data: loads evaluation data
- PyTorch Kubric pipeline → Kubric training data: reads
- PyTorch evaluation data → Evaluation datasets: reads
- JAX prediction task → JAX tracker: runs prediction
- JAX tracker → JAX ResNet: uses backbone
- JAX prediction task → JAX Kubric pipeline: loads evaluation data
- JAX Kubric pipeline → Kubric training data: loads
- TAP-Vid preparation → TAP-Vid data: converts videos and tracks
- Evaluation datasets → JAX prediction task: evaluates tracker