SpatialTracker：在3D空间中跟踪任意2D像素

SpatialTracker：在3D空间中跟踪任意2D像素,
肖宇希*、王倩倩*、张尚展、薛楠、彭思达、沈昱君、周晓伟,
CVPR 2024, 亮点论文 论文链接 arxiv

新闻和待办事项

发布 SpatialTracker-v2 (即将推出)。
发布 SpatialTracker 推理代码和检查点。
2024年4月5日：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 也适用，使用我们的代码进行密集跟踪（约10k个点）需要22GB GPU内存。)

设置环境

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

<mark>注意：</mark>请按照 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 上提供了 检查点 和 示例数据。请下载 spaT_final.pth 并将其放入 ./checkpoints/ 目录。

RGB 视频

对于 示例数据，我们提供了 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/。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 (可选，用于可视化所有点)

可视化 3D 轨迹

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

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

例如，sintel_bandage 的效果如下：

引用

如果您在研究中使用了我们的工作，请考虑引用：

@inproceedings{SpatialTracker,
    title={SpatialTracker: Tracking Any 2D Pixels in 3D Space},
    author={Xiao, Yuxi and Wang, Qianqian and Zhang, Shangzhan and Xue, Nan and Peng, Sida and Shen, Yujun and Zhou, Xiaowei},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2024}
}