TAPNet

repository·main·Indexed 23 days ago

https://github.com/google-deepmind/tapnet

A collection of state-of-the-art Tracking Any Point (TAP) models and datasets from Google DeepMind, including TAPIR, TAPNext, and RoboTAP. The codebase provides tools for high-performance point tracking in videos, supporting both JAX and PyTorch implementations. It includes benchmarks such as TAP-Vid, RoboTAP, and TAPVid-3D, as well as specialized models like BootsTAPIR and TRAJAN for motion reconstruction.

Tokens
25.2K
Snippets
73
Records
102
Agent score
84%

What's inside tapnet

  1. Understand the TAP-Vid, RoboTAP, and TAPVid-3D benchmarks

    main

    The repository supports three primary benchmarks for evaluating Tracking Any Point (TAP) algorithms:

    1. TAP-Vid: A dataset of videos with point tracks (manual or simulated) designed to evaluate tracking of any trackable point on any solid physical surface. It requires point-level precision and handles deformable surfaces and long-term tracking.

      • Detailed instructions for downloading, using, and evaluating can be found in the tapnet/tapvid README.
    2. RoboTAP: An extension of TAP-Vid focused on robotics. It uses real, manually annotated robotics videos and follows the same annotation format as TAP-Vid.

    3. TAPVid-3D: A benchmark for long-range Tracking Any Point in 3D (TAP-3D), featuring 4,000+ real-world videos with metric 3D position trajectories.

      • Note: This benchmark has a separate license from TAP-Vid.
      • Detailed instructions are in the tapnet/tapvid3d README.
  2. Use the TRAJAN TrackAutoEncoder

    main

    TRAJAN uses a track autoencoder to assess generated videos. The workflow involves:

    1. Loading the track_autoencoder.TrackAutoEncoder model.
    2. Using a ProcessTracksForTrackAutoencoder preprocessor to format point tracks into the required input schema.
    3. Running the forward pass with the model.

    Required Input Schema (Batch Format):

    • video: float[B, T, H, W, 3]
    • support_tracks: float[B, QS, T, 2] (Support tracks coordinates)
    • support_tracks_visible: float[B, QS, T, 1] (Support track visibility)
    • query_points: float[B, Q, 3] (Query points in [t, y, x] format)
    • target_points: float[B, Q, T, 2] (Target tracks coordinates)
    • target_tracks_visible: float[B, Q, T, 1] (Target track visibility)
    • boundary_frame: float[B]
  3. Coordinate systems for 2D and 3D tracking

    main

    When working with coordinates in this project, be aware of the following conventions to avoid metric errors:

    2D Coordinates

    • Storage: Coordinates are typically stored in normalized raster coordinates, where (0, 0) is the upper-left corner and (1, 1) is the lower-right corner.
    • Code Processing: The internal code immediately converts these to regular raster coordinates (matching the Kubric reader). In this format, (0, 0) is the upper-left corner and (h, w) is the lower-right corner, where h is image height and w is width.
    • Order: 2D coordinates are typically stored as (x, y).

    3D Coordinates

    • Order: 3D coordinates are worked in the order (t, y, x).
    • Components:
      • t: Frame coordinates (e.g., 0 is the first frame, 0.5 is halfway between the first and second frame).
      • y: Raster coordinate (height).
      • x: Raster coordinate (width).
  4. Download and generate the TAPVid-3D dataset

    main

    The TAPVid-3D dataset consists of annotations for Aria Digital Twin (ADT), Panoptic Studio (PStudio), and DriveTrack.

    Aria Digital Twin (ADT) specific steps: You must accept the ADT license and download the dataset via the Project Aria Explorer. Once you have the CDN JSON file, use the aria_dataset_downloader to download main_vrs, main_groundtruth, segmentation, and depth files:

    aria_dataset_downloader --cdn_file /PATH_TO/Aria_Digital_Twin_1720774989.json -o /OUTPUT_PATH -d 0 6 7 8

    Generating annotations: Individual generation scripts are located in the annotation_generation subdirectory. You can view help for each:

    python3 -m tapnet.tapvid3d.annotation_generation.generate_adt --help
    python3 -m tapnet.tapvid3d.annotation_generation.generate_pstudio --help
    python3 -m tapnet.tapvid3d.annotation_generation.generate_drivetrack --help

    Batch generation:

    • To generate all files: Run generate_all.sh. This creates a tapvid3d_dataset folder containing *.npz files.
    • To generate only the minival split: Run generate_all.sh but replace --split=all with --split=minival inside the script.
    • To test the pipeline: Run ./run_all.sh --debug. This generates exactly one *.npz/video annotation per source.
    ./run_all.sh --debug
  5. Download and use TAP-Vid-DAVIS and TAP-Vid-RGB-Stacking

    main

    The DAVIS and RGB-Stacking datasets are provided as pickle files.

    Data Structure

    • TAP-Vid-DAVIS: A dictionary where keys are DAVIS video names. Values are dictionaries containing:
      • frames: 4D uint8 tensor.
      • points: float32 tensor with 3 axes (point id, time, x/y).
      • occlusions: 2-axis bool tensor (point id, time).
    • TAP-Vid-RGB-Stacking: Follows the same structure as DAVIS, but is a list of these dictionaries rather than a dictionary keyed by video name.
  6. Download and use model checkpoints

    main

    Model checkpoints can be downloaded from HuggingFace or directly from the storage links provided in the checkpoint table.

    To use a checkpoint with the project's loading utilities, the tapnet/checkpoint/ directory must contain a file named checkpoint.npy that is compatible with the NumpyFileCheckpointer.

    Resolution Note: While evaluation results are often reported at 256x256, higher resolutions can improve performance. For BootsTAPIR, 512x512 is recommended, and for TAPIR, even higher resolutions may be beneficial.

  7. Evaluate on datasets

    main

    Run evaluation on specific datasets (e.g., tapvid_davis) by setting the --jaxline_mode to an evaluation mode and providing the necessary configuration paths.

    Available evaluation datasets are defined in supervised_point_prediction.py.

    python3 -m tapnet.training.experiment \
      --config=./tapnet/configs/tapir_config.py \
      --jaxline_mode=eval_davis_points \
      --config.checkpoint_dir=./tapnet/checkpoint/ \
      --config.experiment_kwargs.config.davis_points_path=/path/to/tapvid_davis.pkl
  8. Download and process TAP-Vid-Kinetics

    main

    TAP-Vid-Kinetics requires downloading labels (CSV) and raw video clips from the Kinetics700-2020 validation set.

    Video File Requirements

    Videos must be stored in a local folder <video_root_path> as MP4 files. The filename must follow the pattern: f'{youtube_id}_{start_time_sec:06}_{end_time_sec:06}.mp4' (e.g., abcdefghijk_000010_000020.mp4).

    Processing Steps

    1. Download the labels from TAP-Vid-Kinetics.
    2. Organize raw clips in <video_root_path> (e.g., <video_root_path>/<action_label>/<clip_name>).
    3. Run generate_tapvid.py to create a pickle file containing the dataset information.
    pip3 install -r requirements.txt
    python3 generate_tapvid.py \
      --input_csv_path=<path_to_tapvid_kinetics.csv> \
      --output_base_path=<path_to_pickle_folder> \
      --video_root_path=<path_to_raw_videos_root_folder> \
      --alsologtostderr