简介

视频相机位姿估计算子。抽帧后先用 MoGe-2 逐帧估计单目深度与内参，再把 (image, depth, intrinsics) 喂给 MegaSaM 的 DROID-SLAM 做跟踪与全局 bundle adjustment，输出相机内参与每帧 camera-to-world 位姿。仅支持 GPU，三个 CUDA 扩展与全部权重必须离线预置，算子内零网络零编译。

功能描述

• 两级流水线：MoGe-2 出的是归一化内参，乘回 width / height 变成像素单位；第二级把每帧等比缩放到 384×512 面积、裁到 8 的倍数，内参同步按缩放比修正，深度用 nearest-exact 插值到同尺寸
• DROID 吃的是原始通道序（BGR），本实现不做转换；落 npz 时才转回 (T,H,W,3) RGB，与 mega-sam cvd_opt 的约定一致
• 逐帧 track → track_final → terminate(_opt_intr=opt_intrinsics, full_ba=full_ba) 做全局 BA；内参取 droid.video.intrinsics[0] * 8.0（SLAM 在 1/8 特征分辨率上算），位姿用 lietorch.SE3(traj).inv().matrix() 从 w2c 转成 c2w
• DROID 全部超参（buffer / filter_thresh / warmup / keyframe_thresh / frontend 四项 / backend 三项 / beta）逐个开放为算子参数，以普通对象而非 dict 传入（DROID 内部直接读属性）
• count_k默认 64：DROID 的 warmup=8 意味着 3 帧根本起不了轨迹；另加前置检查，抽到的帧数少于 warmup 时直接给出「调大 count_k / max_frames」的错误信息，而不是让 SLAM 崩在内部
• 大张量不进列：intrinsic / cam_c2w 以 {shape, data}（形状 + 一维 float32）存出，depths 默认关闭，images 完全不进列（一段 300 帧视频的 images 与 depths 是百 MB 量级）。需要完整产物时设 output_npz_dir，落成键名与 mega-sam cvd_opt 一致（images / depths / intrinsic / cam_c2w）的 npz，文件名为 <视频名>_droid.npz（视频名取路径/URL 去参数后的 stem，取不到用 video_<行号>）
• torch.load 打补丁：mega-sam 的 Droid.load_weights 是裸 torch.load，torch>=2.6 默认 weights_only=True 会失败，构建 Droid 期间临时包一层 weights_only=False，构建完立即还原
• 两个源自 DROID 本身的既有行为：① full_ba=True 时 DROID 返回的是关键帧轨迹 video.poses[:counter]，逐帧的 traj_filler 结果在这条分支里被算出来又丢弃，所以 frame_count 通常小于抽帧数（实测抽 32 帧出 31 个关键帧）；② 输出的 depths 是 MoGe 深度按 SLAM 分辨率缩放后的结果，不是 BA 优化后的 depth_est（BA 返回的优化深度未被使用）
• 抽帧走 VideoFrameSampler 在内存里完成，不把帧写盘；输入支持 video_url（含 BOS 路径）、video_base64、video_binary
• 抽不到帧或单个视频估计失败时，该行张量为空、npz_path 为空串、status 记录原因，不中断整列；每个视频处理完调 torch.cuda.empty_cache()

算子参数

输入

输出

输出为一个 struct 列，字段如下（tensor 均指 struct<shape: list<int32>, data: list<float32>>，data 是按行优先展平的 float32；缺失时为 {"shape": [], "data": []}）：

参数

注意事项

• 三个 CUDA 扩展需要提前按目标 GPU 架构编好：droid_backends、lietorch、torch_scatter。算子只做 import 检查，缺失时抛 ImportError，运行期不做编译。
• 编译产物与 python 版本、torch 版本、GPU sm 架构三重绑定：换 python 小版本、换 torch 或换卡都得重编；要同时覆盖多种架构的卡，需要编出含对应架构的 fat binary。
• 必须有 GPU：无 CUDA 时 __init__ 直接 RuntimeError，DROID-SLAM 的 CUDA 扩展没有 CPU 实现。
• moge_model_name必须指到.pt文件：MoGe 的 from_pretrained 传目录会被当成 HF repo id 处理而报错。megasam_ckpt_name 同样指向具体的 .pth 文件，两个路径任一不存在即 FileNotFoundError；megasam_repo_path 下缺 base/droid_slam 也会直接报错。
• 多卡靠 worker 级隔离：Droid.load_weights 内部把网络固定放到 cuda:0，算子的 rank 只影响 MoGe 落卡，多卡场景需要 worker 级 CUDA_VISIBLE_DEVICES。
• 按concurrency=1起：测试用法是 num_cpus=4, num_gpus=1, concurrency=1, batch_size=1；buffer 默认 1024 个关键帧槽位占显存，短片建议下调（测试用 256）。实测 test_video.mp4 抽 32 帧全程 85 秒，npz 42.7M。
• 首帧位姿只是接近单位矩阵：DROID 以第一个关键帧为世界原点，但 full BA 会把所有位姿（含首帧）一起再优化、不做 gauge 固定，实测首帧 c2w 与单位矩阵偏差约 2e-3。做校验时要留容差。
• mega-sam README 里要求的 DepthAnything / RAFT 权重是它自己 mono-depth 与 cvd_opt 流程用的，本算子这条路径不需要。

调用示例

1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.embodied.perception.video_camera_pose import VideoCameraPose
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18ds = daft.from_pydict({"video": ["bos://your-bucket/sample.mp4"]})
19    ds = ds.with_column(
20        "pose",
21        aihc_udf(
22            VideoCameraPose,
23            construct_args={
24                "model_path": "/path/to/models",
25                "megasam_repo_path": "/path/to/thirdparty/mega-sam",
26                "count_k": 32,
27                "buffer": 256,
28                "output_npz_dir": "/tmp/camera_pose_npz",
29            },
30            num_cpus=4,
31            num_gpus=1,
32            concurrency=1,
33            batch_size=1,
34        )(col("video")),
35    )
36    ds = ds.with_column("frame_count", col("pose")["frame_count"])
37    ds = ds.with_column("status", col("pose")["status"])
38    ds.show()