SpatialTracker：在三维空间中追踪任意二维像素

SpatialTracker：在三维空间中追踪任意二维像素，
肖宇熙*，王倩倩*，张尚展，薛楠，彭思达，沈宇军，周晓伟，
CVPR 2024，亮点论文
论文见arxiv

SpatialTrackerV2重磅发布！！！🎉🎉🎉

SpatialTrackerv2 是一款统一模型，可一次性同时生成一致的深度、相机位姿以及像素级的三维追踪结果，相比 V1 实现了 100% 的性能提升。网页地址：SpatialTrackerV2

最新动态与待办事项

• 2025年7月8日：SpatialTracker-v2 正式发布！快来体验吧：🤗 Huggingface Space
• 发布 SpatialTracker 推理代码与检查点。
• 2024年5月4日：SpatialTracker 被选为亮点论文！
• 2024年2月26日：SpatialTracker 被 CVPR 2024 接收！

需求环境

推理代码已在以下环境中测试通过：

• Ubuntu 20.04
• Python 3.10
• PyTorch 2.1.1
• 1 张 NVIDIA GPU（RTX A6000），CUDA 版本 11.8。（其他 GPU 也可使用，22GB 显存足以支持密集追踪——约 1 万个点——我们的代码运行。）

环境搭建

conda create -n SpaTrack python==3.10  
conda activate SpaTrack  

安装 PyTorch

pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118  

其他依赖项

pip install -r requirements.txt  

注意：请严格按照 requirements.txt 中的版本安装依赖，以避免潜在冲突。

深度估计器

在我们的默认设置中，需要使用单目深度估计器从视频输入中获取度量深度。目前有多种模型可供选择（ZoeDepth，Metric3D，UniDepth，以及 DepthAnything）。
我们默认采用 ZoeDepth 模型。请将 dpt_beit_large_384.pt、ZoeD_M12_K.pt、ZoeD_M12_NK.pt 下载至 models/monoD/zoeDepth/ckpts。

数据

我们的方法支持 RGB 或 RGBD 视频输入。我们已在 Google Drive 提供了 checkpoints 和 example_data。请下载 spaT_final.pth 并将其放入 ./checkpoints/。

RGB 视频

对于 example_data，我们提供了 butterfly.mp4 和 butterfly_mask.png 作为示例。请将 butterfly.mp4 和 butterfly_mask.png 下载至 ./assets/，然后运行以下命令：

python demo.py --model spatracker --downsample 1 --vid_name butterfly --len_track 1 --fps_vis 15 --fps 1 --grid_size 40 --gpu ${GPU_id}

RGBD 视频

我们提供了 sintel_bandage.mp4、sintel_bandage.png 以及 sintel_bandage/ 在 example_data 中。其中 sintel_bandage/ 包含了 sintel_bandage.mp4 的深度图。请将 sintel_bandage.mp4、sintel_bandage.png 和 sintel_bandage/ 下载至 ./assets/，然后运行以下命令：

python demo.py --model spatracker --downsample 1 --vid_name sintel_bandage --len_track 1 --fps_vis 15 --fps 1 --grid_size 60 --gpu ${GPU_id} --point_size 1 --rgbd # --vis_support（可选参数，用于可视化所有点）

可视化三维轨迹

首先，请确保已安装 Blender。
我们提供了 Blender 可视化代码：

/Applications/Blender.app/Contents/MacOS/Blender -P create.py -- --input ./vis_results/sintel_bandage_3d.npy

例如，sintel_bandage 的效果如下：

引用

如果您觉得我们的工作对您的研究有所帮助，请考虑引用：

@inproceedings{SpatialTracker,
    title={SpatialTracker：在三维空间中追踪任意二维像素},
    author={肖宇熙和王倩倩和张尚展和薛楠和彭思达和沈宇军和周晓伟},
    booktitle={IEEE/CVF 计算机视觉与模式识别会议（CVPR）论文集},
    year={2024}
}

致谢

Spatialtracker 基于 Cotracker 的代码库构建。我们感谢作者们的卓越贡献，并遵循 Cotracker 的许可协议。

SpaTracker 快速上手指南

环境准备

• 系统要求：Ubuntu 20.04，Python 3.10，NVIDIA GPU（22GB显存足够处理约1万点跟踪）
• CUDA版本：11.8

安装步骤

1. 创建并激活 Conda 环境：

   conda create -n SpaTrack python=3.10
   conda activate SpaTrack
   

2. 安装 PyTorch（使用清华镜像加速）：

   pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装其他依赖：

   pip install -r requirements.txt
   

注意：严格遵循 requirements.txt 中的依赖版本

数据准备

• 下载检查点和示例数据：
  • Google Drive 链接
  • 将 spaT_final.pth 放入 ./checkpoints/
  • RGB 示例：将 butterfly.mp4 和 butterfly_mask.png 放入 ./assets/
  • RGBD 示例：将 sintel_bandage.mp4、sintel_bandage.png 和 sintel_bandage/ 文件夹放入 ./assets/
• （可选）若使用默认深度估计器（ZoeDepth），需下载以下模型到 models/monoD/zoeDepth/ckpts/：
  • dpt_beit_large_384.pt
  • ZoeD_M12_K.pt
  • ZoeD_M12_NK.pt
  • 模型可从 ZoeDepth GitHub 获取，国内建议使用加速下载工具

基本使用

RGB 视频示例

python demo.py --model spatracker --downsample 1 --vid_name butterfly --len_track 1 --fps_vis 15 --fps 1 --grid_size 40 --gpu 0

RGBD 视频示例

python demo.py --model spatracker --downsample 1 --vid_name sintel_bandage --len_track 1 --fps_vis 15 --fps 1 --grid_size 60 --gpu 0 --point_size 1 --rgbd

可视化 3D 轨迹

1. 安装 Blender
2. 运行可视化命令（路径根据系统调整）：

   blender -P create.py -- --input ./vis_results/sintel_bandage_3d.npy