Pi3 (Pi-Cubed) — 图像到 3D 重建模型 (昇腾 NPU 版)

Pi3 (Pi-Cubed, arXiv:2507.13347) 是单模型的多视图图像 → 3D 重建模型。给定同一物体的多张视图图片 (目录) 或一段环绕视频 (.mp4), Pi3 重建出稠密 3D 点云与相机位姿。本仓库为 昇腾 Ascend NPU (torch_npu) 适配版, 全流程无 CUDA。

架构

Pi3 = DINOv2 ViT-L/14 编码器 (pretrained=False)
    + 36 层 RoPE Transformer 解码器 (dim=1024, heads=16, qk_norm, FlashAttentionRope)
    + point_decoder / conf_decoder / camera_decoder (各 5 层 RoPE)
    + point_head(LinearPts3d) + conf_head(LinearPts3d) + camera_head(ResConvBlock + MLP + SVD)

输入: (B, N, 3, H, W) 图像, ImageNet 归一化
输出: points、local_points、conf、camera_poses (4x4 相机位姿)
配置: {"decoder_size": "large", "pos_type": "rope100"}
权重: 单个 model.safetensors (3.8GB, 1210 个 F32 张量, ~958.7M 参数)

NPU 适配点

问题	处理
`SDPBackend.FLASH_ATTENTION` 在 NPU 不可用	Monkeypatch `FlashAttentionRope` → `SDPBackend.MATH / EFFICIENT_ATTENTION`
`torch.amp.autocast(device_type='cuda')`	全部改为 `'npu'`
`torch.svd` 在 NPU 触发 AOE/TBE 编译失败	`torch.linalg.svd` + SVD-free polar 分解回退 (`eigh` 实现)
`PyTorchModelHubMixin.from_pretrained` 走 HF 网络	直接用 `safetensors.safe_open` 本地加载 state dict (strict=True)
DINOv2 编码器 `pretrained=False` 可能访问 torch.hub	本地架构构建, 零网络调用
缺失的 `__init__.py` / `norm.py` / `swiglu_ffn.py` / `vitl14.yaml` / `vision_transformer.py`	全部补齐

环境要求

Python 3.10+
昇腾 NPU (Ascend910) + torch_npu 2.9 + CANN
依赖: torch torch_npu safetensors huggingface_hub numpy pillow torchvision opencv-python plyfile fastapi uvicorn python-multipart

快速开始

# 命令行: 输入图片目录 (多视图) 或 .mp4 视频
python3 inference.py --img sample_views/
python3 inference.py --img sample_views.mp4

# 启动 FastAPI 服务 (根路径为上传表单, POST /predict 返回 3D JSON)
python3 inference.py --server

命令行参数

参数	说明
`--img`	输入图片目录路径或 `.mp4` 视频文件路径 (必需)
`--server`	启动 FastAPI Web 服务 (默认端口 8000)

输出格式 (CLI / FastAPI JSON)

{
  "model": "yyfz233/Pi3",
  "device": "npu:0",
  "input": "sample_views/",
  "num_views": 5,
  "image_size": [504, 504],
  "points_shape": [5, 504, 504, 3],
  "local_points_shape": [5, 504, 504, 3],
  "conf_shape": [5, 504, 504, 1],
  "camera_poses_shape": [5, 4, 4],
  "num_points": 1270080,
  "points_stats": {"min": [...], "max": [...], "mean": [...]},
  "camera_poses": [[[4x4 矩阵], ...], ...],
  "inference_time_s": 1.2,
  "param_count_m": 958.7,
  "flash_attn_backend": "SDPBackend.MATH/EFFICIENT (NPU)",
  "precision": "fp32"
}

points 为每视图稠密 3D 点云 (N, H, W, 3) (归一化坐标)
camera_poses 为每视图 4x4 相机位姿矩阵 (N, 4, 4)
图像预处理与官方一致: 统一缩放到 ≤255000 像素、宽高为 14 的倍数 (ImageNet 归一化)

权重下载

权重 model.safetensors (3.8GB) 自动从镜像下载并缓存到 ./model_weights/yyfz233/Pi3/:

优先检查本地缓存 ./model_weights/yyfz233/Pi3/model.safetensors
缺失时从镜像下载: ai.gitcode.com/hf_mirrors → hf-mirror.com (带断点续传与 429 退避)
直接 safetensors.safe_open 加载 state dict, 不经过 HF Hub 网络

示例数据

sample_views/ — 5 张彩色立方体多视图图片 (由 make_sample_views.py 生成, 纯 PIL 渲染)
sample_views.mp4 — 由上述视图合成的环绕视频

项目结构

Pi3/
├── inference.py            # 主推理脚本 (CLI + FastAPI 双模式)
├── make_sample_views.py    # 示例多视图数据生成器
├── sample_views/           # 示例输入图像目录
├── sample_views.mp4        # 示例输入视频
├── pi3/                    # 昇腾 NPU 适配后的 vendored 模型包
│   ├── models/
│   │   ├── pi3.py
│   │   ├── layers/         # attention/block/pos_embed/camera_head/transformer_head
│   │   └── dinov2/         # DINOv2 ViT-L/14 编码器 (补齐缺失文件)
│   └── utils/              # basic.py / geometry.py
├── model_weights/          # 权重缓存 (symlink 到共享缓存)
├── assets/                 # 截图
│   ├── agent_workflow.png
│   ├── npu_device_call.png
│   └── model_result.png
├── README.md
└── .gitignore

截图

Agent 工作流 (权重下载 + 模型构建)

agent_workflow

CLI 推理 (device=npu:0)

npu_device_call

FastAPI 返回 3D 重建 JSON

model_result