feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe
- Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing. - Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations. - Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Added model synchronization script `sync_models.py` for downloading and verifying ONNX models. - Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality. - Updated `vendor.lock.toml` to remove unused artifacts.
This commit is contained in:
1
.gitignore
vendored
1
.gitignore
vendored
@@ -8,3 +8,4 @@ build/
|
||||
.pytest_cache/
|
||||
|
||||
/vendor/
|
||||
/models/
|
||||
|
||||
14
README.md
14
README.md
@@ -16,11 +16,11 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一
|
||||
2. 打开 Blender，进入 `Preferences > Extensions`
|
||||
3. 点击 `Install from Disk`，选择下载的压缩包
|
||||
4. 在扩展偏好设置中确认 Storage Root 与 Device（保持默认即可）；Device 默认取本机
|
||||
最快的 CUDA/MPS，否则 CPU。下载超时会自动重试国内镜像，无需手动配置
|
||||
最快的 DirectML/CoreML provider，否则 CPU。下载超时会自动重试国内镜像，无需手动配置
|
||||
5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Environment`
|
||||
|
||||
DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时直接从本地解包，不需要访问 GitHub。首次使用仍会自动
|
||||
下载运行环境依赖（PyTorch 等），随后自动下载 `MoGe-2 ViT-B Normal` 和 `BEN2 Base`。其他 AI 模型可在
|
||||
扩展只内置对应平台的 uv。首次使用会自动下载轻量的 ONNX Runtime 环境，随后自动下载 `MoGe-2 ViT-B Normal` 和
|
||||
`BEN2 Base`；不安装 PyTorch，也不下载各模型的 Python 源码包。其他 AI 模型可在
|
||||
`Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。
|
||||
|
||||
## 快速上手
|
||||
@@ -46,14 +46,18 @@ DA3、MoGe-2、utils3d、BEN2 源码包和 uv 已内置在扩展里，安装时
|
||||
|
||||
选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。可选择 MoGe-2
|
||||
ViT-S、ViT-B、ViT-L 或 DA3 Small、Base，其中 MoGe-2 ViT-B 为默认模型。SnapForm 会按模型输出的
|
||||
相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 直接使用官方
|
||||
`utils3d.build_mesh_from_map` 与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes
|
||||
相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 使用内置 NumPy
|
||||
网格构建与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes
|
||||
修改器完成剔面、快速重拓扑、Relax 与 UV 重投。开启 Remove Background 时，背景三角面会在转换期间
|
||||
直接删除，不再交给 Geometry Nodes。匹配相机默认不创建；开启后会作为 Surface 的子级。
|
||||
MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 Surface Layer`，并与原图
|
||||
Bump 合并；网格本身不写入 custom normal，因此重拓扑降低顶点数后仍能保留细腻的法线细节。
|
||||
当前此功能仅支持单张图片。
|
||||
|
||||
MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base 与 BEN2 全部使用 ONNX。模型执行、相机与深度恢复、
|
||||
点云重投影和网格生成均不依赖 PyTorch。Windows 使用 DirectML 加速，provider 不可用时会回退到 ONNX Runtime CPU，
|
||||
并在任务状态中显示实际 provider。
|
||||
|
||||
转换对话框提供模型与 `Create Matching Camera`、`Remove Background`、`Unshaded`；DA3 显示
|
||||
`Resolution Scale`，MoGe-2 显示 `Resolution Level`，
|
||||
以及分割线下方的 `Keep Original`。
|
||||
|
||||
@@ -4,17 +4,17 @@
|
||||
|
||||
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。
|
||||
2. 右键打开 `SnapForm` 子菜单并选择 `Convert to Depth Plane`，通过 `Depth from` 选择深度来源：
|
||||
- `DA3`：显示 Model、Use Ray Pose 和 Resolution Scale。Process Res 先以原图长边乘以
|
||||
- `DA3`：显示 Model 和 Resolution Scale。ONNX 模型使用 camera decoder 恢复相机。Process Res 先以原图长边乘以
|
||||
Resolution Scale，再钳制到 `252–1008`，最后对齐到最接近的 `14` 倍数。Resolution Scale 以百分比
|
||||
显示，范围为 `1–100%`，默认 `100%`；界面只在标签中实时显示最终 Process Res。
|
||||
转换导出 `mini_npz`，使用前需要在
|
||||
`Models` panel 下载所选模型，并固定使用 `process_res_method="upper_bound_resize"`；设备和参考视图策略
|
||||
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型；若没有可用模型，对话框会引导到
|
||||
`Models` panel 下载。
|
||||
- `MoGe-2`：使用所选模型输出 metric depth、normal、validity mask 与 normalized intrinsics。Depth Plane 不保存体积较大的 points；Surface 会在 AI 进程内用 points 和官方 `utils3d.build_mesh_from_map` 直接生成 `mesh.npz`。Debug 始终保留原始 `results.npz`，并按 Output Format 单独生成 Maps、纹理 GLB 或彩色 PLY。
|
||||
- `MoGe-2`：使用所选模型输出 metric depth、normal、validity mask 与 normalized intrinsics。Depth Plane 不保存体积较大的 points；Surface 会在 AI 进程内用内置 NumPy 网格构建生成 `mesh.npz`。Debug 始终保留原始 `results.npz`，并按 Output Format 单独生成 Maps、纹理 GLB 或彩色 PLY。
|
||||
- `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2
|
||||
取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。
|
||||
`Original` 是默认深度来源。选择 `DA3` 时，其 Model、Use Ray Pose 与 Resolution Scale
|
||||
`Original` 是默认深度来源。选择 `DA3` 时，其 Model 与 Resolution Scale
|
||||
参数会集中显示在独立边框中。`Remove Background` 与 `Unshaded` 直接显示，不再使用 Options 折叠；
|
||||
`Keep Original` 默认关闭，并在分割线下方保持为最后一项。仅当开启 `Remove Background` 时才需要
|
||||
`BEN2 Base`，并把 `.foreground` 的 Alpha 作为独立
|
||||
@@ -63,4 +63,3 @@
|
||||
节点组根据 `Frame` 计算从 `1` 开始的 Image Texture Frame：场景帧与 Start 都为 `1` 时读取图片第 `1` 帧；非循环输入
|
||||
钳制到有效范围；循环输入使用 Floored Modulo，让 `Frame Start` 前后的负、正相对帧都映射到
|
||||
`1…Frames`，因此开始帧之前也会循环播放。Geometry 置换会像材质 Image Texture 一样随场景帧刷新。
|
||||
|
||||
|
||||
@@ -12,11 +12,10 @@ uv run python pack.py
|
||||
|
||||
```text
|
||||
.
|
||||
├── vendor.lock.toml 第三方构建输入的版本、URL 与 SHA-256
|
||||
├── prepare_vendor.py 下载并校验第三方构建输入
|
||||
├── vendor.lock.toml 各平台 uv 的版本、URL 与 SHA-256
|
||||
├── prepare_vendor.py 下载并校验各平台 uv
|
||||
├── pack.py
|
||||
├── patches/ 进入 Git 的 DA3 源码覆盖文件
|
||||
├── vendor/ 不进入 Git 的下载缓存；本地调试也直接读取这里
|
||||
├── vendor/ 不进入 Git 的 uv 下载缓存
|
||||
├── src/
|
||||
│ └── snapform/
|
||||
│ ├── blender_manifest.toml
|
||||
@@ -36,7 +35,6 @@ uv run python pack.py
|
||||
## 上游与许可证
|
||||
|
||||
- Depth Anything 3 代码采用 Apache 2.0。
|
||||
- BEN2 代码与模型权重采用 MIT License。
|
||||
- BEN2 ONNX 模型采用其仓库声明的 Apache 2.0 License。
|
||||
- 本 Extension 代码采用 GPL-3.0-or-later。
|
||||
- 模型权重许可证由所选 Hugging Face 模型仓库决定，Extension 下载前会在界面显示。
|
||||
|
||||
|
||||
@@ -12,14 +12,14 @@
|
||||
## 当前功能
|
||||
|
||||
- 使用 Blender Extension manifest，可通过 `Install from Disk` 安装
|
||||
- 在独立 Environment 中通过常驻 Service 运行 Torch、DA3、MoGe-2 与 BEN2，避免污染 Blender Python 进程
|
||||
- 在独立 Environment 中通过常驻 Service 运行 ONNX Runtime；MoGe-2、DA3 与 BEN2 的模型执行及后处理完全不使用 PyTorch，避免污染 Blender Python 进程
|
||||
- Service 的深度模型缓存只保留 DA3 或 MoGe-2 其中一个，切换时主动释放前一个；BEN2 使用独立缓存，
|
||||
可与当前深度模型共存
|
||||
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Service；请求带随机 Token
|
||||
- 常驻 Service 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型，连续执行时不再重复加载
|
||||
- 常驻 Service 复用 ONNX Runtime session、BEN2 与最近使用的深度模型，连续执行时不再重复加载
|
||||
- 多个 Blender 实例自动连接同一个共享 Service；`AI Service` 面板显示状态、端口、客户端与模型缓存，并提供安装、启动、清理、重启和停止控制
|
||||
- Environment、模型和任务结果默认保存在 `~/.snapform`，可由多个 Blender 版本共享
|
||||
- 当前启用 DA3 Small 和 Base；Large 1.1、Giant 1.1 与 Nested Giant Large 1.1 暂停使用
|
||||
- 当前仅启用 FP32 DA3 Small 和 Base；它们不随首次安装下载，Giant 1.1 与 Nested Giant Large 1.1 暂停使用
|
||||
- 支持 BEN2 Base 单图背景剔除测试，输出透明 PNG 与推理耗时
|
||||
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
|
||||
- 视频按顺序截取前若干帧，可设置参与重建的最大帧数
|
||||
@@ -52,17 +52,16 @@
|
||||
1. 先运行 `uv run python prepare_vendor.py`，按 `vendor.lock.toml` 下载并校验固定版本的
|
||||
第三方构建输入；再运行 `uv run python pack.py`，分别构建 `windows-x64`、`macos-arm64`、
|
||||
`linux-x64` 三个平台的 Extension ZIP（如 `SnapForm.v0.10.12.macos-arm64.zip`）。
|
||||
每个包只含对应平台的 uv 二进制，并内置 DA3、BEN2 源码包。
|
||||
每个包只含对应平台的 uv 二进制，不再携带 AI 项目的源码包。
|
||||
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
|
||||
3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能
|
||||
（Image Empty 右键菜单）使用的推理设备，默认取本机最快的 CUDA/MPS，否则 CPU，且只列出
|
||||
（Image Empty 右键菜单）使用的推理设备，默认取本机最快的 DirectML/CoreML provider，否则 CPU，且只列出
|
||||
本机支持的设备。
|
||||
4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、PyTorch、FastAPI、Uvicorn；
|
||||
DA3、BEN2 源码包与 uv 已内置在扩展中，安装时从 安装任务的临时目录中使用，安装结束后自动清理，
|
||||
并在 Environment 安装后自动下载 `DA3 Small` 和 `BEN2 Base` 推荐模型。Environment 安装
|
||||
固定显示 `Environment`、`PyTorch`、`AI components` 三个用户阶段；xFormers（macOS 跳过）、
|
||||
FastAPI 等内部安装步骤合并在 `AI components` 中。耗时最长的 PyTorch 与 AI 组件占主要进度
|
||||
区间。安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证。推荐模型属于安装
|
||||
4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、ONNX Runtime、FastAPI、Uvicorn；
|
||||
uv 已内置在扩展中，DA3、BEN2 与 MoGe 不安装源码包，
|
||||
并在 Environment 安装后自动下载 `MoGe-2 ViT-S Normal` 和 FP16 `BEN2 Base` 推荐模型。Environment 安装
|
||||
固定显示 `Environment`、`ONNX Runtime` 两个用户阶段；FastAPI 等内部安装步骤合并在 `ONNX Runtime` 中。
|
||||
安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证。推荐模型属于安装
|
||||
后的自动步骤，不计入阶段编号；`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。
|
||||
状态文字使用 `阶段/总数 阶段名 | 当前操作 | 预计剩余时间`，依赖安装会结合预设耗时与实际
|
||||
进度修正预期；模型下载根据字节数和速度计算剩余时间，例如
|
||||
|
||||
@@ -1,31 +1,29 @@
|
||||
# 输出格式
|
||||
|
||||
## 输出格式说明
|
||||
## DA3 ONNX
|
||||
|
||||
`Output Type` 提供以下结果：
|
||||
|
||||
- Point Cloud (GLB)：DA3 `glb`
|
||||
- Gaussian Splats (PLY)：DA3 `gs_ply`，同时传入 `infer_gs=True`
|
||||
- Depth Preview (Images)：DA3 `depth_vis`
|
||||
- Depth Data (NPZ)：DA3 `mini_npz`
|
||||
- Textured Surface (GLB)：SnapForm 根据 DA3 Prediction 构建并输出 `mesh.glb`
|
||||
- Point Cloud (GLB)：由 depth、confidence、相机参数和原图生成彩色点云
|
||||
- Depth Preview (Images)：保存深度着色预览
|
||||
- Depth Data (NPZ)：保存 depth、confidence、extrinsics、intrinsics 与处理后的图像
|
||||
- Textured Surface (GLB)：SnapForm 根据 Prediction 构建纹理网格
|
||||
|
||||
Point Cloud 专属设置仅在对应类型下显示，`gs_views_interval` 仅用于 Gaussian Splats。
|
||||
Textured Surface 使用 0–100% 的 Confidence Threshold 删除置信度最低的一定比例：
|
||||
DA3 Small 与 Base 的 ONNX 图不包含 Gaussian head，因此不提供 Gaussian Splats。默认推理参数为：
|
||||
|
||||
- `process_res=504`
|
||||
- `process_res_method="upper_bound_resize"`
|
||||
- `use_ray_pose=False`
|
||||
- `ref_view_strategy="saddle_balanced"`
|
||||
- `conf_thresh_percentile=40.0`
|
||||
- `num_max_points=1_000_000`
|
||||
- `show_cameras=True`
|
||||
- `export_kwargs["gs_ply"]["gs_views_interval"]=1`
|
||||
|
||||
DA3 官方 `glb` 是彩色点云与可选相机线框；`gs_ply` 是带 Gaussian 属性的 PLY；
|
||||
`depth_vis` 是原图与着色深度图的并排预览；`mini_npz` 保存深度、置信度和可用的相机参数。
|
||||
Point Cloud GLB 任务会关闭其默认附带的深度预览，使五种输出产物保持分离。
|
||||
相机外参与内参由模型的 camera decoder 输出恢复，不运行 ray-pose 求解器。
|
||||
|
||||
官方格式与限制可参阅
|
||||
[DA3 Python API 文档](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/docs/API.md)。
|
||||
## MoGe-2 ONNX
|
||||
|
||||
MoGe-2 Debug 可保存原始 NPZ、depth/normal/mask maps、纹理 GLB 或彩色 PLY。Surface 输出由模型 points、mask 与 depth edge 在内部构建三角网格。
|
||||
|
||||
## BEN2 ONNX
|
||||
|
||||
BEN2 输出 `foreground.png`，RGB 保留原始输入，Alpha 来自 ONNX mask 并恢复到输入分辨率。
|
||||
|
||||
125
docs/service.md
125
docs/service.md
@@ -1,98 +1,61 @@
|
||||
# AI Service
|
||||
|
||||
## AI Service
|
||||
## 运行架构
|
||||
|
||||
Environment 指磁盘上已安装的 Python、依赖及模型执行组件；Service 是由 Environment 启动、接收 Blender 任务的
|
||||
常驻 HTTP 进程。`manifest.json` 记录当前 Environment 的安装信息；Service 的 `/health` 返回进程状态和
|
||||
实例标识。项目发布前不维护安装清单或任务协议的兼容版本号。
|
||||
SnapForm 在独立的 Python 3.12 Environment 中运行常驻 HTTP Service，Blender 只通过标准库 `urllib.request` 提交任务，不把 AI 运行库载入 Blender Python。Service 只监听 `127.0.0.1`，请求使用随机 Token；同一 Storage Root 由 `service.lock` 保证只有一个实例。
|
||||
|
||||
Extension 本身不包含 Torch。`Install Environment` 使用 `uv` 创建独立 Python 3.12
|
||||
环境，并安装。机器上即使没有安装 `uv` 或 Python 也没关系：安装程序由 Blender
|
||||
自带的 Python 启动，缺 `uv` 时会解包扩展内置的对应平台二进制，再用它下载
|
||||
Python 3.12 与全部依赖。`uv` 会缓存在扩展存储目录下，后续安装直接复用。
|
||||
安装不做网络预检，直接开始；环境变量仍可切换镜像与代理：`PIP_INDEX_URL`（PyPI 镜像）、
|
||||
`SNAPFORM_TORCH_INDEX`（PyTorch 索引）、`HF_ENDPOINT`（Hugging Face 镜像）、
|
||||
`HTTP(S)_PROXY`/`ALL_PROXY`（代理）。未指定 PyPI 镜像时，
|
||||
pip 安装失败会自动依次重试镜像（清华、阿里云）；模型下载不预探测，官方 Hugging Face
|
||||
失败后自动回退 OmooLab R2（`models.omoolab.xyz`），并校验文件大小和 SHA-256。
|
||||
GitHub 不再参与安装：DA3、BEN2、uv 已内置，缺少内置文件会直接报错提示重新安装，不再下载兜底。
|
||||
单独下载模型时遵循同样的顺序：Hugging Face（可用 `HF_ENDPOINT` 指定其他入口）失败后自动切换到
|
||||
OmooLab R2，不再使用 ModelScope。
|
||||
Windows 上代理若只在系统设置里而未写入环境变量，安装程序会尝试自动从系统代理注入子进程。
|
||||
|
||||
- PyTorch 2.9.1
|
||||
- TorchVision 0.24.1
|
||||
- xFormers 0.0.33.post2（与 PyTorch 2.9.1 对应，仅 Windows/Linux；macOS 无 wheel 自动跳过）
|
||||
- 固定到已验证 commit 的 Depth Anything 3
|
||||
- 固定到已验证 commit 的 BEN2
|
||||
- FastAPI 0.139.2 与 Uvicorn 0.51.0
|
||||
- DA3 与 BEN2 源码包随扩展打包，安装时复制到临时目录后由 uv 本地安装，完成后立即清理
|
||||
（无需访问 GitHub）；SnapForm 直接调用 DA3
|
||||
底层 export 接口，`gs_ply` 导出不会触发 `gs_video` 渲染，因此不需要 gsplat
|
||||
与 C++ 编译工具链
|
||||
- `vendor/` 是不进入 Git 的本地缓存；`vendor.lock.toml` 固定下载 URL、commit/version 与
|
||||
SHA-256，`patches/<包名>/` 进入 Git 并存放 SnapForm 对源码的覆盖文件（移除 xformers
|
||||
硬依赖、export 入口改懒加载、
|
||||
射线位姿的 3×3 QR 分解在 MPS 上改走 CPU），
|
||||
`pack.py` 打包时自动套入官方包，扩展 zip 内才是修补后的版本；更新上游只需替换官方包
|
||||
并核对覆盖文件是否仍适用
|
||||
- `addict` 2.4.0（补充 DA3 上游包元数据暂未声明的运行依赖）
|
||||
|
||||
安装 Environment 时会自动选择 PyTorch 后端：检测到 NVIDIA 驱动则安装 CUDA 12.8 版本，否则安装
|
||||
CPU 版；macOS 始终安装 CPU 版 PyTorch（MPS 包含其中）。CPU 可以用于验证，但实际推理通常很慢。
|
||||
Extension Preferences 的 `Device` 决定正式功能的推理设备，默认取本机最快的 CUDA/MPS，否则
|
||||
CPU；Debug 面板的 `Debug Device` 只影响 Debug 任务。两处设备选项都只列出本机支持的设备：
|
||||
Windows 不显示 MPS，机器上没有可用 NVIDIA 驱动时不显示 CUDA。Environment 安装完成时会检测真实
|
||||
可用的设备并写入 `manifest.json`，设备列表以该检测结果为准。
|
||||
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库，无需安装 CUDA Toolkit；只需要
|
||||
NVIDIA 驱动 ≥ 570（数据中心显卡可用更老驱动）。驱动不满足或没有 NVIDIA 显卡时，
|
||||
设备会自动回退到 CPU（macOS 为 MPS），只有显式选择 CUDA 才会报错。
|
||||
|
||||
macOS（Apple Silicon）会自动从 PyPI 安装 CPU 版 PyTorch，推理设备默认使用 MPS
|
||||
（Apple Metal）加速，也可以在设备选项里切换 CPU。PyTorch 自 2.3 起不再提供 Intel
|
||||
Mac（x86_64）轮子，Intel Mac 暂不支持。DA3 射线位姿估计中的 3×3 QR 分解在 MPS
|
||||
上尚未实现，SnapForm 的覆盖文件会单独把该分解放到 CPU 执行，其余推理仍走 MPS。
|
||||
|
||||
Extension 启用后只会查找已有的共享 Service，不会主动启动新进程。第一次执行 AI 操作或在侧边栏
|
||||
手动点击 `Start` 时才会启动 Service。Service 只绑定回环地址，并让操作系统分配动态端口；第一个实例把
|
||||
端口、进程和随机 Token 写入 `service.json`，其余 Blender 通过标准库
|
||||
`urllib.request` 连接同一个 Service。`service.lock` 保证同一 Storage Root 同时只有一个 Service，Token
|
||||
则阻止同机其他进程直接调用接口。Service 全局同时只执行一个 AI 任务。
|
||||
|
||||
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时先注销自身；如果当前 Service
|
||||
就是由这个 Blender 启动的，会直接终止它，无需等待。异常退出的客户端会在 10 秒后过期；Service
|
||||
在最后一个客户端离开后仍会保留模型缓存等待 20 秒再退出，作为 Service 由其他进程启动等场景的兜底。
|
||||
取消任务不会重启 Service，并会保留已经加载的模型；
|
||||
如果取消发生在一次不可中断的 Torch 调用中，任务会在该调用返回后的最近检查点结束。
|
||||
|
||||
侧边栏 `AI Service` 面板以 `Running (Clients: n)` 表示空闲服务，以 `Running (Busy)` 表示正在执行任务，
|
||||
并显示缓存模型和活动任务。Blender 完成结果导入时会立即释放当前任务对应的本地 Busy 快照，无需等待
|
||||
下一次 Service 轮询。空闲且只有当前 Blender 连接时，可以清理模型缓存、重启或停止 Service；
|
||||
多客户端或任务执行期间不会允许破坏共享服务。
|
||||
|
||||
BEN2 在首次使用后保持加载；DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型；加载时如果
|
||||
显存不足，会先释放另一类缓存模型并重试。
|
||||
|
||||
Environment Manifest Revision 变化时，旧 Environment 需要重新执行 `Install Environment`。
|
||||
Service 的启动和崩溃日志写入 `~/.snapform/service.log`。
|
||||
|
||||
Environment、模型和任务结果默认保存在用户主目录的 `~/.snapform`：
|
||||
Environment、模型和任务结果默认保存在 `~/.snapform`：
|
||||
|
||||
```text
|
||||
~/.snapform/
|
||||
├── .venv/
|
||||
├── models/
|
||||
├── jobs/
|
||||
├── manifest.json
|
||||
├── service.log
|
||||
├── service.json # Service 运行时存在
|
||||
└── service.lock # Service 运行时持有
|
||||
```
|
||||
|
||||
`models/` 下每个模型一个文件夹，命名统一带前缀：DA3 为 `da3_*`（如 `da3_small`、`da3_base`），
|
||||
BEN2 为 `ben2_*`（如 `ben2_base`）。旧版无前缀的 DA3 文件夹会在首次使用时自动改名为新名称。
|
||||
不同 Blender 版本可共享此目录。也可以在 Extension Preferences 的 `Storage Root` 中修改位置。
|
||||
|
||||
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
|
||||
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录，
|
||||
因为该目录可能只读，并且升级时会被替换。不要同时从多个 Blender 实例执行
|
||||
`Install Environment`；Service、任务队列和模型读取可以共享。
|
||||
## ONNX-only Environment
|
||||
|
||||
扩展只内置对应平台的 `uv` 可执行文件。`Install Environment` 使用它创建环境并安装：
|
||||
|
||||
- ONNX Runtime 1.24.4；Windows 使用 `onnxruntime-directml`
|
||||
- NumPy 2.4.2、Pillow 12.1.1、OpenCV Headless 4.13.0.92、trimesh 4.11.2
|
||||
- FastAPI 0.139.2、Uvicorn 0.51.0、huggingface-hub 1.4.1
|
||||
|
||||
运行环境不安装 PyTorch、TorchVision、xFormers，也不安装 DA3、MoGe、BEN2 或 utils3d 的 Python 源码包。模型推理和 SnapForm 所需的相机、点云、网格后处理均由 ONNX Runtime、NumPy 与内部实现完成。
|
||||
|
||||
Windows 当前发行环境提供 DirectML 与 CPU；`Auto` 优先使用 `DmlExecutionProvider`。显式指定不可用设备时会安全回退到 CPU。其他平台安装标准 ONNX Runtime，并按可用 provider 选择 CoreML、OpenVINO 或 CPU。
|
||||
|
||||
`manifest.json` 记录 Environment 版本、ONNX Runtime 版本、runtime variant 与设备列表。Manifest Revision 改变后需要重新执行 `Install Environment`。
|
||||
|
||||
## 模型
|
||||
|
||||
所有正式功能都使用 ONNX：
|
||||
|
||||
- MoGe-2 ViT-S、ViT-B、ViT-L Normal
|
||||
- DA3 Small、DA3 Base
|
||||
- BEN2 Base
|
||||
|
||||
首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`，其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`，并保留仓库内部结构。下载顺序为 Hugging Face，然后回退到 `https://models.omoolab.xyz/<repository-name>/<relative-path>`；R2 文件会校验大小与 SHA-256。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。
|
||||
|
||||
DA3 ONNX 使用模型的 camera decoder 恢复外参与内参；旧版 `Use Ray Pose` 参数只为任务协议兼容保留，不参与 ONNX 推理。MoGe-2 的 metric depth、normal、mask、normalized intrinsics，以及 BEN2 的 Alpha，均在 NumPy 后处理中恢复。
|
||||
|
||||
|
||||
## 安装网络与镜像
|
||||
|
||||
安装器直接开始工作，不做阻塞式网络预检。依赖安装失败时会依次重试清华与阿里云 PyPI 镜像，也可用 `PIP_INDEX_URL` 或 Preferences 中的 `PyPI Index` 显式指定。Windows 系统代理会注入子进程；同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。
|
||||
|
||||
GitHub 不参与用户安装流程。`prepare_vendor.py` 只为扩展打包准备各平台的 `uv` 二进制。
|
||||
|
||||
## 生命周期
|
||||
|
||||
扩展启用后只查找已有 Service；第一次执行 AI 操作或点击 `Start` 时才启动。每个 Blender 客户端每 5 秒发送心跳，正常关闭时注销。最后一个客户端离开后，Service 保留模型缓存 20 秒再退出。
|
||||
|
||||
Service 全局同时执行一个 AI 任务。BEN2 独立缓存，深度模型只保留最近使用的一个；切换模型时释放旧会话。取消任务不会重启 Service，ONNX Runtime 当前调用返回后在最近检查点结束。
|
||||
|
||||
启动与崩溃信息写入 `~/.snapform/service.log`。
|
||||
|
||||
@@ -25,7 +25,7 @@
|
||||
- Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个
|
||||
`Smooth` 程度系数；节点组不再额外执行 Flip Faces；
|
||||
- 从修改前的原始表面以 nearest-surface 重新采样 `UVMap`，最终明确写为 Corner 域 `2D Vector`。
|
||||
MoGe-2 使用独立的 `O MoGe2 Surface`。Service 由 `utils3d.np.build_mesh_from_map(..., tri=True)`
|
||||
MoGe-2 使用独立的 `O MoGe2 Surface`。Service 由内置 NumPy 网格构建函数
|
||||
生成有限深度范围内的紧凑三角网格，并保留模型 mask、逐点 depth 及相邻点相对深度差。Geometry Nodes
|
||||
按官方 `mask & ~depth_map_edge(depth, rtol=0.04)` 的含义，先由 `Mask` 删除 mask 外的面，再由
|
||||
`Depth Continuity` 按同一份 3×3 邻域相对深度差删除断层面。DA3 与 MoGe-2 使用相同的属性、
|
||||
@@ -163,7 +163,7 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN
|
||||
浮点深度；保存 `.blend` 后再次打开仍需对应 job 目录存在。
|
||||
|
||||
选择 PLY 后，模型菜单只会提供 `DA3 Giant 1.1` 和
|
||||
`DA3 Nested Giant Large 1.1`。Small、Base 和 Large 没有 Gaussian head，无法由
|
||||
`DA3 Nested Giant Large 1.1`。Small 和 Base 没有 Gaussian head，无法由
|
||||
DA3 生成 `gs_ply`。
|
||||
|
||||
每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹。文件夹统一采用
|
||||
@@ -185,7 +185,7 @@ Blender 为 Packed Image 或序列帧生成的输入副本也归档到对应 job
|
||||
- `remove-image-background`
|
||||
- `generate-original-depth-plane`
|
||||
|
||||
DA3 官方 exporter 产生的 `gs_ply/`、`depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺，
|
||||
DA3 ONNX exporter 产生的 `depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺，
|
||||
不会保留只含一层产物的子目录：
|
||||
|
||||
```text
|
||||
@@ -218,6 +218,4 @@ remove-background_20260807-153000-123456/
|
||||
└── depth.0001.exr… AI 深度序列帧
|
||||
```
|
||||
|
||||
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重；`DA3 Large 1.1` 权重采用
|
||||
CC BY-NC 4.0，仅限非商业用途。
|
||||
|
||||
建议从 `DA3 Base` 开始。它使用 Apache 2.0 模型权重。
|
||||
|
||||
63
pack.py
63
pack.py
@@ -1,6 +1,4 @@
|
||||
import hashlib
|
||||
import io
|
||||
import tarfile
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
from zipfile import ZIP_DEFLATED, ZipFile
|
||||
@@ -10,7 +8,6 @@ PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
|
||||
VENDOR_DIR = PROJECT_ROOT / "vendor"
|
||||
DIST_DIR = PROJECT_ROOT / "dist"
|
||||
PATCHES_DIR = PROJECT_ROOT / "patches"
|
||||
LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml"
|
||||
|
||||
|
||||
@@ -28,15 +25,6 @@ PLATFORMS = {
|
||||
platform: artifact["filename"]
|
||||
for platform, artifact in VENDOR_LOCK["platforms"].items()
|
||||
}
|
||||
BUNDLED_ARCHIVES = {
|
||||
f"vendor/{artifact['filename']}": artifact["filename"]
|
||||
for artifact in VENDOR_LOCK["artifacts"].values()
|
||||
}
|
||||
OVERLAYS = {
|
||||
artifact["filename"]: artifact["patches"]
|
||||
for artifact in VENDOR_LOCK["artifacts"].values()
|
||||
if "patches" in artifact
|
||||
}
|
||||
|
||||
|
||||
def file_sha256(path):
|
||||
@@ -63,41 +51,6 @@ def validate_artifact(path, expected_sha256):
|
||||
)
|
||||
|
||||
|
||||
def bundled_artifact_bytes(filename):
|
||||
"""Return a bundled artifact with its configured overlay applied."""
|
||||
archive_path = VENDOR_DIR / filename
|
||||
overlay_dir = PATCHES_DIR / OVERLAYS.get(filename, filename)
|
||||
if not overlay_dir.is_dir():
|
||||
return archive_path.read_bytes()
|
||||
if not filename.endswith(".tar.gz"):
|
||||
raise ValueError(f"Unsupported overlay archive: {filename}")
|
||||
return _apply_tar_overlay(archive_path, overlay_dir)
|
||||
|
||||
|
||||
def _apply_tar_overlay(archive_path, overlay_dir):
|
||||
with tarfile.open(archive_path, "r:gz") as source:
|
||||
members = source.getmembers()
|
||||
top_level = members[0].name.split("/", 1)[0]
|
||||
buffer = io.BytesIO()
|
||||
with tarfile.open(fileobj=buffer, mode="w:gz") as output:
|
||||
for member in members:
|
||||
if member.isdir() or member.islnk() or member.issym():
|
||||
output.addfile(member)
|
||||
continue
|
||||
relative = member.name[len(top_level) + 1:]
|
||||
overlay_file = overlay_dir / relative
|
||||
if overlay_file.is_file():
|
||||
info = tarfile.TarInfo(member.name)
|
||||
info.size = overlay_file.stat().st_size
|
||||
info.mode = member.mode
|
||||
with overlay_file.open("rb") as overlay_stream:
|
||||
output.addfile(info, overlay_stream)
|
||||
else:
|
||||
with source.extractfile(member) as extracted:
|
||||
output.addfile(member, extracted)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def manifest_for_platform(platform):
|
||||
"""Return the manifest text declaring only the given platform."""
|
||||
text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
|
||||
@@ -121,23 +74,13 @@ def manifest_for_platform(platform):
|
||||
|
||||
|
||||
def bundled_vendor_files(platform):
|
||||
entries = {
|
||||
name: VENDOR_DIR / filename
|
||||
for name, filename in BUNDLED_ARCHIVES.items()
|
||||
return {
|
||||
f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform]
|
||||
}
|
||||
entries[f"vendor/{PLATFORMS[platform]}"] = (
|
||||
VENDOR_DIR / PLATFORMS[platform]
|
||||
)
|
||||
return entries
|
||||
|
||||
|
||||
def validate_vendor_files(platform):
|
||||
"""Validate common archives and the selected uv archive."""
|
||||
for artifact in VENDOR_LOCK["artifacts"].values():
|
||||
validate_artifact(
|
||||
VENDOR_DIR / artifact["filename"],
|
||||
artifact["sha256"],
|
||||
)
|
||||
uv_artifact = VENDOR_LOCK["platforms"][platform]
|
||||
validate_artifact(
|
||||
VENDOR_DIR / uv_artifact["filename"],
|
||||
@@ -174,7 +117,7 @@ def build_extension(platform):
|
||||
for archive_name, source_path in vendor_files.items():
|
||||
archive.writestr(
|
||||
archive_name,
|
||||
bundled_artifact_bytes(source_path.name),
|
||||
source_path.read_bytes(),
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
@@ -1,93 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["hatchling>=1.25", "hatch-vcs>=0.4"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "depth-anything-3"
|
||||
version = "0.0.0"
|
||||
description = "Depth Anything 3"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9, <=3.13"
|
||||
license = { text = "Apache-2.0" }
|
||||
authors = [{ name = "Your Name" }]
|
||||
|
||||
dependencies = [
|
||||
"pre-commit",
|
||||
"trimesh",
|
||||
"torch>=2",
|
||||
"torchvision",
|
||||
"einops",
|
||||
"huggingface_hub",
|
||||
"imageio",
|
||||
"numpy<2",
|
||||
"opencv-python",
|
||||
"open3d",
|
||||
"fastapi",
|
||||
"uvicorn",
|
||||
"requests",
|
||||
"typer",
|
||||
"pillow",
|
||||
"omegaconf",
|
||||
"evo",
|
||||
"e3nn",
|
||||
"moviepy",
|
||||
"plyfile",
|
||||
"pillow_heif",
|
||||
"safetensors",
|
||||
"uvicorn",
|
||||
"moviepy==1.0.3",
|
||||
"typer>=0.9.0",
|
||||
"pycolmap",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
app = ["gradio>=5", "pillow>=9.0"] # requires that python3>=3.10
|
||||
gs = ["gsplat @ git+https://github.com/nerfstudio-project/gsplat.git@0b4dddf04cb687367602c01196913cde6a743d70"]
|
||||
all = ["depth-anything-3[app,gs]"]
|
||||
|
||||
|
||||
[project.scripts]
|
||||
da3 = "depth_anything_3.cli:app"
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/ByteDance-Seed/Depth-Anything-3"
|
||||
|
||||
[tool.hatch.version]
|
||||
source = "vcs"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/depth_anything_3"]
|
||||
|
||||
[tool.hatch.build.targets.sdist]
|
||||
include = [
|
||||
"/README.md",
|
||||
"/pyproject.toml",
|
||||
"/src/depth_anything_3",
|
||||
]
|
||||
|
||||
[tool.hatch.metadata]
|
||||
allow-direct-references = true
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ["jaxtyping.mypy_plugin"]
|
||||
|
||||
[tool.black]
|
||||
line-length = 99
|
||||
target-version = ['py37', 'py38', 'py39', 'py310', 'py311']
|
||||
include = '\.pyi?$'
|
||||
exclude = '''
|
||||
/(
|
||||
| \.git
|
||||
)/
|
||||
'''
|
||||
|
||||
[tool.isort]
|
||||
profile = "black"
|
||||
multi_line_output = 3
|
||||
include_trailing_comma = true
|
||||
known_third_party = ["bson","cruise","cv2","dataloader","diffusers","omegaconf","tensorflow","torch","torchvision","transformers","gsplat"]
|
||||
known_first_party = ["common", "data", "models", "projects"]
|
||||
sections = ["FUTURE","STDLIB","THIRDPARTY","FIRSTPARTY","LOCALFOLDER"]
|
||||
skip_gitignore = true
|
||||
line_length = 99
|
||||
no_lines_before="THIRDPARTY"
|
||||
@@ -1,23 +0,0 @@
|
||||
pre-commit
|
||||
trimesh
|
||||
torch>=2
|
||||
torchvision
|
||||
einops
|
||||
huggingface_hub
|
||||
imageio
|
||||
numpy<2
|
||||
opencv-python
|
||||
open3d
|
||||
fastapi
|
||||
uvicorn
|
||||
requests
|
||||
typer
|
||||
pillow
|
||||
omegaconf
|
||||
evo
|
||||
e3nn
|
||||
moviepy
|
||||
plyfile
|
||||
pillow_heif
|
||||
safetensors
|
||||
pycolmap
|
||||
@@ -1,84 +0,0 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from depth_anything_3.specs import Prediction
|
||||
|
||||
# Canonical set of export formats this dispatcher understands. Treat this as the
|
||||
# server-trusted enum for any request-facing validation (e.g. the HTTP backend) instead
|
||||
# of letting a client-supplied export_format string reach the dispatcher unchecked.
|
||||
SUPPORTED_EXPORT_FORMATS = frozenset(
|
||||
{
|
||||
"glb",
|
||||
"mini_npz",
|
||||
"npz",
|
||||
"feat_vis",
|
||||
"depth_vis",
|
||||
"gs_ply",
|
||||
"gs_video",
|
||||
"colmap",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def export(
|
||||
prediction: Prediction,
|
||||
export_format: str,
|
||||
export_dir: str,
|
||||
**kwargs,
|
||||
):
|
||||
if "-" in export_format:
|
||||
export_formats = export_format.split("-")
|
||||
for export_format in export_formats:
|
||||
export(prediction, export_format, export_dir, **kwargs)
|
||||
return # Prevent falling through to single-format handling
|
||||
|
||||
if export_format == "glb":
|
||||
from .glb import export_to_glb
|
||||
|
||||
export_to_glb(prediction, export_dir, **kwargs.get(export_format, {}))
|
||||
elif export_format == "mini_npz":
|
||||
from .npz import export_to_mini_npz
|
||||
|
||||
export_to_mini_npz(prediction, export_dir)
|
||||
elif export_format == "npz":
|
||||
from .npz import export_to_npz
|
||||
|
||||
export_to_npz(prediction, export_dir)
|
||||
elif export_format == "feat_vis":
|
||||
from .feat_vis import export_to_feat_vis
|
||||
|
||||
export_to_feat_vis(prediction, export_dir, **kwargs.get(export_format, {}))
|
||||
elif export_format == "depth_vis":
|
||||
from .depth_vis import export_to_depth_vis
|
||||
|
||||
export_to_depth_vis(prediction, export_dir)
|
||||
elif export_format == "gs_ply":
|
||||
from .gs import export_to_gs_ply
|
||||
|
||||
export_to_gs_ply(prediction, export_dir, **kwargs.get(export_format, {}))
|
||||
elif export_format == "gs_video":
|
||||
from .gs import export_to_gs_video
|
||||
|
||||
export_to_gs_video(prediction, export_dir, **kwargs.get(export_format, {}))
|
||||
elif export_format == "colmap":
|
||||
from .colmap import export_to_colmap
|
||||
|
||||
export_to_colmap(prediction, export_dir, **kwargs.get(export_format, {}))
|
||||
else:
|
||||
raise ValueError(f"Unsupported export format: {export_format}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
export,
|
||||
]
|
||||
@@ -1,531 +0,0 @@
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import torch
|
||||
from einops import repeat
|
||||
from .geometry import unproject_depth
|
||||
|
||||
|
||||
def compute_optimal_rotation_intrinsics_batch(
|
||||
rays_origin, rays_target, z_threshold=1e-4, reproj_threshold=0.2, weights=None,
|
||||
n_sample = None,
|
||||
n_iter=100,
|
||||
num_sample_for_ransac=8,
|
||||
rand_sample_iters_idx=None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
rays_origin (torch.Tensor): (B, N, 3)
|
||||
rays_target (torch.Tensor): (B, N, 3)
|
||||
z_threshold (float): Threshold for z value to be considered valid.
|
||||
|
||||
Returns:
|
||||
R (torch.tensor): (3, 3)
|
||||
focal_length (torch.tensor): (2,)
|
||||
principal_point (torch.tensor): (2,)
|
||||
"""
|
||||
device = rays_origin.device
|
||||
B, N, _ = rays_origin.shape
|
||||
z_mask = torch.logical_and(
|
||||
torch.abs(rays_target[:, :, 2]) > z_threshold, torch.abs(rays_origin[:, :, 2]) > z_threshold
|
||||
) # (B, N, 1)
|
||||
rays_origin = rays_origin.clone()
|
||||
rays_target = rays_target.clone()
|
||||
rays_origin[:, :, 0][z_mask] /= rays_origin[:, :, 2][z_mask]
|
||||
rays_origin[:, :, 1][z_mask] /= rays_origin[:, :, 2][z_mask]
|
||||
rays_target[:, :, 0][z_mask] /= rays_target[:, :, 2][z_mask]
|
||||
rays_target[:, :, 1][z_mask] /= rays_target[:, :, 2][z_mask]
|
||||
|
||||
rays_origin = rays_origin[:, :, :2]
|
||||
rays_target = rays_target[:, :, :2]
|
||||
assert weights is not None, "weights must be provided"
|
||||
weights[~z_mask] = 0
|
||||
|
||||
A_list = []
|
||||
max_chunk_size = 2
|
||||
for i in range(0, rays_origin.shape[0], max_chunk_size):
|
||||
A = ransac_find_homography_weighted_fast_batch(
|
||||
rays_origin[i:i+max_chunk_size],
|
||||
rays_target[i:i+max_chunk_size],
|
||||
weights[i:i+max_chunk_size],
|
||||
n_iter=n_iter,
|
||||
n_sample = n_sample,
|
||||
num_sample_for_ransac=num_sample_for_ransac,
|
||||
reproj_threshold=reproj_threshold,
|
||||
rand_sample_iters_idx=rand_sample_iters_idx,
|
||||
max_inlier_num=8000,
|
||||
)
|
||||
A = A.to(device)
|
||||
A_need_inv_mask = torch.linalg.det(A) < 0
|
||||
A[A_need_inv_mask] = -A[A_need_inv_mask]
|
||||
A_list.append(A)
|
||||
|
||||
A = torch.cat(A_list, dim=0)
|
||||
|
||||
R_list = []
|
||||
f_list = []
|
||||
pp_list = []
|
||||
for i in range(A.shape[0]):
|
||||
R, L = ql_decomposition(A[i])
|
||||
L = L / L[2][2]
|
||||
|
||||
f = torch.stack((L[0][0], L[1][1]))
|
||||
pp = torch.stack((L[2][0], L[2][1]))
|
||||
R_list.append(R)
|
||||
f_list.append(f)
|
||||
pp_list.append(pp)
|
||||
|
||||
R = torch.stack(R_list)
|
||||
f = torch.stack(f_list)
|
||||
pp = torch.stack(pp_list)
|
||||
|
||||
return R, f, pp
|
||||
|
||||
|
||||
# https://www.reddit.com/r/learnmath/comments/v1crd7/linear_algebra_qr_to_ql_decomposition/
|
||||
def ql_decomposition(A):
|
||||
device = A.device
|
||||
P = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 0, 0]], device=device).float()
|
||||
A_tilde = torch.matmul(A, P)
|
||||
# MPS 尚未实现 torch.linalg.qr，仅把 3x3 分解放到 CPU 执行后移回
|
||||
use_cpu = device.type == "mps"
|
||||
if use_cpu:
|
||||
A_tilde = A_tilde.cpu()
|
||||
Q_tilde, R_tilde = torch.linalg.qr(A_tilde)
|
||||
if use_cpu:
|
||||
Q_tilde = Q_tilde.to(device=device)
|
||||
R_tilde = R_tilde.to(device=device)
|
||||
Q = torch.matmul(Q_tilde, P)
|
||||
L = torch.matmul(torch.matmul(P, R_tilde), P)
|
||||
d = torch.diag(L)
|
||||
Q[:, 0] *= torch.sign(d[0])
|
||||
Q[:, 1] *= torch.sign(d[1])
|
||||
Q[:, 2] *= torch.sign(d[2])
|
||||
L[0] *= torch.sign(d[0])
|
||||
L[1] *= torch.sign(d[1])
|
||||
L[2] *= torch.sign(d[2])
|
||||
return Q, L
|
||||
|
||||
def find_homography_least_squares_weighted_torch(src_pts, dst_pts, confident_weight):
|
||||
"""
|
||||
src_pts: (N,2) source points (torch.Tensor, float32/float64)
|
||||
dst_pts: (N,2) target points (torch.Tensor, float32/float64)
|
||||
confident_weight: (N,) weights (torch.Tensor)
|
||||
Returns: (3,3) homography matrix H (torch.Tensor)
|
||||
"""
|
||||
assert src_pts.shape == dst_pts.shape
|
||||
N = src_pts.shape[0]
|
||||
if N < 4:
|
||||
raise ValueError("At least 4 points are required to compute homography.")
|
||||
assert confident_weight.shape == (N,)
|
||||
|
||||
w = confident_weight.sqrt().unsqueeze(1) # (N,1)
|
||||
|
||||
x = src_pts[:, 0:1] # (N,1)
|
||||
y = src_pts[:, 1:2] # (N,1)
|
||||
u = dst_pts[:, 0:1]
|
||||
v = dst_pts[:, 1:2]
|
||||
|
||||
zeros = torch.zeros_like(x)
|
||||
|
||||
# Construct A matrix (2N, 9)
|
||||
A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=1)
|
||||
A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=1)
|
||||
A = torch.cat([A1, A2], dim=0) # (2N, 9)
|
||||
|
||||
# SVD
|
||||
# Note: torch.linalg.svd returns U, S, Vh, where Vh is the transpose of V
|
||||
_, _, Vh = torch.linalg.svd(A)
|
||||
H = Vh[-1].reshape(3, 3)
|
||||
H = H / H[-1, -1]
|
||||
return H
|
||||
|
||||
|
||||
def ransac_find_homography_weighted(
|
||||
src_pts,
|
||||
dst_pts,
|
||||
confident_weight,
|
||||
n_iter=100,
|
||||
sample_ratio=0.2,
|
||||
reproj_threshold=3.0,
|
||||
num_sample_for_ransac=16,
|
||||
random_seed=None,
|
||||
):
|
||||
"""
|
||||
RANSAC version of weighted Homography estimation.
|
||||
Sample 4 points from the top 50% weighted points each time.
|
||||
reproj_threshold: points with reprojection error less than this value are inliers
|
||||
Returns: best_H
|
||||
"""
|
||||
if random_seed is not None:
|
||||
torch.manual_seed(random_seed)
|
||||
N = src_pts.shape[0]
|
||||
assert N >= 4
|
||||
# 1. Select top 50% weighted points
|
||||
sorted_idx = torch.argsort(confident_weight, descending=True)
|
||||
n_sample = max(num_sample_for_ransac, int(N * sample_ratio))
|
||||
candidate_idx = sorted_idx[:n_sample]
|
||||
best_inlier_mask = None
|
||||
best_score = 0
|
||||
for _ in range(n_iter):
|
||||
# 2. Randomly sample 4 points
|
||||
idx = candidate_idx[torch.randperm(n_sample)[:num_sample_for_ransac]]
|
||||
# 3. Compute Homography
|
||||
try:
|
||||
H = find_homography_least_squares_weighted_torch(
|
||||
src_pts[idx], dst_pts[idx], confident_weight[idx]
|
||||
)
|
||||
except Exception:
|
||||
H = torch.eye(3, dtype=src_pts.dtype, device=src_pts.device)
|
||||
# 4. Compute reprojection error for all points
|
||||
src_homo = torch.cat(
|
||||
[src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1
|
||||
)
|
||||
proj = (H @ src_homo.T).T
|
||||
proj = proj[:, :2] / proj[:, 2:3]
|
||||
error = ((proj - dst_pts) ** 2).sum(dim=1).sqrt() # Euclidean distance
|
||||
inlier_mask = error < reproj_threshold
|
||||
total_score = (inlier_mask * confident_weight).sum().item()
|
||||
n_inlier = inlier_mask.sum().item()
|
||||
if n_inlier < 4:
|
||||
continue # At least 4 inliers required for fitting
|
||||
|
||||
if total_score > best_score:
|
||||
best_score = total_score
|
||||
best_inlier_mask = inlier_mask
|
||||
|
||||
# 5. Refit Homography using inliers
|
||||
H_inlier = find_homography_least_squares_weighted_torch(
|
||||
src_pts[best_inlier_mask], dst_pts[best_inlier_mask], confident_weight[best_inlier_mask]
|
||||
)
|
||||
|
||||
return H_inlier
|
||||
|
||||
|
||||
def find_homography_least_squares_weighted_torch_batch(
|
||||
src_pts_batch, dst_pts_batch, confident_weight_batch
|
||||
):
|
||||
"""
|
||||
Batch version of weighted least squares Homography
|
||||
src_pts_batch: (B, K, 2)
|
||||
dst_pts_batch: (B, K, 2)
|
||||
confident_weight_batch: (B, K)
|
||||
Returns: (B, 3, 3)
|
||||
"""
|
||||
B, K, _ = src_pts_batch.shape
|
||||
w = confident_weight_batch.sqrt().unsqueeze(2) # (B,K,1)
|
||||
x = src_pts_batch[:, :, 0:1]
|
||||
y = src_pts_batch[:, :, 1:2]
|
||||
u = dst_pts_batch[:, :, 0:1]
|
||||
v = dst_pts_batch[:, :, 1:2]
|
||||
zeros = torch.zeros_like(x)
|
||||
A1 = torch.cat([-x * w, -y * w, -w, zeros, zeros, zeros, x * u * w, y * u * w, u * w], dim=2)
|
||||
A2 = torch.cat([zeros, zeros, zeros, -x * w, -y * w, -w, x * v * w, y * v * w, v * w], dim=2)
|
||||
A = torch.cat([A1, A2], dim=1) # (B, 2K, 9)
|
||||
# SVD: torch.linalg.svd supports batch
|
||||
_, _, Vh = torch.linalg.svd(A)
|
||||
H = Vh[:, -1].reshape(B, 3, 3)
|
||||
H = H / H[:, 2:3, 2:3]
|
||||
return H
|
||||
|
||||
|
||||
def ransac_find_homography_weighted_fast(
|
||||
src_pts,
|
||||
dst_pts,
|
||||
confident_weight,
|
||||
n_sample,
|
||||
n_iter=100,
|
||||
reproj_threshold=3.0,
|
||||
num_sample_for_ransac=8,
|
||||
random_seed=None,
|
||||
rand_sample_iters_idx=None,
|
||||
):
|
||||
"""
|
||||
Batch version of RANSAC weighted Homography estimation.
|
||||
Returns: H_inlier
|
||||
"""
|
||||
if random_seed is not None:
|
||||
torch.manual_seed(random_seed)
|
||||
N = src_pts.shape[0]
|
||||
device = src_pts.device
|
||||
assert N >= 4
|
||||
# 1. Select top weighted points by sample_ratio
|
||||
sorted_idx = torch.argsort(confident_weight, descending=True)
|
||||
candidate_idx = sorted_idx[:n_sample] # (n_sample,)
|
||||
if rand_sample_iters_idx is None:
|
||||
rand_sample_iters_idx = torch.stack(
|
||||
[torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)],
|
||||
dim=0,
|
||||
) # (n_iter, num_sample_for_ransac)
|
||||
# 2. Generate all sampling groups at once
|
||||
# shape: (n_iter, num_sample_for_ransac)
|
||||
rand_idx = candidate_idx[rand_sample_iters_idx] # (n_iter, num_sample_for_ransac)
|
||||
# 3. Construct batch input
|
||||
src_pts_batch = src_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2)
|
||||
dst_pts_batch = dst_pts[rand_idx] # (n_iter, num_sample_for_ransac, 2)
|
||||
confident_weight_batch = confident_weight[rand_idx] # (n_iter, num_sample_for_ransac)
|
||||
# 4. Batch fit Homography
|
||||
H_batch = find_homography_least_squares_weighted_torch_batch(
|
||||
src_pts_batch, dst_pts_batch, confident_weight_batch
|
||||
) # (n_iter, 3, 3)
|
||||
# 5. Batch evaluate inliers for all H
|
||||
src_homo = torch.cat(
|
||||
[src_pts, torch.ones(N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=1
|
||||
) # (N,3)
|
||||
src_homo_expand = src_homo.unsqueeze(0).expand(n_iter, N, 3) # (n_iter, N, 3)
|
||||
dst_pts_expand = dst_pts.unsqueeze(0).expand(n_iter, N, 2) # (n_iter, N, 2)
|
||||
confident_weight_expand = confident_weight.unsqueeze(0).expand(n_iter, N) # (n_iter, N)
|
||||
# H_batch: (n_iter, 3, 3)
|
||||
proj = torch.bmm(src_homo_expand, H_batch.transpose(1, 2)) # (n_iter, N, 3)
|
||||
proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (n_iter, N, 2)
|
||||
error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=2).sqrt() # (n_iter, N)
|
||||
inlier_mask = error < reproj_threshold # (n_iter, N)
|
||||
total_score = (inlier_mask * confident_weight_expand).sum(dim=1) # (n_iter,)
|
||||
# 6. Select the sampling group with the highest score
|
||||
best_idx = torch.argmax(total_score)
|
||||
best_inlier_mask = inlier_mask[best_idx] # (N,)
|
||||
inlier_src_pts = src_pts[best_inlier_mask]
|
||||
inlier_dst_pts = dst_pts[best_inlier_mask]
|
||||
inlier_confident_weight = confident_weight[best_inlier_mask]
|
||||
|
||||
max_inlier_num = 10000
|
||||
sorted_idx = torch.argsort(inlier_confident_weight, descending=True)
|
||||
|
||||
# method 1: sort according to confident_weight, and only keep max_inlier_num pts
|
||||
# sorted_idx = sorted_idx[:max_inlier_num]
|
||||
|
||||
# method 2: random choose max_inlier_num pts
|
||||
sorted_idx = sorted_idx[torch.randperm(len(sorted_idx))[:max_inlier_num]]
|
||||
|
||||
inlier_src_pts = inlier_src_pts[sorted_idx]
|
||||
inlier_dst_pts = inlier_dst_pts[sorted_idx]
|
||||
inlier_confident_weight = inlier_confident_weight[sorted_idx]
|
||||
# 7. Refit Homography using inliers
|
||||
H_inlier = find_homography_least_squares_weighted_torch(
|
||||
inlier_src_pts, inlier_dst_pts, inlier_confident_weight
|
||||
)
|
||||
return H_inlier
|
||||
|
||||
|
||||
def ransac_find_homography_weighted_fast_batch(
|
||||
src_pts, # (B, N, 3)
|
||||
dst_pts, # (B, N, 2)
|
||||
confident_weight, # (B, N)
|
||||
n_sample,
|
||||
n_iter=100,
|
||||
reproj_threshold=3.0,
|
||||
num_sample_for_ransac=8,
|
||||
max_inlier_num=10000,
|
||||
random_seed=None,
|
||||
rand_sample_iters_idx=None,
|
||||
):
|
||||
"""
|
||||
Batch version of RANSAC weighted Homography estimation (supports batch).
|
||||
Input:
|
||||
src_pts: (B, N, 2)
|
||||
dst_pts: (B, N, 2)
|
||||
confident_weight: (B, N)
|
||||
Returns:
|
||||
H_inlier: (B, 3, 3)
|
||||
"""
|
||||
if random_seed is not None:
|
||||
torch.manual_seed(random_seed)
|
||||
B, N, _ = src_pts.shape
|
||||
assert N >= 4
|
||||
|
||||
device = src_pts.device
|
||||
|
||||
# 1. Select top weighted points by sample_ratio
|
||||
sorted_idx = torch.argsort(confident_weight, descending=True, dim=1) # (B, N)
|
||||
candidate_idx = sorted_idx[:, :n_sample] # (B, n_sample)
|
||||
|
||||
# 2. Generate all sampling groups at once
|
||||
# rand_idx: (B, n_iter, num_sample_for_ransac)
|
||||
if rand_sample_iters_idx is None:
|
||||
rand_sample_iters_idx = torch.stack(
|
||||
[torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)],
|
||||
dim=0,
|
||||
) # (n_iter, num_sample_for_ransac)
|
||||
|
||||
rand_idx = candidate_idx[:, rand_sample_iters_idx] # (B, n_iter, num_sample_for_ransac)
|
||||
|
||||
# 3. Construct batch input
|
||||
# Indexing method below: (B, n_iter, num_sample_for_ransac, ...)
|
||||
b_idx = torch.arange(B, device=device).view(B, 1, 1).expand(B, n_iter, num_sample_for_ransac)
|
||||
src_pts_batch = src_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2)
|
||||
dst_pts_batch = dst_pts[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac, 2)
|
||||
confident_weight_batch = confident_weight[b_idx, rand_idx] # (B, n_iter, num_sample_for_ransac)
|
||||
|
||||
# 4. Batch fit Homography
|
||||
# Need to implement batch version that supports (B, n_iter, num_sample_for_ransac, ...) input
|
||||
# Output H_batch: (B, n_iter, 3, 3)
|
||||
cB, cN = src_pts_batch.shape[:2]
|
||||
H_batch = find_homography_least_squares_weighted_torch_batch(
|
||||
src_pts_batch.flatten(0, 1), dst_pts_batch.flatten(0, 1), confident_weight_batch.flatten(0, 1)
|
||||
) # (B, n_iter, 3, 3)
|
||||
H_batch = H_batch.unflatten(0, (cB, cN))
|
||||
|
||||
# 5. Batch evaluate inliers for all H
|
||||
src_homo = torch.cat(
|
||||
[src_pts, torch.ones(B, N, 1, dtype=src_pts.dtype, device=src_pts.device)], dim=2
|
||||
) # (B, N, 3)
|
||||
src_homo_expand = src_homo.unsqueeze(1).expand(B, n_iter, N, 3) # (B, n_iter, N, 3)
|
||||
dst_pts_expand = dst_pts.unsqueeze(1).expand(B, n_iter, N, 2) # (B, n_iter, N, 2)
|
||||
confident_weight_expand = confident_weight.unsqueeze(1).expand(B, n_iter, N) # (B, n_iter, N)
|
||||
|
||||
# H_batch: (B, n_iter, 3, 3)
|
||||
# Need to reshape H_batch to (B*n_iter, 3, 3), src_homo_expand to (B*n_iter, N, 3)
|
||||
H_batch_flat = H_batch.reshape(-1, 3, 3)
|
||||
src_homo_expand_flat = src_homo_expand.reshape(-1, N, 3)
|
||||
proj = torch.bmm(src_homo_expand_flat, H_batch_flat.transpose(1, 2)) # (B*n_iter, N, 3)
|
||||
proj_xy = proj[:, :, :2] / proj[:, :, 2:3] # (B*n_iter, N, 2)
|
||||
proj_xy = proj_xy.reshape(B, n_iter, N, 2)
|
||||
error = ((proj_xy - dst_pts_expand) ** 2).sum(dim=3).sqrt() # (B, n_iter, N)
|
||||
inlier_mask = error < reproj_threshold # (B, n_iter, N)
|
||||
total_score = (inlier_mask * confident_weight_expand).sum(dim=2) # (B, n_iter)
|
||||
|
||||
# 6. Select the sampling group with the highest score
|
||||
best_idx = torch.argmax(total_score, dim=1) # (B,)
|
||||
best_inlier_mask = inlier_mask[torch.arange(B, device=device), best_idx] # (B, N)
|
||||
|
||||
# 7. Refit Homography using inliers
|
||||
H_inlier_list = []
|
||||
for b in range(B):
|
||||
mask = best_inlier_mask[b]
|
||||
inlier_src_pts = src_pts[b][mask] # (?, 3)
|
||||
inlier_dst_pts = dst_pts[b][mask] # (?, 2)
|
||||
inlier_confident_weight = confident_weight[b][mask] # (?)
|
||||
|
||||
sorted_idx = torch.argsort(inlier_confident_weight, descending=True)
|
||||
# # method 1: sort according to confident_weight, and only keep max_inlier_num pts
|
||||
# sorted_idx = sorted_idx[:max_inlier_num]
|
||||
# method 2: random choose max_inlier_num pts
|
||||
if len(sorted_idx) > max_inlier_num:
|
||||
# random choose from first 95% confident pts
|
||||
keep_len = max(int(len(sorted_idx) * 0.95), max_inlier_num)
|
||||
sorted_idx = sorted_idx[:keep_len]
|
||||
perm = torch.randperm(len(sorted_idx), device=device)[:max_inlier_num]
|
||||
sorted_idx = sorted_idx[perm]
|
||||
inlier_src_pts = inlier_src_pts[sorted_idx]
|
||||
inlier_dst_pts = inlier_dst_pts[sorted_idx]
|
||||
inlier_confident_weight = inlier_confident_weight[sorted_idx]
|
||||
|
||||
H_inlier = find_homography_least_squares_weighted_torch(
|
||||
inlier_src_pts, inlier_dst_pts, inlier_confident_weight
|
||||
) # (3, 3)
|
||||
H_inlier_list.append(H_inlier)
|
||||
H_inlier = torch.stack(H_inlier_list, dim=0) # (B, 3, 3)
|
||||
return H_inlier
|
||||
|
||||
def get_params_for_ransac(N, device):
|
||||
n_iter=100
|
||||
sample_ratio=0.3
|
||||
num_sample_for_ransac=8
|
||||
n_sample = max(num_sample_for_ransac, int(N * sample_ratio))
|
||||
rand_sample_iters_idx = torch.stack(
|
||||
[torch.randperm(n_sample, device=device)[:num_sample_for_ransac] for _ in range(n_iter)],
|
||||
dim=0,
|
||||
) # (n_iter, num_sample_for_ransac)
|
||||
return n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx
|
||||
|
||||
|
||||
def camray_to_caminfo(camray, confidence=None, reproj_threshold=0.2, training=False):
|
||||
"""
|
||||
Args:
|
||||
camray: (B, S, num_patches_y, num_patches_x, 6)
|
||||
confidence: (B, S, num_patches_y, num_patches_x)
|
||||
Returns:
|
||||
R: (B, S, 3, 3)
|
||||
T: (B, S, 3)
|
||||
focal_lengths: (B, S, 2)
|
||||
principal_points: (B, S, 2)
|
||||
"""
|
||||
if confidence is None:
|
||||
confidence = torch.ones_like(camray[:, :, :, :, 0])
|
||||
B, S, num_patches_y, num_patches_x, _ = camray.shape
|
||||
# identity K, assume imw=imh=2.0
|
||||
I_K = torch.eye(3, dtype=camray.dtype, device=camray.device)
|
||||
I_K[0, 2] = 1.0
|
||||
I_K[1, 2] = 1.0
|
||||
# repeat I_K to match camray
|
||||
I_K = I_K.unsqueeze(0).unsqueeze(0).expand(B, S, -1, -1)
|
||||
|
||||
cam_plane_depth = torch.ones(
|
||||
B, S, num_patches_y, num_patches_x, 1, dtype=camray.dtype, device=camray.device
|
||||
)
|
||||
I_cam_plane_unproj = unproject_depth(
|
||||
cam_plane_depth,
|
||||
I_K,
|
||||
c2w=None,
|
||||
ixt_normalized=True,
|
||||
num_patches_x=num_patches_x,
|
||||
num_patches_y=num_patches_y,
|
||||
) # (B, S, num_patches_y, num_patches_x, 3)
|
||||
|
||||
camray = camray.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x, 6)
|
||||
I_cam_plane_unproj = I_cam_plane_unproj.flatten(0, 1).flatten(
|
||||
1, 2
|
||||
) # (B*S, num_patches_y*num_patches_x, 3)
|
||||
confidence = confidence.flatten(0, 1).flatten(1, 2) # (B*S, num_patches_y*num_patches_x)
|
||||
|
||||
# Compute optimal rotation to align rays
|
||||
N = camray.shape[-2]
|
||||
device = camray.device
|
||||
n_iter, num_sample_for_ransac, n_sample, rand_sample_iters_idx = get_params_for_ransac(N, device)
|
||||
|
||||
# Use batch processing (confidence is guaranteed to be not None at this point)
|
||||
if training:
|
||||
camray = camray.clone().detach()
|
||||
I_cam_plane_unproj = I_cam_plane_unproj.clone().detach()
|
||||
confidence = confidence.clone().detach()
|
||||
R, focal_lengths, principal_points = compute_optimal_rotation_intrinsics_batch(
|
||||
I_cam_plane_unproj,
|
||||
camray[:, :, :3],
|
||||
reproj_threshold=reproj_threshold,
|
||||
weights=confidence,
|
||||
n_sample = n_sample,
|
||||
n_iter=n_iter,
|
||||
num_sample_for_ransac=num_sample_for_ransac,
|
||||
rand_sample_iters_idx=rand_sample_iters_idx,
|
||||
)
|
||||
|
||||
T = torch.sum(camray[:, :, 3:] * confidence.unsqueeze(-1), dim=1) / torch.sum(
|
||||
confidence, dim=-1, keepdim=True
|
||||
)
|
||||
|
||||
R = R.reshape(B, S, 3, 3)
|
||||
T = T.reshape(B, S, 3)
|
||||
focal_lengths = focal_lengths.reshape(B, S, 2)
|
||||
principal_points = principal_points.reshape(B, S, 2)
|
||||
|
||||
return R, T, 1.0 / focal_lengths, principal_points + 1.0
|
||||
|
||||
def get_extrinsic_from_camray(camray, conf, patch_size_y, patch_size_x, training=False):
|
||||
pred_R, pred_T, pred_focal_lengths, pred_principal_points = camray_to_caminfo(
|
||||
camray, confidence=conf.squeeze(-1), training=training
|
||||
)
|
||||
|
||||
pred_extrinsic = torch.cat(
|
||||
[
|
||||
torch.cat([pred_R, pred_T.unsqueeze(-1)], dim=-1),
|
||||
repeat(
|
||||
torch.tensor([0, 0, 0, 1], dtype=pred_R.dtype, device=pred_R.device),
|
||||
"c -> b s 1 c",
|
||||
b=pred_R.shape[0],
|
||||
s=pred_R.shape[1],
|
||||
),
|
||||
],
|
||||
dim=-2,
|
||||
) # B, S, 4, 4
|
||||
return pred_extrinsic, pred_focal_lengths, pred_principal_points
|
||||
@@ -1,22 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "moge"
|
||||
version = "2.0.0"
|
||||
description = "MoGe monocular geometry estimation"
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
dependencies = [
|
||||
"huggingface_hub",
|
||||
"numpy",
|
||||
"pillow",
|
||||
"torch>=2.0.0",
|
||||
"utils3d",
|
||||
]
|
||||
requires-python = ">=3.9"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["moge*"]
|
||||
@@ -1,16 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "utils3d"
|
||||
version = "1.7"
|
||||
description = "Utilities for 3D computer vision"
|
||||
readme = "README.md"
|
||||
license = {text = "MIT"}
|
||||
dependencies = ["numpy", "scipy"]
|
||||
requires-python = ">=3.8"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["utils3d*"]
|
||||
@@ -65,7 +65,7 @@ def download_artifact(artifact, destination_dir=VENDOR_DIR):
|
||||
|
||||
|
||||
def prepare(platforms=None, lock_path=LOCK_PATH, destination_dir=VENDOR_DIR):
|
||||
"""Prepare common third-party archives for selected platforms."""
|
||||
"""Prepare bundled uv executables for selected platforms."""
|
||||
lock = load_lock(lock_path)
|
||||
available = lock["platforms"]
|
||||
selected = list(available) if platforms is None else platforms
|
||||
@@ -73,8 +73,7 @@ def prepare(platforms=None, lock_path=LOCK_PATH, destination_dir=VENDOR_DIR):
|
||||
if unknown:
|
||||
raise ValueError(f"Unsupported platform: {', '.join(unknown)}")
|
||||
|
||||
artifacts = list(lock["artifacts"].values())
|
||||
artifacts.extend(available[platform] for platform in selected)
|
||||
artifacts = [available[platform] for platform in selected]
|
||||
return [
|
||||
download_artifact(artifact, destination_dir)
|
||||
for artifact in artifacts
|
||||
|
||||
@@ -21,9 +21,9 @@ from .operators import (
|
||||
)
|
||||
from .panel import (
|
||||
Ben2RemovalDebugPanel,
|
||||
AiServicePanel,
|
||||
Da3OutputDebugPanel,
|
||||
DebugPanel,
|
||||
AiServicePanel,
|
||||
ModelsPanel,
|
||||
Moge2ApiDebugPanel,
|
||||
)
|
||||
@@ -50,8 +50,8 @@ CLASSES = (
|
||||
AiServicePanel,
|
||||
ModelsPanel,
|
||||
DebugPanel,
|
||||
Da3OutputDebugPanel,
|
||||
Moge2ApiDebugPanel,
|
||||
Da3OutputDebugPanel,
|
||||
Ben2RemovalDebugPanel,
|
||||
)
|
||||
|
||||
|
||||
50
src/snapform/ben2_onnx.py
Normal file
50
src/snapform/ben2_onnx.py
Normal file
@@ -0,0 +1,50 @@
|
||||
"""BEN2 ONNX inference with Pillow and NumPy preprocessing."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from .onnx_runtime import create_session as create_runtime_session
|
||||
except ImportError:
|
||||
from onnx_runtime import create_session as create_runtime_session
|
||||
|
||||
|
||||
ONNX_PATH = Path("onnx/model_fp16.onnx")
|
||||
INPUT_SIZE = (1024, 1024)
|
||||
|
||||
|
||||
def model_path(model_directory):
|
||||
return Path(model_directory) / ONNX_PATH
|
||||
|
||||
|
||||
def create_session(model_directory, requested_device="auto"):
|
||||
return create_runtime_session(model_path(model_directory), requested_device)
|
||||
|
||||
|
||||
def preprocess(image):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
resized = image.convert("RGB").resize(INPUT_SIZE, Image.Resampling.LANCZOS)
|
||||
array = np.asarray(resized, dtype=np.float32) / 255.0
|
||||
mean = np.array((0.485, 0.456, 0.406), dtype=np.float32)
|
||||
std = np.array((0.229, 0.224, 0.225), dtype=np.float32)
|
||||
array = (array - mean) / std
|
||||
return np.ascontiguousarray(array.transpose(2, 0, 1)[None])
|
||||
|
||||
|
||||
def infer(session, image):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
input_name = session.get_inputs()[0].name
|
||||
output_name = session.get_outputs()[0].name
|
||||
alpha = session.run([output_name], {input_name: preprocess(image)})[0]
|
||||
alpha = np.asarray(alpha, dtype=np.float32).squeeze()
|
||||
minimum = float(alpha.min())
|
||||
span = max(float(alpha.max()) - minimum, np.finfo(np.float32).eps)
|
||||
alpha = np.clip((alpha - minimum) / span * 255.0, 0.0, 255.0)
|
||||
mask = Image.fromarray(alpha.astype(np.uint8), mode="L")
|
||||
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
|
||||
foreground = image.convert("RGBA")
|
||||
foreground.putalpha(mask)
|
||||
return foreground
|
||||
@@ -26,10 +26,8 @@ from .models import (
|
||||
)
|
||||
|
||||
|
||||
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
|
||||
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
|
||||
MANIFEST_FILENAME = "manifest.json"
|
||||
MANIFEST_REVISION = 5
|
||||
MANIFEST_REVISION = 7
|
||||
SERVICE_ADDRESS = "127.0.0.1"
|
||||
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
|
||||
JOB_FOLDER_NAMES = {
|
||||
@@ -172,8 +170,8 @@ def env_devices():
|
||||
def machine_devices():
|
||||
"""Devices the machine can run before environment installation."""
|
||||
devices = ["cpu"]
|
||||
if shutil.which("nvidia-smi"):
|
||||
devices.insert(0, "cuda")
|
||||
if sys.platform == "win32":
|
||||
devices.insert(0, "directml")
|
||||
elif sys.platform == "darwin" and platform.machine().lower() in {
|
||||
"arm64",
|
||||
"aarch64",
|
||||
@@ -189,6 +187,7 @@ def supported_devices():
|
||||
|
||||
DEVICE_LABELS = {
|
||||
"cuda": ("CUDA", "NVIDIA CUDA"),
|
||||
"directml": ("DirectML", "Windows GPU (DirectML)"),
|
||||
"mps": ("MPS", "Apple Metal (MPS)"),
|
||||
"cpu": ("CPU", "CPU"),
|
||||
}
|
||||
@@ -283,7 +282,6 @@ def uv_cache_path():
|
||||
|
||||
def install_env_command(status_path):
|
||||
executable = uv_executable()
|
||||
torch_variant = "cu128" if "cuda" in machine_devices() else "cpu"
|
||||
script = Path(__file__).with_name("install_env.py")
|
||||
command = [
|
||||
sys.executable,
|
||||
@@ -296,12 +294,6 @@ def install_env_command(status_path):
|
||||
str(storage_root() / MANIFEST_FILENAME),
|
||||
"--models-dir",
|
||||
str(models_directory()),
|
||||
"--torch-variant",
|
||||
torch_variant,
|
||||
"--da3-commit",
|
||||
DA3_COMMIT,
|
||||
"--ben2-commit",
|
||||
BEN2_COMMIT,
|
||||
"--status",
|
||||
str(status_path),
|
||||
]
|
||||
|
||||
221
src/snapform/da3_onnx.py
Normal file
221
src/snapform/da3_onnx.py
Normal file
@@ -0,0 +1,221 @@
|
||||
"""Depth Anything 3 inference with ONNX Runtime and NumPy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from .onnx_runtime import create_session as create_onnx_session
|
||||
except ImportError:
|
||||
from onnx_runtime import create_session as create_onnx_session
|
||||
|
||||
|
||||
ONNX_PATH = "model.onnx"
|
||||
PATCH_SIZE = 14
|
||||
STRATEGY_CODES = {
|
||||
"first": 0,
|
||||
"middle": 1,
|
||||
"saddle_balanced": 2,
|
||||
"saddle_sim_range": 3,
|
||||
}
|
||||
IMAGE_MEAN = np.asarray((0.485, 0.456, 0.406), dtype=np.float32)
|
||||
IMAGE_STD = np.asarray((0.229, 0.224, 0.225), dtype=np.float32)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Prediction:
|
||||
depth: np.ndarray
|
||||
conf: np.ndarray
|
||||
extrinsics: np.ndarray
|
||||
intrinsics: np.ndarray
|
||||
processed_images: np.ndarray
|
||||
is_metric: int = 0
|
||||
sky: np.ndarray | None = None
|
||||
gaussians: object | None = None
|
||||
aux: dict | None = None
|
||||
scale_factor: float | None = None
|
||||
|
||||
|
||||
def create_session(model_dir, requested_device="auto"):
|
||||
path = Path(model_dir) / ONNX_PATH
|
||||
return create_onnx_session(path, requested_device)
|
||||
|
||||
|
||||
def nearest_multiple(value, multiple=PATCH_SIZE):
|
||||
lower = value // multiple * multiple
|
||||
upper = lower + multiple
|
||||
return upper if abs(upper - value) <= abs(value - lower) else lower
|
||||
|
||||
|
||||
def target_size(size, process_res, method):
|
||||
width, height = size
|
||||
boundary = max(width, height)
|
||||
if method.startswith("lower_bound"):
|
||||
boundary = min(width, height)
|
||||
scale = float(process_res) / boundary
|
||||
width = max(1, int(round(width * scale)))
|
||||
height = max(1, int(round(height * scale)))
|
||||
if method.endswith("resize"):
|
||||
return (
|
||||
max(PATCH_SIZE, nearest_multiple(width)),
|
||||
max(PATCH_SIZE, nearest_multiple(height)),
|
||||
)
|
||||
if method.endswith("crop"):
|
||||
return (
|
||||
max(PATCH_SIZE, width // PATCH_SIZE * PATCH_SIZE),
|
||||
max(PATCH_SIZE, height // PATCH_SIZE * PATCH_SIZE),
|
||||
)
|
||||
raise ValueError(f"Unsupported process_res_method: {method}")
|
||||
|
||||
|
||||
def resize_and_crop(image, process_res, method):
|
||||
from PIL import Image
|
||||
width, height = image.size
|
||||
boundary = max(width, height)
|
||||
if method.startswith("lower_bound"):
|
||||
boundary = min(width, height)
|
||||
scale = float(process_res) / boundary
|
||||
resized_size = (
|
||||
max(1, int(round(width * scale))),
|
||||
max(1, int(round(height * scale))),
|
||||
)
|
||||
resampling = (
|
||||
Image.Resampling.BICUBIC
|
||||
if scale > 1.0
|
||||
else Image.Resampling.BOX
|
||||
)
|
||||
image = image.resize(resized_size, resample=resampling)
|
||||
final_size = target_size((width, height), process_res, method)
|
||||
if method.endswith("resize"):
|
||||
if image.size != final_size:
|
||||
upscale = final_size[0] > image.width or final_size[1] > image.height
|
||||
resampling = (
|
||||
Image.Resampling.BICUBIC
|
||||
if upscale
|
||||
else Image.Resampling.BOX
|
||||
)
|
||||
image = image.resize(final_size, resample=resampling)
|
||||
return image
|
||||
left = (image.width - final_size[0]) // 2
|
||||
top = (image.height - final_size[1]) // 2
|
||||
return image.crop((left, top, left + final_size[0], top + final_size[1]))
|
||||
|
||||
|
||||
def preprocess(image_paths, process_res=504, process_res_method="upper_bound_resize"):
|
||||
from PIL import Image
|
||||
processed = []
|
||||
for path in image_paths:
|
||||
with Image.open(path) as source:
|
||||
image = resize_and_crop(
|
||||
source.convert("RGB"),
|
||||
process_res,
|
||||
process_res_method,
|
||||
)
|
||||
processed.append(np.asarray(image, dtype=np.uint8))
|
||||
min_height = min(image.shape[0] for image in processed)
|
||||
min_width = min(image.shape[1] for image in processed)
|
||||
cropped = []
|
||||
for image in processed:
|
||||
top = (image.shape[0] - min_height) // 2
|
||||
left = (image.shape[1] - min_width) // 2
|
||||
cropped.append(
|
||||
image[top : top + min_height, left : left + min_width]
|
||||
)
|
||||
images = np.stack(cropped)
|
||||
normalized = images.astype(np.float32) / 255.0
|
||||
normalized = (normalized - IMAGE_MEAN) / IMAGE_STD
|
||||
model_input = normalized.transpose(0, 3, 1, 2)[None]
|
||||
return np.ascontiguousarray(model_input), images
|
||||
|
||||
|
||||
def quaternion_to_matrix(quaternion):
|
||||
i, j, k, real = np.moveaxis(quaternion, -1, 0)
|
||||
two_scale = 2.0 / np.sum(quaternion * quaternion, axis=-1)
|
||||
values = (
|
||||
1 - two_scale * (j * j + k * k),
|
||||
two_scale * (i * j - k * real),
|
||||
two_scale * (i * k + j * real),
|
||||
two_scale * (i * j + k * real),
|
||||
1 - two_scale * (i * i + k * k),
|
||||
two_scale * (j * k - i * real),
|
||||
two_scale * (i * k - j * real),
|
||||
two_scale * (j * k + i * real),
|
||||
1 - two_scale * (i * i + j * j),
|
||||
)
|
||||
return np.stack(values, axis=-1).reshape(quaternion.shape[:-1] + (3, 3))
|
||||
|
||||
|
||||
def camera_from_pose(pose_encoding, height, width):
|
||||
translation = pose_encoding[..., :3]
|
||||
rotation = quaternion_to_matrix(pose_encoding[..., 3:7])
|
||||
rotation_inverse = np.swapaxes(rotation, -1, -2)
|
||||
translation_inverse = -rotation_inverse @ translation[..., None]
|
||||
extrinsics = np.concatenate(
|
||||
(rotation_inverse, translation_inverse),
|
||||
axis=-1,
|
||||
).astype(np.float32)
|
||||
|
||||
field_of_view_height = pose_encoding[..., 7]
|
||||
field_of_view_width = pose_encoding[..., 8]
|
||||
focal_y = (height / 2.0) / np.maximum(
|
||||
np.tan(field_of_view_height / 2.0),
|
||||
1e-6,
|
||||
)
|
||||
focal_x = (width / 2.0) / np.maximum(
|
||||
np.tan(field_of_view_width / 2.0),
|
||||
1e-6,
|
||||
)
|
||||
intrinsics = np.zeros(pose_encoding.shape[:-1] + (3, 3), dtype=np.float32)
|
||||
intrinsics[..., 0, 0] = focal_x
|
||||
intrinsics[..., 1, 1] = focal_y
|
||||
intrinsics[..., 0, 2] = width / 2.0
|
||||
intrinsics[..., 1, 2] = height / 2.0
|
||||
intrinsics[..., 2, 2] = 1.0
|
||||
return extrinsics, intrinsics
|
||||
|
||||
|
||||
def infer(
|
||||
session,
|
||||
image_paths,
|
||||
process_res=504,
|
||||
process_res_method="upper_bound_resize",
|
||||
ref_view_strategy="saddle_balanced",
|
||||
use_ray_pose=False,
|
||||
):
|
||||
del use_ray_pose
|
||||
inputs = session.get_inputs()
|
||||
model_input, processed_images = preprocess(
|
||||
image_paths,
|
||||
process_res,
|
||||
process_res_method,
|
||||
)
|
||||
strategy_code = np.asarray(
|
||||
[STRATEGY_CODES.get(ref_view_strategy, 2)],
|
||||
dtype=np.int64,
|
||||
)
|
||||
input_names = [value.name for value in inputs]
|
||||
outputs = session.run(
|
||||
None,
|
||||
{
|
||||
input_names[0]: model_input,
|
||||
input_names[1]: strategy_code,
|
||||
},
|
||||
)
|
||||
depth, confidence, pose_encoding = outputs[:3]
|
||||
height, width = depth.shape[-2:]
|
||||
extrinsics, intrinsics = camera_from_pose(
|
||||
pose_encoding[0],
|
||||
height,
|
||||
width,
|
||||
)
|
||||
return Prediction(
|
||||
depth=np.asarray(depth[0], dtype=np.float32),
|
||||
conf=np.asarray(confidence[0], dtype=np.float32),
|
||||
extrinsics=extrinsics,
|
||||
intrinsics=intrinsics,
|
||||
processed_images=processed_images,
|
||||
aux={},
|
||||
)
|
||||
@@ -2239,12 +2239,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
name="Depth From",
|
||||
description="Model or original image used to build depth",
|
||||
items=image_depth_model_items,
|
||||
default=1,
|
||||
default=0,
|
||||
)
|
||||
snapform_use_ray_pose: bpy.props.BoolProperty(
|
||||
name="Use Ray Pose",
|
||||
description="Use ray-based camera pose estimation",
|
||||
default=True,
|
||||
description="Legacy task compatibility; ONNX uses the camera decoder",
|
||||
default=False,
|
||||
)
|
||||
snapform_process_res_factor: process_resolution_scale_property()
|
||||
snapform_moge_resolution_level: bpy.props.IntProperty(
|
||||
@@ -2278,14 +2278,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
layout = self.layout
|
||||
layout.prop(self, "snapform_model")
|
||||
if self.snapform_model in {"SMALL", "BASE"}:
|
||||
layout.prop(self, "snapform_use_ray_pose")
|
||||
draw_process_resolution(layout, self, context.object.data)
|
||||
elif self.snapform_model.startswith("MOGE2_"):
|
||||
from .client import model_ready
|
||||
|
||||
if not model_ready(self.snapform_model):
|
||||
layout.label(text="MoGe-2 model is not installed", icon="ERROR")
|
||||
layout.label(text="Download it in Models")
|
||||
layout.prop(self, "snapform_moge_resolution_level")
|
||||
layout.prop(self, "snapform_remove_background")
|
||||
layout.prop(self, "snapform_unshaded")
|
||||
|
||||
@@ -608,12 +608,12 @@ class ConvertImageToMesh(bpy.types.Operator):
|
||||
name="Depth From",
|
||||
description="Model used to estimate image depth",
|
||||
items=image_surface_model_items,
|
||||
default=1,
|
||||
default=0,
|
||||
)
|
||||
snapform_use_ray_pose: bpy.props.BoolProperty(
|
||||
name="Use Ray Pose",
|
||||
description="Use ray-based camera pose estimation",
|
||||
default=True,
|
||||
description="Legacy task compatibility; ONNX uses the camera decoder",
|
||||
default=False,
|
||||
)
|
||||
snapform_process_res_factor: process_resolution_scale_property()
|
||||
snapform_moge_resolution_level: bpy.props.IntProperty(
|
||||
@@ -670,7 +670,6 @@ class ConvertImageToMesh(bpy.types.Operator):
|
||||
if self.snapform_model.startswith("MOGE2_"):
|
||||
layout.prop(self, "snapform_moge_resolution_level")
|
||||
else:
|
||||
layout.prop(self, "snapform_use_ray_pose")
|
||||
draw_process_resolution(layout, self, context.object.data)
|
||||
layout.prop(self, "snapform_create_camera")
|
||||
layout.prop(self, "snapform_remove_background")
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import argparse
|
||||
import filecmp
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
@@ -28,47 +27,36 @@ PYPI_MIRRORS = (
|
||||
"https://mirrors.aliyun.com/pypi/simple/",
|
||||
)
|
||||
PYPI_DEFAULT_INDEX = "https://pypi.org/simple/"
|
||||
TORCH_VERSION = "2.9.1"
|
||||
TORCHVISION_VERSION = "0.24.1"
|
||||
XFORMERS_VERSION = "0.0.33.post2"
|
||||
ADDICT_VERSION = "2.4.0"
|
||||
FASTAPI_VERSION = "0.139.2"
|
||||
UVICORN_VERSION = "0.51.0"
|
||||
MANIFEST_REVISION = 5
|
||||
TORCH_INDEXES = {
|
||||
"cu128": "https://download.pytorch.org/whl/cu128",
|
||||
"cpu": "https://download.pytorch.org/whl/cpu",
|
||||
}
|
||||
ONNXRUNTIME_VERSION = "1.24.4"
|
||||
NUMPY_VERSION = "2.4.2"
|
||||
PILLOW_VERSION = "12.1.1"
|
||||
OPENCV_VERSION = "4.13.0.92"
|
||||
TRIMESH_VERSION = "4.11.2"
|
||||
HUGGINGFACE_HUB_VERSION = "1.4.1"
|
||||
MANIFEST_REVISION = 7
|
||||
_resolved_pypi_index = None
|
||||
DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key)
|
||||
RUNTIME_STAGE_COUNT = 3
|
||||
RUNTIME_STAGE_COUNT = 2
|
||||
ENVIRONMENT_SLICE = (0.02, 0.06)
|
||||
PYTORCH_SLICE = (0.06, 0.48)
|
||||
XFORMERS_SLICE = (0.48, 0.62)
|
||||
AI_COMPONENTS_SLICE = (0.62, 0.85)
|
||||
AI_COMPONENTS_SLICE = (0.10, 0.85)
|
||||
MODELS_SLICE_START = 0.85
|
||||
MODELS_SLICE_END = 1.0
|
||||
RAMP_RATE = 0.02
|
||||
CREEP_CAP = 0.60
|
||||
INSTALL_ESTIMATES = {
|
||||
"PyTorch": 8 * 60,
|
||||
"AI components": 8 * 60,
|
||||
"ONNX Runtime": 3 * 60,
|
||||
}
|
||||
INSTALL_ESTIMATE_RANGES = {
|
||||
"PyTorch": PYTORCH_SLICE,
|
||||
"AI components": (XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1]),
|
||||
"ONNX Runtime": AI_COMPONENTS_SLICE,
|
||||
}
|
||||
CHECK_CODE = (
|
||||
"import json; import torch; "
|
||||
"from depth_anything_3.api import DepthAnything3; "
|
||||
"from ben2 import BEN_Base; "
|
||||
"from moge.model.v2 import MoGeModel; "
|
||||
"import fastapi, uvicorn; "
|
||||
"mps = getattr(torch.backends, 'mps', None); "
|
||||
"import json; import fastapi, uvicorn, onnxruntime; "
|
||||
"import cv2, numpy, PIL, trimesh; "
|
||||
"print('SNAPFORM_DEVICES=' + json.dumps({"
|
||||
"'torch': torch.__version__, "
|
||||
"'cuda': torch.cuda.is_available(), "
|
||||
"'mps': bool(mps is not None and mps.is_available())}))"
|
||||
"'onnxruntime': onnxruntime.__version__, "
|
||||
"'providers': onnxruntime.get_available_providers()}))"
|
||||
)
|
||||
|
||||
|
||||
@@ -158,10 +146,6 @@ def estimated_remaining(label, progress, elapsed=0.0):
|
||||
return max(expected_seconds * remaining_fraction - elapsed, 0.0)
|
||||
|
||||
|
||||
def effective_torch_index(torch_variant):
|
||||
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(torch_variant)
|
||||
|
||||
|
||||
def resolve_indexes(args):
|
||||
"""???????????????????????????????????"""
|
||||
global _resolved_pypi_index
|
||||
@@ -320,52 +304,13 @@ def venv_python(venv_dir):
|
||||
return Path(venv_dir) / "bin" / "python"
|
||||
|
||||
|
||||
def install_stages(args, venv_dir, python, uv, vendor_dir=None):
|
||||
# macOS 的 PyTorch wheel 只在 PyPI 上，download.pytorch.org 没有 macosx 构建
|
||||
torch_index = (
|
||||
None
|
||||
if sys.platform == "darwin"
|
||||
else effective_torch_index(args.torch_variant)
|
||||
def install_stages(args, venv_dir, python, uv):
|
||||
"""Build the lightweight ONNX-only runtime and model download stages."""
|
||||
runtime_package = (
|
||||
f"onnxruntime-directml=={ONNXRUNTIME_VERSION}"
|
||||
if sys.platform == "win32"
|
||||
else f"onnxruntime=={ONNXRUNTIME_VERSION}"
|
||||
)
|
||||
vendor_dir = Path(vendor_dir or Path(venv_dir).parent / "vendor")
|
||||
|
||||
def archive_requirement(project_name, filename):
|
||||
local = Path(vendor_dir) / filename
|
||||
if not local.is_file():
|
||||
raise RuntimeError(
|
||||
"Bundled vendor artifact is missing; reinstall the extension: "
|
||||
f"{local}"
|
||||
)
|
||||
return f"{project_name} @ {local.as_uri()}"
|
||||
|
||||
da3_requirement = archive_requirement(
|
||||
"depth-anything-3",
|
||||
"depth-anything-3.tar.gz",
|
||||
)
|
||||
ben2_requirement = archive_requirement(
|
||||
"ben2",
|
||||
"ben2.tar.gz",
|
||||
)
|
||||
utils3d_requirement = archive_requirement(
|
||||
"utils3d",
|
||||
"utils3d.tar.gz",
|
||||
)
|
||||
moge_requirement = archive_requirement(
|
||||
"moge",
|
||||
"moge.tar.gz",
|
||||
)
|
||||
torch_command = [
|
||||
uv,
|
||||
"pip",
|
||||
"install",
|
||||
"-v",
|
||||
"--python",
|
||||
python,
|
||||
f"torch=={TORCH_VERSION}",
|
||||
f"torchvision=={TORCHVISION_VERSION}",
|
||||
]
|
||||
if torch_index is not None:
|
||||
torch_command.extend(["--index-url", torch_index])
|
||||
stages = [
|
||||
(
|
||||
"Environment",
|
||||
@@ -374,13 +319,7 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None):
|
||||
"plain",
|
||||
),
|
||||
(
|
||||
"PyTorch",
|
||||
torch_command,
|
||||
None,
|
||||
"install",
|
||||
),
|
||||
(
|
||||
"AI components",
|
||||
"ONNX Runtime",
|
||||
[
|
||||
uv,
|
||||
"pip",
|
||||
@@ -388,38 +327,19 @@ def install_stages(args, venv_dir, python, uv, vendor_dir=None):
|
||||
"-v",
|
||||
"--python",
|
||||
python,
|
||||
da3_requirement,
|
||||
f"addict=={ADDICT_VERSION}",
|
||||
ben2_requirement,
|
||||
utils3d_requirement,
|
||||
moge_requirement,
|
||||
f"fastapi=={FASTAPI_VERSION}",
|
||||
f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}",
|
||||
f"numpy=={NUMPY_VERSION}",
|
||||
f"opencv-python-headless=={OPENCV_VERSION}",
|
||||
runtime_package,
|
||||
f"pillow=={PILLOW_VERSION}",
|
||||
f"trimesh=={TRIMESH_VERSION}",
|
||||
f"uvicorn=={UVICORN_VERSION}",
|
||||
"setuptools",
|
||||
"wheel",
|
||||
],
|
||||
None,
|
||||
"install",
|
||||
),
|
||||
]
|
||||
if sys.platform != "darwin":
|
||||
xformers_command = [
|
||||
uv,
|
||||
"pip",
|
||||
"install",
|
||||
"-v",
|
||||
"--python",
|
||||
python,
|
||||
f"xformers=={XFORMERS_VERSION}",
|
||||
]
|
||||
if args.torch_variant == "cu128":
|
||||
xformers_command.extend(["--index-url", torch_index])
|
||||
else:
|
||||
xformers_command.append("--no-deps")
|
||||
stages.insert(
|
||||
2,
|
||||
("AI components", xformers_command, None, "install"),
|
||||
)
|
||||
models = DOWNLOADABLE_MODELS
|
||||
worker_script = Path(__file__).with_name("worker.py")
|
||||
models_dir = Path(args.models_dir)
|
||||
@@ -475,13 +395,13 @@ def env_installed(venv_dir, manifest_path=None):
|
||||
)
|
||||
|
||||
|
||||
def installed_devices(args):
|
||||
def installed_devices():
|
||||
devices = []
|
||||
machine = platform.machine().lower()
|
||||
if sys.platform == "darwin" and machine in {"arm64", "aarch64"}:
|
||||
devices.append("mps")
|
||||
elif args.torch_variant.startswith("cu"):
|
||||
devices.append("cuda")
|
||||
elif sys.platform == "win32":
|
||||
devices.append("directml")
|
||||
devices.append("cpu")
|
||||
return devices
|
||||
|
||||
@@ -592,82 +512,6 @@ def bundled_uv_archive():
|
||||
return path if path.is_file() else None
|
||||
|
||||
|
||||
def vendor_overlay_directory(filename):
|
||||
names = {
|
||||
"depth-anything-3.tar.gz": "depth-anything-3",
|
||||
"moge.tar.gz": "moge",
|
||||
"utils3d.tar.gz": "utils3d",
|
||||
}
|
||||
overlay_name = names.get(filename)
|
||||
if overlay_name is None:
|
||||
return None
|
||||
project_root = Path(__file__).resolve().parents[2]
|
||||
overlay = project_root / "patches" / overlay_name
|
||||
return overlay if overlay.is_dir() else None
|
||||
|
||||
|
||||
def copy_archive_with_overlay(source, target, overlay):
|
||||
"""开发目录使用原始 vendor 时，应用与正式打包相同的源码覆盖。"""
|
||||
with tarfile.open(source, "r:gz") as input_archive:
|
||||
members = input_archive.getmembers()
|
||||
top_level = members[0].name.split("/", 1)[0]
|
||||
with tarfile.open(target, "w:gz") as output_archive:
|
||||
for member in members:
|
||||
if member.isdir() or member.islnk() or member.issym():
|
||||
output_archive.addfile(member)
|
||||
continue
|
||||
relative = member.name[len(top_level) + 1 :]
|
||||
replacement = overlay / relative
|
||||
if replacement.is_file():
|
||||
info = tarfile.TarInfo(member.name)
|
||||
info.size = replacement.stat().st_size
|
||||
info.mode = member.mode
|
||||
with replacement.open("rb") as stream:
|
||||
output_archive.addfile(info, stream)
|
||||
continue
|
||||
with input_archive.extractfile(member) as stream:
|
||||
output_archive.addfile(member, stream)
|
||||
|
||||
|
||||
def copy_vendor_inputs(bundled_dir, vendor_dir):
|
||||
"""Copy bundled installation inputs into a temporary directory."""
|
||||
for name in (
|
||||
"depth-anything-3.tar.gz",
|
||||
"ben2.tar.gz",
|
||||
"utils3d.tar.gz",
|
||||
"moge.tar.gz",
|
||||
):
|
||||
bundled = Path(bundled_dir) / name
|
||||
if not bundled.is_file():
|
||||
continue
|
||||
target = Path(vendor_dir) / name
|
||||
overlay = vendor_overlay_directory(name)
|
||||
if overlay is None and target.is_file() and filecmp.cmp(
|
||||
bundled,
|
||||
target,
|
||||
shallow=False,
|
||||
):
|
||||
continue
|
||||
Path(vendor_dir).mkdir(parents=True, exist_ok=True)
|
||||
if overlay is not None and tarfile.is_tarfile(bundled):
|
||||
copy_archive_with_overlay(bundled, target, overlay)
|
||||
else:
|
||||
shutil.copy2(bundled, target)
|
||||
|
||||
|
||||
def remove_legacy_vendor(storage_root):
|
||||
"""Remove installation inputs persisted by older development builds."""
|
||||
root = Path(storage_root)
|
||||
shutil.rmtree(root / "vendor", ignore_errors=True)
|
||||
for name in (
|
||||
"depth-anything-3.tar.gz",
|
||||
"ben2.tar.gz",
|
||||
"utils3d.tar.gz",
|
||||
"moge.tar.gz",
|
||||
):
|
||||
(root / name).unlink(missing_ok=True)
|
||||
|
||||
|
||||
def archive_is_valid(path):
|
||||
if zipfile.is_zipfile(path):
|
||||
return True
|
||||
@@ -800,14 +644,6 @@ def run_stage(
|
||||
|
||||
|
||||
def install(args):
|
||||
if sys.platform == "darwin" and platform.machine().lower() not in {
|
||||
"arm64",
|
||||
"aarch64",
|
||||
}:
|
||||
raise RuntimeError(
|
||||
"SnapForm requires Apple Silicon macOS; PyTorch stopped "
|
||||
"shipping Intel Mac (x86_64) wheels after 2.2.2"
|
||||
)
|
||||
resolve_indexes(args)
|
||||
venv_dir = Path(args.venv_dir)
|
||||
manifest_path = Path(
|
||||
@@ -823,13 +659,7 @@ def install(args):
|
||||
uv = Path(resolved) if resolved else uv
|
||||
uv_missing = not uv.is_file()
|
||||
python = venv_python(venv_dir)
|
||||
vendor_temp = tempfile.TemporaryDirectory(prefix="snapform_vendor_")
|
||||
vendor_dir = Path(vendor_temp.name)
|
||||
copy_vendor_inputs(
|
||||
bundled_vendor_directory(),
|
||||
vendor_dir,
|
||||
)
|
||||
stages = install_stages(args, venv_dir, python, uv, vendor_dir)
|
||||
stages = install_stages(args, venv_dir, python, uv)
|
||||
env_stages = [stage for stage in stages if stage[3] != "download"]
|
||||
model_stages = [stage for stage in stages if stage[3] == "download"]
|
||||
model_span = (
|
||||
@@ -838,23 +668,15 @@ def install(args):
|
||||
else 0.0
|
||||
)
|
||||
installed_packages = []
|
||||
component_stages = env_stages[2:]
|
||||
component_ranges = (
|
||||
[XFORMERS_SLICE, AI_COMPONENTS_SLICE]
|
||||
if len(component_stages) == 2
|
||||
else [(XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1])]
|
||||
)
|
||||
for index, (label, command, environment, kind) in enumerate(env_stages):
|
||||
if index == 0:
|
||||
stage = 1
|
||||
stage_start, stage_end = ENVIRONMENT_SLICE
|
||||
elif index == 1:
|
||||
stage = 2
|
||||
stage_start, stage_end = PYTORCH_SLICE
|
||||
stage_start, stage_end = AI_COMPONENTS_SLICE
|
||||
else:
|
||||
stage = 3
|
||||
component_index = index - 2
|
||||
stage_start, stage_end = component_ranges[component_index]
|
||||
raise RuntimeError(f"Unexpected runtime stage: {label}")
|
||||
if uv_missing and index == 0:
|
||||
midpoint = stage_start + (stage_end - stage_start) * 0.5
|
||||
install_uv(
|
||||
@@ -880,9 +702,7 @@ def install(args):
|
||||
)
|
||||
if isinstance(stage_installed, list):
|
||||
installed_packages.extend(stage_installed)
|
||||
vendor_temp.cleanup()
|
||||
remove_legacy_vendor(venv_dir.parent)
|
||||
devices = installed_devices(args)
|
||||
devices = installed_devices()
|
||||
venv_dir.mkdir(parents=True, exist_ok=True)
|
||||
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
ready_path = manifest_path
|
||||
@@ -891,14 +711,13 @@ def install(args):
|
||||
{
|
||||
"manifest_revision": MANIFEST_REVISION,
|
||||
"python": "3.12",
|
||||
"torch": TORCH_VERSION,
|
||||
"xformers": XFORMERS_VERSION,
|
||||
"onnxruntime": ONNXRUNTIME_VERSION,
|
||||
"fastapi": FASTAPI_VERSION,
|
||||
"uvicorn": UVICORN_VERSION,
|
||||
"torch_variant": args.torch_variant,
|
||||
"runtime_variant": (
|
||||
"directml" if sys.platform == "win32" else "default"
|
||||
),
|
||||
"devices": devices,
|
||||
"da3_commit": args.da3_commit,
|
||||
"ben2_commit": args.ben2_commit,
|
||||
"packages": sorted(set(installed_packages)),
|
||||
},
|
||||
indent=2,
|
||||
@@ -952,9 +771,6 @@ def parse_args():
|
||||
parser.add_argument("--venv-dir", required=True)
|
||||
parser.add_argument("--manifest", required=True)
|
||||
parser.add_argument("--models-dir", required=True)
|
||||
parser.add_argument("--torch-variant", choices=tuple(TORCH_INDEXES), required=True)
|
||||
parser.add_argument("--da3-commit", required=True)
|
||||
parser.add_argument("--ben2-commit", required=True)
|
||||
parser.add_argument("--status", required=True)
|
||||
parser.add_argument("--pypi-index", default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -21,42 +21,33 @@ MODELS = {
|
||||
"SMALL": ModelSpec(
|
||||
key="SMALL",
|
||||
label="DA3 Small",
|
||||
directory_name="da3_small",
|
||||
directory_name="DA3-SMALL",
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
license_name="Apache 2.0",
|
||||
description="Fastest any-view model with Apache 2.0 weights",
|
||||
description="Fastest any-view model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
r2_directory="DA3Small",
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="DA3-SMALL",
|
||||
r2_files=(
|
||||
("config.json", 1202, "a486e29e82b7ab4a7d4cefc1ea4526cfe2ae438a572c8ca98917cfbcde7447d2"),
|
||||
("model.safetensors", 137248940, "364492e38a3a06d221ac75da7f6621ada3f2361cd24fde11ba79091e9f40efcf"),
|
||||
("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"),
|
||||
),
|
||||
),
|
||||
"BASE": ModelSpec(
|
||||
key="BASE",
|
||||
label="DA3 Base",
|
||||
directory_name="da3_base",
|
||||
directory_name="DA3-BASE",
|
||||
repository="depth-anything/DA3-BASE",
|
||||
license_name="Apache 2.0",
|
||||
description="Balanced any-view model with Apache 2.0 weights",
|
||||
description="Balanced any-view model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
r2_directory="DA3Base",
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="DA3-BASE",
|
||||
r2_files=(
|
||||
("config.json", 1205, "5e34115ebc17bd2d8d43033c5f72e9446ac8833fd61d3fa160b7e67e0bb5b7b5"),
|
||||
("configuration.json", 74, "b684d553257b2ca379fd5dfdf96a6e4a633eec3b969c97119c9875465a7c07a8"),
|
||||
("model.safetensors", 541518028, "e01067dc1659613083d9145a9a2547ccdbe6ccbbf83c4fe7b3e8a4e2bdae78b5"),
|
||||
("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"),
|
||||
),
|
||||
),
|
||||
"LARGE_1_1": ModelSpec(
|
||||
key="LARGE_1_1",
|
||||
label="DA3 Large 1.1",
|
||||
directory_name="da3_large_1_1",
|
||||
repository="depth-anything/DA3-LARGE-1.1",
|
||||
license_name="CC BY-NC 4.0",
|
||||
description="Higher quality non-commercial model with larger VRAM use",
|
||||
supports_gs=False,
|
||||
enabled=False,
|
||||
),
|
||||
"GIANT_1_1": ModelSpec(
|
||||
key="GIANT_1_1",
|
||||
label="DA3 Giant 1.1",
|
||||
@@ -82,52 +73,76 @@ MODELS = {
|
||||
BEN2_MODEL = ModelSpec(
|
||||
key="BEN2_BASE",
|
||||
label="BEN2 Base",
|
||||
directory_name="ben2_base",
|
||||
repository="PramaLLC/BEN2",
|
||||
license_name="MIT",
|
||||
description="Fast high-quality background removal model",
|
||||
directory_name="BEN2-ONNX",
|
||||
repository="onnx-community/BEN2-ONNX",
|
||||
license_name="Apache 2.0",
|
||||
description="Fast high-quality FP16 ONNX background removal model",
|
||||
supports_gs=False,
|
||||
download_patterns=("config.json", "model.safetensors"),
|
||||
r2_directory="BEN2Base",
|
||||
r2_files=(
|
||||
("config.json", 124, "01ac876925577b52fa7c6af3f5b1d1420bc9b40016b6c25e2de89b38d8fff865"),
|
||||
("model.safetensors", 380577976, "ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb"),
|
||||
),
|
||||
)
|
||||
def moge2_model(key, label, directory_name, repository, description, r2_directory, size, checksum):
|
||||
return ModelSpec(
|
||||
key=key,
|
||||
label=label,
|
||||
directory_name=directory_name,
|
||||
repository=repository,
|
||||
license_name="MIT",
|
||||
description=description,
|
||||
supports_gs=False,
|
||||
download_patterns=("model.pt",),
|
||||
ready_patterns=("model.pt",),
|
||||
r2_directory=r2_directory,
|
||||
r2_files=(("model.pt", size, checksum),),
|
||||
download_patterns=("onnx/model_fp16.onnx",),
|
||||
ready_patterns=("onnx/model_fp16.onnx",),
|
||||
r2_directory="BEN2-ONNX",
|
||||
r2_files=((
|
||||
"onnx/model_fp16.onnx",
|
||||
219121675,
|
||||
"dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1",
|
||||
),),
|
||||
)
|
||||
|
||||
|
||||
MOGE2_MODELS = {
|
||||
"MOGE2_VITS_NORMAL": moge2_model(
|
||||
"MOGE2_VITS_NORMAL", "MoGe-2 ViT-S Normal", "moge2_vits_normal",
|
||||
"Ruicheng/moge-2-vits-normal", "Fastest MoGe-2 model", "MoGe2ViTSNormal",
|
||||
140550416, "79a16621928c2bf0ed04659218c55c01075e950507f40bb3332fb4c873d3e1dc",
|
||||
"MOGE2_VITS_NORMAL": ModelSpec(
|
||||
key="MOGE2_VITS_NORMAL",
|
||||
label="MoGe-2 ViT-S Normal",
|
||||
directory_name="moge-2-vits-normal-onnx",
|
||||
repository="Ruicheng/moge-2-vits-normal-onnx",
|
||||
license_name="MIT",
|
||||
description="Fastest MoGe-2 model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="moge-2-vits-normal-onnx",
|
||||
r2_files=((
|
||||
"model.onnx",
|
||||
140852051,
|
||||
"24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0",
|
||||
),),
|
||||
),
|
||||
"MOGE2_VITB_NORMAL": moge2_model(
|
||||
"MOGE2_VITB_NORMAL", "MoGe-2 ViT-B Normal", "moge2_vitb_normal",
|
||||
"Ruicheng/moge-2-vitb-normal", "Balanced MoGe-2 model", "MoGe2ViTBNormal",
|
||||
419110160, "16b8110e86d5dc5a849db120ca96ef3a223fd30b0c9146d1d81db504073da5f6",
|
||||
"MOGE2_VITB_NORMAL": ModelSpec(
|
||||
key="MOGE2_VITB_NORMAL",
|
||||
label="MoGe-2 ViT-B Normal",
|
||||
directory_name="moge-2-vitb-normal-onnx",
|
||||
repository="Ruicheng/moge-2-vitb-normal-onnx",
|
||||
license_name="MIT",
|
||||
description="Balanced MoGe-2 model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="moge-2-vitb-normal-onnx",
|
||||
r2_files=((
|
||||
"model.onnx",
|
||||
419411850,
|
||||
"bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21",
|
||||
),),
|
||||
),
|
||||
"MOGE2_VITL_NORMAL": moge2_model(
|
||||
"MOGE2_VITL_NORMAL", "MoGe-2 ViT-L Normal", "moge2_vitl_normal",
|
||||
"Ruicheng/moge-2-vitl-normal", "Highest-quality MoGe-2 model", "MoGe2ViTLNormal",
|
||||
1323815904, "280741fd09bc3f403ccff9967784c2a391b52d2c0742ae3efdb21d9f90cc1a01",
|
||||
"MOGE2_VITL_NORMAL": ModelSpec(
|
||||
key="MOGE2_VITL_NORMAL",
|
||||
label="MoGe-2 ViT-L Normal",
|
||||
directory_name="moge-2-vitl-normal-onnx",
|
||||
repository="Ruicheng/moge-2-vitl-normal-onnx",
|
||||
license_name="MIT",
|
||||
description="Highest-quality MoGe-2 model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="moge-2-vitl-normal-onnx",
|
||||
r2_files=((
|
||||
"model.onnx",
|
||||
1324265014,
|
||||
"afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
|
||||
),),
|
||||
),
|
||||
}
|
||||
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITB_NORMAL"]
|
||||
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"]
|
||||
DOWNLOADABLE_MODELS = {
|
||||
**MODELS,
|
||||
BEN2_MODEL.key: BEN2_MODEL,
|
||||
|
||||
165
src/snapform/moge_onnx.py
Normal file
165
src/snapform/moge_onnx.py
Normal file
@@ -0,0 +1,165 @@
|
||||
"""MoGe-2 ONNX inference and NumPy post-processing."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
|
||||
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
|
||||
if str(SCRIPT_DIRECTORY) not in sys.path:
|
||||
sys.path.insert(0, str(SCRIPT_DIRECTORY))
|
||||
|
||||
try:
|
||||
from .onnx_runtime import create_session as create_runtime_session
|
||||
from .onnx_runtime import provider_label, select_providers
|
||||
except ImportError:
|
||||
from onnx_runtime import create_session as create_runtime_session
|
||||
from onnx_runtime import provider_label, select_providers
|
||||
|
||||
|
||||
ONNX_FILENAME = "model.onnx"
|
||||
MIN_TOKENS = 1200
|
||||
MAX_TOKENS = 3600
|
||||
|
||||
|
||||
def is_onnx_model(model_directory):
|
||||
return (Path(model_directory) / ONNX_FILENAME).is_file()
|
||||
|
||||
|
||||
def create_session(model_directory, requested_device="auto"):
|
||||
path = Path(model_directory) / ONNX_FILENAME
|
||||
return create_runtime_session(path, requested_device)
|
||||
|
||||
|
||||
def normalized_view_plane_uv(width, height, dtype):
|
||||
import numpy as np
|
||||
|
||||
aspect_ratio = width / height
|
||||
span_x = aspect_ratio / (1.0 + aspect_ratio**2) ** 0.5
|
||||
span_y = 1.0 / (1.0 + aspect_ratio**2) ** 0.5
|
||||
u = np.linspace(
|
||||
-span_x * (width - 1) / width,
|
||||
span_x * (width - 1) / width,
|
||||
width,
|
||||
dtype=dtype,
|
||||
)
|
||||
v = np.linspace(
|
||||
-span_y * (height - 1) / height,
|
||||
span_y * (height - 1) / height,
|
||||
height,
|
||||
dtype=dtype,
|
||||
)
|
||||
u, v = np.meshgrid(u, v, indexing="xy")
|
||||
return np.stack((u, v), axis=-1)
|
||||
|
||||
|
||||
def solve_focal_shift(uv, points):
|
||||
import numpy as np
|
||||
|
||||
uv = uv.reshape(-1, 2)
|
||||
xy = points[..., :2].reshape(-1, 2)
|
||||
z = points[..., 2].reshape(-1)
|
||||
matrix = np.stack((xy, -uv), axis=-1).reshape(-1, 2)
|
||||
target = (uv * z[:, None]).reshape(-1)
|
||||
focal, shift = np.linalg.lstsq(matrix, target, rcond=None)[0]
|
||||
epsilon = np.finfo(np.float32).eps
|
||||
for _iteration in range(20):
|
||||
denominator = z + shift
|
||||
denominator = np.where(
|
||||
np.abs(denominator) < epsilon,
|
||||
np.copysign(epsilon, denominator),
|
||||
denominator,
|
||||
)
|
||||
projected = xy / denominator[:, None]
|
||||
residual = (focal * projected - uv).reshape(-1)
|
||||
jacobian = np.stack(
|
||||
(
|
||||
projected,
|
||||
-focal * xy / np.square(denominator)[:, None],
|
||||
),
|
||||
axis=-1,
|
||||
).reshape(-1, 2)
|
||||
delta = np.linalg.lstsq(jacobian, -residual, rcond=None)[0]
|
||||
focal += delta[0]
|
||||
shift += delta[1]
|
||||
if np.max(np.abs(delta)) < 1e-6:
|
||||
break
|
||||
return np.float32(focal), shift
|
||||
|
||||
|
||||
def recover_focal_shift(points, mask, downsample_size=(64, 64)):
|
||||
import numpy as np
|
||||
|
||||
height, width = points.shape[:2]
|
||||
uv = normalized_view_plane_uv(width, height, points.dtype)
|
||||
target_height, target_width = downsample_size
|
||||
y = np.floor(np.arange(target_height) * height / target_height).astype(int)
|
||||
x = np.floor(np.arange(target_width) * width / target_width).astype(int)
|
||||
points_low = points[y[:, None], x[None, :]]
|
||||
uv_low = uv[y[:, None], x[None, :]]
|
||||
mask_low = mask[y[:, None], x[None, :]]
|
||||
points_low = points_low[mask_low]
|
||||
uv_low = uv_low[mask_low]
|
||||
if len(points_low) < 2:
|
||||
return np.float32(1.0), np.float32(0.0)
|
||||
return solve_focal_shift(uv_low, points_low)
|
||||
|
||||
|
||||
def depth_to_points(depth, intrinsics):
|
||||
import numpy as np
|
||||
|
||||
height, width = depth.shape
|
||||
u = np.linspace(0.5 / width, 1.0 - 0.5 / width, width, dtype=depth.dtype)
|
||||
v = np.linspace(0.5 / height, 1.0 - 0.5 / height, height, dtype=depth.dtype)
|
||||
u, v = np.meshgrid(u, v, indexing="xy")
|
||||
x = (u - intrinsics[0, 2]) / intrinsics[0, 0] * depth
|
||||
y = (v - intrinsics[1, 2]) / intrinsics[1, 1] * depth
|
||||
return np.stack((x, y, depth), axis=-1)
|
||||
|
||||
|
||||
def postprocess(raw_output):
|
||||
import numpy as np
|
||||
|
||||
points = np.asarray(raw_output["points"][0], dtype=np.float32)
|
||||
normal = np.asarray(raw_output["normal"][0], dtype=np.float32)
|
||||
mask = np.asarray(raw_output["mask"][0] > 0.5)
|
||||
metric_scale = np.float32(np.asarray(raw_output["metric_scale"]).reshape(-1)[0])
|
||||
height, width = mask.shape
|
||||
aspect_ratio = width / height
|
||||
focal, shift = recover_focal_shift(points, mask)
|
||||
diagonal_factor = (1.0 + aspect_ratio**2) ** 0.5
|
||||
fx = focal / 2.0 * diagonal_factor / aspect_ratio
|
||||
fy = focal / 2.0 * diagonal_factor
|
||||
intrinsics = np.array(
|
||||
((fx, 0.0, 0.5), (0.0, fy, 0.5), (0.0, 0.0, 1.0)),
|
||||
dtype=np.float32,
|
||||
)
|
||||
depth = points[..., 2] + shift
|
||||
mask &= depth > 0.0
|
||||
depth = depth * metric_scale
|
||||
points = depth_to_points(depth, intrinsics)
|
||||
return {
|
||||
"points": points.astype(np.float32),
|
||||
"depth": depth.astype(np.float32),
|
||||
"normal": normal,
|
||||
"mask": mask,
|
||||
"intrinsics": intrinsics,
|
||||
}
|
||||
|
||||
|
||||
def infer(session, image, resolution_level):
|
||||
import numpy as np
|
||||
|
||||
level = min(max(int(resolution_level), 0), 9)
|
||||
num_tokens = int(MIN_TOKENS + level / 9.0 * (MAX_TOKENS - MIN_TOKENS))
|
||||
image_input = np.ascontiguousarray(image.transpose(2, 0, 1)[None], dtype=np.float32)
|
||||
input_feed = {
|
||||
"image": image_input,
|
||||
"num_tokens": np.asarray(num_tokens, dtype=np.int64),
|
||||
}
|
||||
output_names = {output.name for output in session.get_outputs()}
|
||||
scale_name = "metric_scale" if "metric_scale" in output_names else "scale"
|
||||
names = ("points", "normal", "mask", scale_name)
|
||||
values = session.run(list(names), input_feed)
|
||||
raw_output = dict(zip(names, values))
|
||||
raw_output["metric_scale"] = raw_output.pop(scale_name)
|
||||
return postprocess(raw_output)
|
||||
55
src/snapform/onnx_runtime.py
Normal file
55
src/snapform/onnx_runtime.py
Normal file
@@ -0,0 +1,55 @@
|
||||
"""Shared ONNX Runtime session helpers."""
|
||||
|
||||
|
||||
def select_providers(requested, available):
|
||||
available = set(available)
|
||||
providers_for_device = {
|
||||
"cuda": ("CUDAExecutionProvider", "DmlExecutionProvider"),
|
||||
"directml": ("DmlExecutionProvider",),
|
||||
"mps": ("CoreMLExecutionProvider",),
|
||||
"cpu": ("CPUExecutionProvider",),
|
||||
}
|
||||
if requested != "auto":
|
||||
candidates = providers_for_device[requested]
|
||||
for provider in candidates:
|
||||
if provider in available:
|
||||
return [provider]
|
||||
if "CPUExecutionProvider" in available:
|
||||
return ["CPUExecutionProvider"]
|
||||
raise RuntimeError(
|
||||
f"{requested.upper()} was requested but the ONNX Runtime "
|
||||
f"providers {', '.join(candidates)} are not available"
|
||||
)
|
||||
priorities = (
|
||||
"CUDAExecutionProvider",
|
||||
"CoreMLExecutionProvider",
|
||||
"DmlExecutionProvider",
|
||||
"OpenVINOExecutionProvider",
|
||||
"CPUExecutionProvider",
|
||||
)
|
||||
providers = [provider for provider in priorities if provider in available]
|
||||
if not providers:
|
||||
raise RuntimeError("ONNX Runtime has no usable execution provider")
|
||||
return providers
|
||||
|
||||
|
||||
def create_session(path, requested_device="auto"):
|
||||
import onnxruntime
|
||||
|
||||
providers = select_providers(
|
||||
requested_device,
|
||||
onnxruntime.get_available_providers(),
|
||||
)
|
||||
return onnxruntime.InferenceSession(str(path), providers=providers)
|
||||
|
||||
|
||||
def provider_label(session):
|
||||
provider = session.get_providers()[0]
|
||||
labels = {
|
||||
"CUDAExecutionProvider": "CUDA",
|
||||
"CoreMLExecutionProvider": "CoreML",
|
||||
"DmlExecutionProvider": "DirectML",
|
||||
"OpenVINOExecutionProvider": "OpenVINO",
|
||||
"CPUExecutionProvider": "CPU",
|
||||
}
|
||||
return labels.get(provider, provider)
|
||||
@@ -186,7 +186,6 @@ def draw_da3_input_settings(layout, settings, model_property):
|
||||
layout.prop(settings, "video_frames")
|
||||
layout.prop(settings, "process_res")
|
||||
layout.prop(settings, "process_res_method")
|
||||
layout.prop(settings, "use_ray_pose")
|
||||
layout.prop(settings, "ref_view_strategy")
|
||||
|
||||
|
||||
@@ -198,7 +197,7 @@ class Da3OutputDebugPanel(bpy.types.Panel):
|
||||
bl_region_type = "UI"
|
||||
bl_category = "SnapForm"
|
||||
bl_options = {"DEFAULT_CLOSED"}
|
||||
bl_order = 0
|
||||
bl_order = 1
|
||||
|
||||
def draw(self, context):
|
||||
layout = self.layout
|
||||
@@ -248,7 +247,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel):
|
||||
bl_region_type = "UI"
|
||||
bl_category = "SnapForm"
|
||||
bl_options = {"DEFAULT_CLOSED"}
|
||||
bl_order = 1
|
||||
bl_order = 0
|
||||
|
||||
def draw(self, context):
|
||||
layout = self.layout
|
||||
|
||||
@@ -76,7 +76,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
|
||||
name="Model",
|
||||
description="MoGe-2 model used for depth estimation",
|
||||
items=moge2_model_enum_items,
|
||||
default=1,
|
||||
default=0,
|
||||
)
|
||||
moge_export_format: bpy.props.EnumProperty(
|
||||
name="Output Format",
|
||||
@@ -110,7 +110,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
|
||||
)
|
||||
use_ray_pose: bpy.props.BoolProperty(
|
||||
name="Use Ray Pose",
|
||||
description="Use the slower and more accurate ray-based pose estimation",
|
||||
description="Legacy task compatibility; ONNX uses the camera decoder",
|
||||
default=False,
|
||||
)
|
||||
ref_view_strategy: bpy.props.EnumProperty(
|
||||
|
||||
@@ -74,9 +74,6 @@ class ServiceLock:
|
||||
|
||||
class ModelCache:
|
||||
def __init__(self):
|
||||
import torch
|
||||
|
||||
self.torch = torch
|
||||
self.ben2_model = None
|
||||
self.ben2_key = None
|
||||
self.da3_model = None
|
||||
@@ -85,75 +82,51 @@ class ModelCache:
|
||||
self.moge_key = None
|
||||
|
||||
def get_ben2(self, model_dir, device):
|
||||
key = (str(Path(model_dir).resolve()), device.type)
|
||||
requested_device = getattr(device, "type", device)
|
||||
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
|
||||
if self.ben2_model is not None and self.ben2_key == key:
|
||||
return self.ben2_model, 0.0
|
||||
self.release_ben2()
|
||||
try:
|
||||
return self._load_ben2(model_dir, device, key)
|
||||
except self.torch.cuda.OutOfMemoryError:
|
||||
self.release_da3()
|
||||
self.release_moge()
|
||||
return self._load_ben2(model_dir, device, key)
|
||||
return self._load_ben2(model_dir, requested_device, key)
|
||||
|
||||
def _load_ben2(self, model_dir, device, key):
|
||||
from ben2 import BEN_Base
|
||||
def _load_ben2(self, model_dir, requested_device, key):
|
||||
from ben2_onnx import create_session
|
||||
|
||||
started = time.perf_counter()
|
||||
model = BEN_Base.from_pretrained(model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
model = create_session(model_dir, requested_device)
|
||||
self.ben2_model = model
|
||||
self.ben2_key = key
|
||||
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
||||
return model, elapsed_ms
|
||||
|
||||
def get_da3(self, model_dir, device):
|
||||
key = (str(Path(model_dir).resolve()), device.type)
|
||||
requested_device = getattr(device, "type", device)
|
||||
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
|
||||
if self.moge_model is not None:
|
||||
self.release_moge()
|
||||
if self.da3_model is not None and self.da3_key == key:
|
||||
return self.da3_model
|
||||
self.release_da3()
|
||||
try:
|
||||
return self._load_da3(model_dir, device, key)
|
||||
except self.torch.cuda.OutOfMemoryError:
|
||||
self.release_ben2()
|
||||
self.release_moge()
|
||||
return self._load_da3(model_dir, device, key)
|
||||
return self._load_da3(model_dir, requested_device, key)
|
||||
|
||||
def _load_da3(self, model_dir, device, key):
|
||||
from depth_anything_3.api import DepthAnything3
|
||||
def _load_da3(self, model_dir, requested_device, key):
|
||||
from da3_onnx import create_session
|
||||
|
||||
model = DepthAnything3.from_pretrained(model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
model = create_session(model_dir, requested_device)
|
||||
self.da3_model = model
|
||||
self.da3_key = key
|
||||
return model
|
||||
|
||||
def get_moge(self, model_dir, device):
|
||||
key = (str(Path(model_dir).resolve()), device.type)
|
||||
def get_moge_onnx(self, model_dir, requested_device):
|
||||
from moge_onnx import create_session
|
||||
|
||||
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
|
||||
if self.da3_model is not None:
|
||||
self.release_da3()
|
||||
if self.moge_model is not None and self.moge_key == key:
|
||||
return self.moge_model
|
||||
self.release_moge()
|
||||
try:
|
||||
return self._load_moge(model_dir, device, key)
|
||||
except self.torch.cuda.OutOfMemoryError:
|
||||
self.release_da3()
|
||||
self.release_ben2()
|
||||
return self._load_moge(model_dir, device, key)
|
||||
|
||||
def _load_moge(self, model_dir, device, key):
|
||||
worker.prepare_moge_runtime()
|
||||
from moge.model.v2 import MoGeModel
|
||||
|
||||
checkpoint = Path(model_dir) / "model.pt"
|
||||
model = MoGeModel.from_pretrained(checkpoint)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
model = create_session(model_dir, requested_device)
|
||||
self.moge_model = model
|
||||
self.moge_key = key
|
||||
return model
|
||||
@@ -201,8 +174,6 @@ class ModelCache:
|
||||
|
||||
def _release_cuda_cache(self):
|
||||
gc.collect()
|
||||
if self.torch.cuda.is_available():
|
||||
self.torch.cuda.empty_cache()
|
||||
|
||||
|
||||
class ServiceState:
|
||||
|
||||
@@ -6,11 +6,17 @@ import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
|
||||
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
|
||||
if str(SCRIPT_DIRECTORY) not in sys.path:
|
||||
sys.path.insert(0, str(SCRIPT_DIRECTORY))
|
||||
|
||||
|
||||
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
|
||||
VIDEO_SUFFIXES = {
|
||||
".mp4",
|
||||
@@ -434,6 +440,92 @@ def export_prediction_attributes(prediction, output_directory):
|
||||
return output_path
|
||||
|
||||
|
||||
def export_da3_mini_npz(prediction, output_directory):
|
||||
import numpy as np
|
||||
|
||||
values = {"depth": np.round(prediction.depth, 8)}
|
||||
for name in ("conf", "extrinsics", "intrinsics"):
|
||||
value = getattr(prediction, name, None)
|
||||
if value is not None:
|
||||
values[name] = np.round(value, 2) if name == "conf" else value
|
||||
output_path = Path(output_directory) / "results.npz"
|
||||
np.savez_compressed(output_path, **values)
|
||||
return output_path
|
||||
|
||||
|
||||
def export_da3_depth_preview(prediction, output_directory):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
target = Path(output_directory) / "depth_vis"
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
for index, depth in enumerate(prediction.depth):
|
||||
valid = np.isfinite(depth) & (depth > 0.0)
|
||||
normalized = np.zeros(depth.shape, dtype=np.float32)
|
||||
if valid.any():
|
||||
low, high = np.percentile(depth[valid], (2.0, 98.0))
|
||||
normalized[valid] = np.clip(
|
||||
(depth[valid] - low) / max(float(high - low), 1e-6),
|
||||
0.0,
|
||||
1.0,
|
||||
)
|
||||
color = np.stack(
|
||||
(
|
||||
normalized,
|
||||
1.0 - np.abs(normalized * 2.0 - 1.0),
|
||||
1.0 - normalized,
|
||||
),
|
||||
axis=-1,
|
||||
)
|
||||
color = np.asarray(color * 255.0, dtype=np.uint8)
|
||||
source = np.asarray(prediction.processed_images[index], dtype=np.uint8)
|
||||
combined = np.concatenate((source, color), axis=1)
|
||||
Image.fromarray(combined, mode="RGB").save(
|
||||
target / f"{index:04d}.jpg",
|
||||
quality=95,
|
||||
)
|
||||
|
||||
|
||||
def export_da3_point_cloud(
|
||||
prediction,
|
||||
output_directory,
|
||||
confidence_percentile=40.0,
|
||||
num_max_points=1_000_000,
|
||||
):
|
||||
import numpy as np
|
||||
import trimesh
|
||||
|
||||
depth = np.asarray(prediction.depth, dtype=np.float32)
|
||||
confidence = np.asarray(prediction.conf, dtype=np.float32)
|
||||
threshold = confidence_percentile_threshold(
|
||||
depth,
|
||||
confidence,
|
||||
confidence_percentile,
|
||||
)
|
||||
point_sets = []
|
||||
color_sets = []
|
||||
for index in range(depth.shape[0]):
|
||||
points = unproject_depth_frame(
|
||||
depth[index],
|
||||
prediction.intrinsics[index],
|
||||
prediction.extrinsics[index],
|
||||
)
|
||||
valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0)
|
||||
if threshold is not None:
|
||||
valid &= confidence[index] >= threshold
|
||||
point_sets.append(points[valid])
|
||||
color_sets.append(prediction.processed_images[index][valid])
|
||||
points = np.concatenate(point_sets, axis=0)
|
||||
colors = np.concatenate(color_sets, axis=0)
|
||||
maximum = max(int(num_max_points), 1)
|
||||
if len(points) > maximum:
|
||||
indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64)
|
||||
points = points[indices]
|
||||
colors = colors[indices]
|
||||
scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors))
|
||||
scene.export(Path(output_directory) / SCENE_FILENAME)
|
||||
|
||||
|
||||
def remove_background_faces(faces, foreground_mask):
|
||||
"""Return only triangles whose vertices all belong to the foreground."""
|
||||
foreground_points = foreground_mask.reshape(-1)
|
||||
@@ -1133,29 +1225,10 @@ def format_rate(rate):
|
||||
return f"{size:.1f} {units[unit_index]}"
|
||||
|
||||
|
||||
def select_device(requested, torch):
|
||||
if requested == "cpu":
|
||||
return torch.device("cpu")
|
||||
if requested == "cuda":
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA was requested but is not available")
|
||||
return torch.device("cuda")
|
||||
if requested == "mps":
|
||||
if not torch.backends.mps.is_available():
|
||||
raise RuntimeError("MPS was requested but is not available")
|
||||
return torch.device("mps")
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
if torch.backends.mps.is_available():
|
||||
return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def infer_foreground(
|
||||
input_path,
|
||||
model_dir,
|
||||
device,
|
||||
torch,
|
||||
model_cache=None,
|
||||
cancel_check=None,
|
||||
):
|
||||
@@ -1165,14 +1238,13 @@ def infer_foreground(
|
||||
check_cancelled(cancel_check)
|
||||
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
|
||||
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
|
||||
owns_model = model_cache is None
|
||||
if owns_model:
|
||||
from ben2 import BEN_Base
|
||||
from ben2_onnx import create_session, infer
|
||||
|
||||
owns_model = model_cache is None
|
||||
requested_device = getattr(device, "type", device)
|
||||
if owns_model:
|
||||
load_started = time.perf_counter()
|
||||
model = BEN_Base.from_pretrained(model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
model = create_session(model_dir, requested_device)
|
||||
load_ms = (time.perf_counter() - load_started) * 1000.0
|
||||
else:
|
||||
model, load_ms = model_cache.get_ben2(model_dir, device)
|
||||
@@ -1180,22 +1252,10 @@ def infer_foreground(
|
||||
check_cancelled(cancel_check)
|
||||
with Image.open(input_path) as opened_image:
|
||||
image = opened_image.convert("RGB")
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
inference_started = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
foreground = model.inference(
|
||||
image,
|
||||
refine_foreground=False,
|
||||
)
|
||||
foreground = infer(model, image)
|
||||
foreground = preserve_source_rgb(image, foreground)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
inference_ms = (time.perf_counter() - inference_started) * 1000.0
|
||||
if owns_model:
|
||||
del model
|
||||
if device.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
check_cancelled(cancel_check)
|
||||
return foreground, load_ms, inference_ms
|
||||
|
||||
@@ -1216,7 +1276,6 @@ def infer_foreground_frames(
|
||||
images,
|
||||
model_dir,
|
||||
device,
|
||||
torch,
|
||||
output_directory,
|
||||
status_path,
|
||||
start_progress,
|
||||
@@ -1238,7 +1297,6 @@ def infer_foreground_frames(
|
||||
image_path,
|
||||
model_dir,
|
||||
device,
|
||||
torch,
|
||||
model_cache=model_cache,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
@@ -1260,14 +1318,9 @@ def remove_background(args, model_cache=None, cancel_check=None):
|
||||
write_status(args.status, 0.02, "Starting Environment")
|
||||
check_cancelled(cancel_check)
|
||||
|
||||
if model_cache is None:
|
||||
import torch
|
||||
else:
|
||||
torch = model_cache.torch
|
||||
|
||||
output_directory = Path(args.output_dir)
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
device = select_device(args.device, torch)
|
||||
device = args.device
|
||||
|
||||
video_frames = getattr(args, "video_frames", 16)
|
||||
start_frame = getattr(args, "start_frame", 1)
|
||||
@@ -1279,14 +1332,13 @@ def remove_background(args, model_cache=None, cancel_check=None):
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
try:
|
||||
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
|
||||
write_status(args.status, 0.12, f"Loading BEN2 for {device.upper()}")
|
||||
if len(images) == 1:
|
||||
write_status(args.status, 0.35, "Removing background")
|
||||
foreground, load_ms, inference_ms = infer_foreground(
|
||||
images[0],
|
||||
args.model_dir,
|
||||
device,
|
||||
torch,
|
||||
model_cache=model_cache,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
@@ -1305,7 +1357,6 @@ def remove_background(args, model_cache=None, cancel_check=None):
|
||||
images,
|
||||
args.model_dir,
|
||||
device,
|
||||
torch,
|
||||
output_directory,
|
||||
args.status,
|
||||
0.12,
|
||||
@@ -1517,41 +1568,33 @@ def export_moge_debug_mesh(
|
||||
export_format,
|
||||
):
|
||||
import numpy as np
|
||||
import utils3d
|
||||
from moge.utils.io import save_glb, save_ply
|
||||
import trimesh
|
||||
|
||||
clean_mask = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04)
|
||||
uv = utils3d.np.uv_map(*depth.shape)
|
||||
result = utils3d.np.build_mesh_from_map(
|
||||
clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04)
|
||||
faces = mesh_triangle_faces(
|
||||
points,
|
||||
image,
|
||||
uv,
|
||||
normal,
|
||||
mask=clean_mask,
|
||||
tri=True,
|
||||
depth,
|
||||
np.ones(depth.shape, dtype=np.float32),
|
||||
None,
|
||||
False,
|
||||
0.0,
|
||||
vertex_mask=clean_mask,
|
||||
)
|
||||
faces, vertices, colors, vertex_uvs, vertex_normals = result
|
||||
vertices = points.reshape(-1, 3)
|
||||
colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8)
|
||||
vertex_normals = normal.reshape(-1, 3)
|
||||
vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
||||
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
|
||||
vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
|
||||
if export_format == "glb":
|
||||
save_glb(
|
||||
output_directory / "mesh.glb",
|
||||
vertices,
|
||||
faces,
|
||||
vertex_uvs,
|
||||
(image * 255.0).astype(np.uint8),
|
||||
vertex_normals,
|
||||
)
|
||||
return
|
||||
if export_format == "ply":
|
||||
save_ply(
|
||||
output_directory / "pointcloud.ply",
|
||||
vertices,
|
||||
np.zeros((0, 3), dtype=np.int32),
|
||||
colors,
|
||||
vertex_normals,
|
||||
mesh = trimesh.Trimesh(
|
||||
vertices=vertices,
|
||||
faces=faces,
|
||||
vertex_colors=colors,
|
||||
vertex_normals=vertex_normals,
|
||||
process=False,
|
||||
)
|
||||
mesh.remove_unreferenced_vertices()
|
||||
filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply"
|
||||
mesh.export(Path(output_directory) / filename)
|
||||
|
||||
|
||||
def export_moge_surface_archive(
|
||||
@@ -1564,7 +1607,6 @@ def export_moge_surface_archive(
|
||||
foreground_mask=None,
|
||||
):
|
||||
import numpy as np
|
||||
import utils3d
|
||||
|
||||
height, width = depth.shape
|
||||
finite_mask = np.isfinite(depth) & (depth > 0.0)
|
||||
@@ -1572,34 +1614,28 @@ def export_moge_surface_archive(
|
||||
if foreground_mask is not None:
|
||||
finite_mask &= foreground_mask
|
||||
depth_discontinuity = depth_discontinuity_ratio(depth)
|
||||
result = utils3d.np.build_mesh_from_map(
|
||||
faces = mesh_triangle_faces(
|
||||
points,
|
||||
utils3d.np.uv_map(height, width),
|
||||
normal,
|
||||
mask.astype(np.float32),
|
||||
depth_discontinuity.astype(np.float32),
|
||||
depth,
|
||||
mask=finite_mask,
|
||||
tri=True,
|
||||
np.ones(depth.shape, dtype=np.float32),
|
||||
None,
|
||||
False,
|
||||
0.0,
|
||||
vertex_mask=finite_mask,
|
||||
)
|
||||
(
|
||||
faces,
|
||||
vertices,
|
||||
vertex_uvs,
|
||||
_vertex_normals,
|
||||
vertex_mask,
|
||||
vertex_depth_discontinuity,
|
||||
point_depth,
|
||||
) = result
|
||||
vertex_uvs = vertex_uvs * np.array([1.0, -1.0]) + np.array([0.0, 1.0])
|
||||
vertices = points.reshape(-1, 3)
|
||||
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
|
||||
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
|
||||
uu, vv = np.meshgrid(u, v)
|
||||
vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2)
|
||||
pixel_intrinsics = intrinsics.copy()
|
||||
pixel_intrinsics[0] *= width
|
||||
pixel_intrinsics[1] *= height
|
||||
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
|
||||
point_depth = point_depth.astype(np.float32)
|
||||
face_mask = vertex_mask[faces].all(axis=1)
|
||||
face_depth_discontinuity = vertex_depth_discontinuity[faces].max(axis=1)
|
||||
masked_depth = point_depth[vertex_mask.astype(bool)]
|
||||
point_depth = depth.reshape(-1).astype(np.float32)
|
||||
face_mask = mask.reshape(-1)[faces].all(axis=1)
|
||||
face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
|
||||
masked_depth = point_depth[mask.reshape(-1).astype(bool)]
|
||||
reference_depth = np.median(masked_depth if len(masked_depth) else point_depth)
|
||||
np.savez(
|
||||
output_path,
|
||||
@@ -1619,14 +1655,13 @@ def export_moge_surface_archive(
|
||||
def run_moge(args, model_cache=None, cancel_check=None):
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from moge_onnx import (
|
||||
create_session,
|
||||
infer as infer_moge_onnx,
|
||||
provider_label,
|
||||
)
|
||||
|
||||
prepare_moge_runtime()
|
||||
write_status(args.status, 0.02, "Starting MoGe-2")
|
||||
if model_cache is None:
|
||||
import torch
|
||||
from moge.model.v2 import MoGeModel
|
||||
else:
|
||||
torch = model_cache.torch
|
||||
output_directory = Path(args.output_dir)
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
images, temporary_directory = prepare_input(
|
||||
@@ -1637,15 +1672,15 @@ def run_moge(args, model_cache=None, cancel_check=None):
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
try:
|
||||
device = select_device(args.device, torch)
|
||||
write_status(args.status, 0.08, f"Loading MoGe-2 on {device.type}")
|
||||
if model_cache is None:
|
||||
checkpoint = Path(args.model_dir) / "model.pt"
|
||||
model = MoGeModel.from_pretrained(checkpoint)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
model = create_session(args.model_dir, args.device)
|
||||
else:
|
||||
model = model_cache.get_moge(args.model_dir, device)
|
||||
model = model_cache.get_moge_onnx(args.model_dir, args.device)
|
||||
write_status(
|
||||
args.status,
|
||||
0.08,
|
||||
f"Loading MoGe-2 on {provider_label(model)}",
|
||||
)
|
||||
|
||||
depths = []
|
||||
normals = []
|
||||
@@ -1662,21 +1697,17 @@ def run_moge(args, model_cache=None, cancel_check=None):
|
||||
)
|
||||
with Image.open(image_path) as opened:
|
||||
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
|
||||
image_tensor = torch.from_numpy(image).permute(2, 0, 1).to(device)
|
||||
source_images.append(image)
|
||||
with torch.inference_mode():
|
||||
prediction = model.infer(
|
||||
image_tensor,
|
||||
resolution_level=int(args.resolution_level),
|
||||
force_projection=True,
|
||||
apply_mask=False,
|
||||
use_fp16=device.type == "cuda",
|
||||
prediction = infer_moge_onnx(
|
||||
model,
|
||||
image,
|
||||
int(args.resolution_level),
|
||||
)
|
||||
depth = prediction["depth"].detach().float().cpu().numpy()
|
||||
normal = prediction["normal"].detach().float().cpu().numpy()
|
||||
mask = prediction["mask"].detach().cpu().numpy().astype(bool)
|
||||
point = prediction["points"].detach().float().cpu().numpy()
|
||||
intrinsic = prediction["intrinsics"].detach().float().cpu().numpy()
|
||||
depth = prediction["depth"]
|
||||
normal = prediction["normal"]
|
||||
mask = prediction["mask"].astype(bool)
|
||||
point = prediction["points"]
|
||||
intrinsic = prediction["intrinsics"]
|
||||
depths.append(depth)
|
||||
normals.append(normal)
|
||||
masks.append(mask)
|
||||
@@ -1727,8 +1758,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
|
||||
foreground, _, _ = infer_foreground(
|
||||
images[0],
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
torch,
|
||||
args.device,
|
||||
model_cache=model_cache,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
@@ -1743,8 +1773,7 @@ def run_moge(args, model_cache=None, cancel_check=None):
|
||||
infer_foreground_frames(
|
||||
images,
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
torch,
|
||||
args.device,
|
||||
output_directory,
|
||||
args.status,
|
||||
0.87,
|
||||
@@ -1781,21 +1810,11 @@ def run_moge(args, model_cache=None, cancel_check=None):
|
||||
shutil.rmtree(temporary_directory, ignore_errors=True)
|
||||
|
||||
|
||||
def prepare_moge_runtime():
|
||||
"""Use PyTorch SDPA because xFormers does not support every CUDA GPU."""
|
||||
os.environ["XFORMERS_DISABLED"] = "1"
|
||||
|
||||
|
||||
def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
write_status(args.status, 0.02, "Starting Environment")
|
||||
check_cancelled(cancel_check)
|
||||
|
||||
if model_cache is None:
|
||||
import torch
|
||||
from depth_anything_3.api import DepthAnything3
|
||||
else:
|
||||
torch = model_cache.torch
|
||||
from depth_anything_3.utils.export import export
|
||||
from da3_onnx import create_session, infer
|
||||
from onnx_runtime import provider_label
|
||||
|
||||
output_directory = Path(args.output_dir)
|
||||
output_directory.mkdir(parents=True, exist_ok=True)
|
||||
@@ -1810,7 +1829,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
|
||||
try:
|
||||
check_cancelled(cancel_check)
|
||||
device = select_device(args.device, torch)
|
||||
remove_background_model_dir = getattr(
|
||||
args,
|
||||
"remove_background_model_dir",
|
||||
@@ -1823,8 +1841,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
foreground, load_ms, inference_ms = infer_foreground(
|
||||
images[0],
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
torch,
|
||||
args.device,
|
||||
model_cache=model_cache,
|
||||
cancel_check=cancel_check,
|
||||
)
|
||||
@@ -1841,8 +1858,7 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
load_ms, inference_ms = infer_foreground_frames(
|
||||
images,
|
||||
remove_background_model_dir,
|
||||
device,
|
||||
torch,
|
||||
args.device,
|
||||
output_directory,
|
||||
args.status,
|
||||
0.12,
|
||||
@@ -1858,13 +1874,16 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
f"BEN2 inference {inference_ms:.0f} ms; "
|
||||
f"model load {load_ms:.0f} ms"
|
||||
)
|
||||
write_status(args.status, 0.15, f"Loading model on {device.type}")
|
||||
write_status(args.status, 0.15, "Loading DA3 ONNX model")
|
||||
if model_cache is None:
|
||||
model = DepthAnything3.from_pretrained(args.model_dir)
|
||||
model = model.to(device=device)
|
||||
model.eval()
|
||||
model = create_session(args.model_dir, args.device)
|
||||
else:
|
||||
model = model_cache.get_da3(args.model_dir, device)
|
||||
model = model_cache.get_da3(args.model_dir, args.device)
|
||||
write_status(
|
||||
args.status,
|
||||
0.20,
|
||||
f"Loaded DA3 on {provider_label(model)}",
|
||||
)
|
||||
|
||||
check_cancelled(cancel_check)
|
||||
write_status(
|
||||
@@ -1872,22 +1891,9 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
0.25,
|
||||
f"Estimating depth for {len(images)} frame(s)",
|
||||
)
|
||||
export_kwargs = {}
|
||||
if args.export_format == "glb":
|
||||
export_kwargs["glb"] = {
|
||||
"conf_thresh_percentile": args.conf_thresh_percentile,
|
||||
"num_max_points": args.num_max_points,
|
||||
"show_cameras": bool(args.show_cameras),
|
||||
"export_depth_vis": False,
|
||||
}
|
||||
if args.export_format == "gs_ply":
|
||||
export_kwargs["gs_ply"] = {
|
||||
"gs_views_interval": args.gs_views_interval,
|
||||
}
|
||||
with torch.inference_mode():
|
||||
prediction = model.inference(
|
||||
prediction = infer(
|
||||
model,
|
||||
[str(path) for path in images],
|
||||
infer_gs=args.export_format == "gs_ply",
|
||||
process_res=args.process_res,
|
||||
process_res_method=args.process_res_method,
|
||||
use_ray_pose=bool(args.use_ray_pose),
|
||||
@@ -1921,12 +1927,20 @@ def reconstruct(args, model_cache=None, cancel_check=None):
|
||||
),
|
||||
progress_callback=mesh_progress,
|
||||
)
|
||||
else:
|
||||
export(
|
||||
elif args.export_format == "mini_npz":
|
||||
export_da3_mini_npz(prediction, output_directory)
|
||||
elif args.export_format == "depth_vis":
|
||||
export_da3_depth_preview(prediction, output_directory)
|
||||
elif args.export_format == "glb":
|
||||
export_da3_point_cloud(
|
||||
prediction,
|
||||
args.export_format,
|
||||
str(output_directory),
|
||||
**export_kwargs,
|
||||
output_directory,
|
||||
args.conf_thresh_percentile,
|
||||
args.num_max_points,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"DA3 ONNX does not support export format: {args.export_format}"
|
||||
)
|
||||
if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
|
||||
write_status(
|
||||
@@ -1989,7 +2003,7 @@ def parser():
|
||||
background.add_argument("--status", required=True)
|
||||
background.add_argument(
|
||||
"--device",
|
||||
choices=("auto", "cuda", "mps", "cpu"),
|
||||
choices=("auto", "cuda", "directml", "mps", "cpu"),
|
||||
required=True,
|
||||
)
|
||||
background.add_argument("--video-frames", type=int, default=16)
|
||||
@@ -2040,7 +2054,7 @@ def parser():
|
||||
reconstruction.add_argument("--start-frame", type=int, default=1)
|
||||
reconstruction.add_argument(
|
||||
"--device",
|
||||
choices=("auto", "cuda", "mps", "cpu"),
|
||||
choices=("auto", "cuda", "directml", "mps", "cpu"),
|
||||
required=True,
|
||||
)
|
||||
reconstruction.add_argument(
|
||||
|
||||
194
sync_models.py
Normal file
194
sync_models.py
Normal file
@@ -0,0 +1,194 @@
|
||||
"""Prepare SnapForm ONNX models and sync them to OmooLab R2."""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
MODELS_DIR = PROJECT_ROOT / "models"
|
||||
R2_ENDPOINT = "https://models.omoolab.xyz"
|
||||
CHUNK_SIZE = 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelFile:
|
||||
folder: str
|
||||
size: int
|
||||
sha256: str
|
||||
sources: tuple[str, ...]
|
||||
|
||||
@property
|
||||
def destination(self):
|
||||
filename = "model_fp16.onnx" if self.folder.endswith("/onnx") else "model.onnx"
|
||||
return MODELS_DIR / self.folder / filename
|
||||
|
||||
|
||||
def hf_url(repository, filename):
|
||||
return f"https://huggingface.co/{repository}/resolve/main/{filename}"
|
||||
|
||||
|
||||
def r2_url(folder, filename="model.onnx"):
|
||||
return f"{R2_ENDPOINT}/{folder}/{filename}"
|
||||
|
||||
|
||||
MODEL_FILES = (
|
||||
ModelFile(
|
||||
"DA3-SMALL",
|
||||
105484029,
|
||||
"7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631",
|
||||
(r2_url("DA3-SMALL"),),
|
||||
),
|
||||
ModelFile(
|
||||
"DA3-BASE",
|
||||
412856697,
|
||||
"31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6",
|
||||
(r2_url("DA3-BASE"),),
|
||||
),
|
||||
ModelFile(
|
||||
"BEN2-ONNX/onnx",
|
||||
219121675,
|
||||
"dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1",
|
||||
(
|
||||
hf_url("onnx-community/BEN2-ONNX", "onnx/model_fp16.onnx"),
|
||||
r2_url("BEN2-ONNX", "onnx/model_fp16.onnx"),
|
||||
),
|
||||
),
|
||||
ModelFile(
|
||||
"moge-2-vits-normal-onnx",
|
||||
140852051,
|
||||
"24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0",
|
||||
(
|
||||
hf_url("Ruicheng/moge-2-vits-normal-onnx", "model.onnx"),
|
||||
r2_url("moge-2-vits-normal-onnx"),
|
||||
),
|
||||
),
|
||||
ModelFile(
|
||||
"moge-2-vitb-normal-onnx",
|
||||
419411850,
|
||||
"bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21",
|
||||
(
|
||||
hf_url("Ruicheng/moge-2-vitb-normal-onnx", "model.onnx"),
|
||||
r2_url("moge-2-vitb-normal-onnx"),
|
||||
),
|
||||
),
|
||||
ModelFile(
|
||||
"moge-2-vitl-normal-onnx",
|
||||
1324265014,
|
||||
"afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
|
||||
(
|
||||
hf_url("Ruicheng/moge-2-vitl-normal-onnx", "model.onnx"),
|
||||
r2_url("moge-2-vitl-normal-onnx"),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def file_sha256(path):
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open("rb") as stream:
|
||||
while chunk := stream.read(CHUNK_SIZE):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def verify_existing(model):
|
||||
path = model.destination
|
||||
if not path.is_file():
|
||||
return False
|
||||
if path.stat().st_size != model.size:
|
||||
raise RuntimeError(f"Existing model has an unexpected size: {path}")
|
||||
if file_sha256(path) != model.sha256:
|
||||
raise RuntimeError(f"Existing model has an unexpected SHA-256: {path}")
|
||||
print(f"Skip existing: {path.relative_to(PROJECT_ROOT)}")
|
||||
return True
|
||||
|
||||
|
||||
def download_once(model, url):
|
||||
destination = model.destination
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = destination.with_suffix(".onnx.part")
|
||||
temporary.unlink(missing_ok=True)
|
||||
request = Request(url, headers={"User-Agent": "SnapForm model sync"})
|
||||
print(f"Download: {url}")
|
||||
downloaded = 0
|
||||
digest = hashlib.sha256()
|
||||
try:
|
||||
with urlopen(request, timeout=60) as response, temporary.open("wb") as output:
|
||||
while chunk := response.read(CHUNK_SIZE):
|
||||
output.write(chunk)
|
||||
digest.update(chunk)
|
||||
downloaded += len(chunk)
|
||||
print(
|
||||
f"\r {downloaded / 1024**2:.1f} / "
|
||||
f"{model.size / 1024**2:.1f} MiB",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
print()
|
||||
if downloaded != model.size:
|
||||
raise RuntimeError(f"size mismatch: {downloaded} != {model.size}")
|
||||
if digest.hexdigest() != model.sha256:
|
||||
raise RuntimeError("SHA-256 mismatch")
|
||||
temporary.replace(destination)
|
||||
except Exception:
|
||||
temporary.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def prepare_model(model):
|
||||
if verify_existing(model):
|
||||
return
|
||||
failures = []
|
||||
for url in model.sources:
|
||||
try:
|
||||
download_once(model, url)
|
||||
return
|
||||
except Exception as error:
|
||||
failures.append(f"{url}: {error}")
|
||||
print(f"Source failed: {error}", file=sys.stderr)
|
||||
raise RuntimeError("All sources failed:\n" + "\n".join(failures))
|
||||
|
||||
|
||||
def sync_models():
|
||||
executable = shutil.which("rclone")
|
||||
if executable is None:
|
||||
raise RuntimeError("rclone is not available on PATH")
|
||||
subprocess.run(
|
||||
[
|
||||
executable,
|
||||
"sync",
|
||||
"models",
|
||||
"omoolab-r2:ai-models",
|
||||
"--progress",
|
||||
],
|
||||
cwd=PROJECT_ROOT,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--download-only",
|
||||
action="store_true",
|
||||
help="Prepare and verify models without running rclone",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
for model in MODEL_FILES:
|
||||
prepare_model(model)
|
||||
if not args.download_only:
|
||||
sync_models()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -280,6 +280,10 @@ class AddonTest(unittest.TestCase):
|
||||
annotations["snapform_process_res_factor"]["subtype"],
|
||||
"PERCENTAGE",
|
||||
)
|
||||
self.assertEqual(
|
||||
annotations["snapform_model"]["default"],
|
||||
0,
|
||||
)
|
||||
|
||||
def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self):
|
||||
image_mesh = self.addon.image_mesh
|
||||
@@ -644,8 +648,11 @@ class AddonTest(unittest.TestCase):
|
||||
def test_conversion_dialog_exposes_requested_settings(self):
|
||||
operator = self.addon.ConvertImageToDepthPlane
|
||||
|
||||
self.assertEqual(operator.__annotations__["snapform_model"]["default"], 1)
|
||||
self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"])
|
||||
self.assertEqual(
|
||||
operator.__annotations__["snapform_model"]["default"],
|
||||
0,
|
||||
)
|
||||
self.assertFalse(operator.__annotations__["snapform_use_ray_pose"]["default"])
|
||||
self.assertEqual(
|
||||
operator.__annotations__["snapform_process_res_factor"]["default"],
|
||||
100.0,
|
||||
@@ -1094,7 +1101,6 @@ class AddonTest(unittest.TestCase):
|
||||
for snapform_model, snapform_ray_pose in (
|
||||
("SMALL", True),
|
||||
("BASE", False),
|
||||
("LARGE_1_1", True),
|
||||
("GIANT_1_1", False),
|
||||
("NESTED_GIANT_LARGE_1_1", True),
|
||||
):
|
||||
@@ -1947,8 +1953,8 @@ class AddonTest(unittest.TestCase):
|
||||
self.addon.DebugPanel.bl_options,
|
||||
)
|
||||
panels = (
|
||||
self.addon.Da3OutputDebugPanel,
|
||||
self.addon.Moge2ApiDebugPanel,
|
||||
self.addon.Da3OutputDebugPanel,
|
||||
self.addon.Ben2RemovalDebugPanel,
|
||||
)
|
||||
for panel in panels:
|
||||
@@ -1956,7 +1962,7 @@ class AddonTest(unittest.TestCase):
|
||||
self.assertIn("DEFAULT_CLOSED", panel.bl_options)
|
||||
self.assertEqual(
|
||||
[panel.bl_label for panel in panels],
|
||||
["DA3", "MoGe-2", "BEN2"],
|
||||
["MoGe-2", "DA3", "BEN2"],
|
||||
)
|
||||
|
||||
def test_service_panel_exposes_shared_service_state_and_controls(self):
|
||||
@@ -2483,7 +2489,7 @@ class AddonTest(unittest.TestCase):
|
||||
for index, value in enumerate(command)
|
||||
if value == "--allow-pattern"
|
||||
]
|
||||
self.assertEqual(patterns, ["config.json", "model.safetensors"])
|
||||
self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
|
||||
label_index = command.index("--label")
|
||||
self.assertEqual(command[label_index + 1], "BEN2 Base")
|
||||
|
||||
@@ -2512,15 +2518,12 @@ class AddonTest(unittest.TestCase):
|
||||
self.assertEqual(command[0], sys.executable)
|
||||
self.assertIn("--uv", command)
|
||||
self.assertIn("C:/uv/uv.exe", command)
|
||||
self.assertEqual(
|
||||
command[command.index("--torch-variant") + 1],
|
||||
"cu128",
|
||||
)
|
||||
self.assertNotIn("--torch-variant", command)
|
||||
self.assertTrue(
|
||||
any("install_env.py" in str(value) for value in command)
|
||||
)
|
||||
|
||||
def test_install_env_command_auto_selects_torch_variant(self):
|
||||
def test_install_env_command_does_not_select_legacy_runtime(self):
|
||||
client = importlib.import_module("snapform.client")
|
||||
with patch.object(
|
||||
client,
|
||||
@@ -2542,10 +2545,7 @@ class AddonTest(unittest.TestCase):
|
||||
command = client.install_env_command(Path("status.json"))
|
||||
|
||||
self.assertNotIn("--pypi-index", command)
|
||||
self.assertEqual(
|
||||
command[command.index("--torch-variant") + 1],
|
||||
"cpu",
|
||||
)
|
||||
self.assertNotIn("--torch-variant", command)
|
||||
self.assertIn("C:/uv/uv.exe", command)
|
||||
|
||||
def test_completed_export_opens_job_directory(self):
|
||||
|
||||
53
tests/test_ben2_onnx.py
Normal file
53
tests/test_ben2_onnx.py
Normal file
@@ -0,0 +1,53 @@
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform"
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
import ben2_onnx
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def get_inputs(self):
|
||||
return [SimpleNamespace(name="input")]
|
||||
|
||||
def get_outputs(self):
|
||||
return [SimpleNamespace(name="alpha")]
|
||||
|
||||
def run(self, output_names, feeds):
|
||||
self.output_names = output_names
|
||||
self.feeds = feeds
|
||||
gradient = np.linspace(0.0, 1.0, 1024, dtype=np.float32)
|
||||
return [np.tile(gradient, (1024, 1))[None, None]]
|
||||
|
||||
|
||||
class Ben2OnnxTest(unittest.TestCase):
|
||||
def test_preprocess_matches_model_layout(self):
|
||||
image = Image.new("RGB", (40, 20), (64, 128, 192))
|
||||
result = ben2_onnx.preprocess(image)
|
||||
|
||||
self.assertEqual(result.shape, (1, 3, 1024, 1024))
|
||||
self.assertEqual(result.dtype, np.float32)
|
||||
self.assertTrue(result.flags.c_contiguous)
|
||||
|
||||
def test_infer_restores_original_size_and_alpha(self):
|
||||
session = FakeSession()
|
||||
image = Image.new("RGB", (40, 20), (10, 20, 30))
|
||||
|
||||
result = ben2_onnx.infer(session, image)
|
||||
|
||||
self.assertEqual(result.mode, "RGBA")
|
||||
self.assertEqual(result.size, image.size)
|
||||
self.assertEqual(session.feeds["input"].shape, (1, 3, 1024, 1024))
|
||||
alpha = np.asarray(result.getchannel("A"))
|
||||
self.assertLess(int(alpha[:, 0].max()), int(alpha[:, -1].min()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
61
tests/test_da3_onnx.py
Normal file
61
tests/test_da3_onnx.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform"
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
import da3_onnx
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def get_inputs(self):
|
||||
return [SimpleNamespace(name="images"), SimpleNamespace(name="strategy")]
|
||||
|
||||
def run(self, output_names, feeds):
|
||||
self.feeds = feeds
|
||||
_, views, _, height, width = feeds["images"].shape
|
||||
depth = np.ones((1, views, height, width), dtype=np.float32)
|
||||
confidence = np.full_like(depth, 0.75)
|
||||
pose = np.zeros((1, views, 9), dtype=np.float32)
|
||||
pose[..., 6] = 1.0
|
||||
pose[..., 7:] = np.pi / 2.0
|
||||
return depth, confidence, pose
|
||||
|
||||
|
||||
class Da3OnnxTest(unittest.TestCase):
|
||||
def test_target_size_uses_patch_multiple(self):
|
||||
self.assertEqual(
|
||||
da3_onnx.target_size((800, 600), 504, "upper_bound_resize"),
|
||||
(504, 378),
|
||||
)
|
||||
|
||||
def test_infer_returns_camera_and_image_arrays(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
image_path = Path(directory) / "input.png"
|
||||
Image.new("RGB", (80, 60), (10, 20, 30)).save(image_path)
|
||||
session = FakeSession()
|
||||
|
||||
result = da3_onnx.infer(
|
||||
session,
|
||||
[image_path],
|
||||
process_res=56,
|
||||
ref_view_strategy="middle",
|
||||
)
|
||||
|
||||
self.assertEqual(result.depth.shape, (1, 42, 56))
|
||||
self.assertEqual(result.conf.shape, result.depth.shape)
|
||||
self.assertEqual(result.processed_images.shape, (1, 42, 56, 3))
|
||||
self.assertEqual(result.extrinsics.shape, (1, 3, 4))
|
||||
self.assertEqual(result.intrinsics.shape, (1, 3, 3))
|
||||
self.assertEqual(int(session.feeds["strategy"][0]), 1)
|
||||
np.testing.assert_allclose(result.extrinsics[0, :3, :3], np.eye(3))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -36,7 +36,7 @@ class InstallEnvTest(unittest.TestCase):
|
||||
(vendor / "moge.tar.gz").write_bytes(b"moge")
|
||||
return vendor
|
||||
|
||||
def test_installs_da3_from_archives_without_gsplat(self):
|
||||
def test_installs_onnx_runtime_without_training_frameworks(self):
|
||||
bootstrap = load_install_env()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
@@ -71,77 +71,31 @@ class InstallEnvTest(unittest.TestCase):
|
||||
commands += [
|
||||
call.args[4] for call in mocked_run_install.call_args_list
|
||||
]
|
||||
da3_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any("depth-anything-3 @" in str(value) for value in command)
|
||||
)
|
||||
self.assertIn(
|
||||
f"addict=={bootstrap.ADDICT_VERSION}",
|
||||
da3_command,
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
"depth-anything-3 @ file:" in str(value)
|
||||
and "snapform_vendor_" in str(value)
|
||||
for value in da3_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("git+" in str(value) for value in da3_command)
|
||||
)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"gsplat" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"--no-build-isolation" in str(value)
|
||||
for command in commands
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
ben2_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any("ben2 @" in str(value) for value in command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
"ben2 @ file:" in str(value)
|
||||
and "snapform_vendor_" in str(value)
|
||||
for value in ben2_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("git+" in str(value) for value in ben2_command)
|
||||
)
|
||||
service_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
|
||||
)
|
||||
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command)
|
||||
self.assertIn("setuptools", service_command)
|
||||
self.assertNotIn("ninja", service_command)
|
||||
self.assertIn(
|
||||
f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}",
|
||||
service_command,
|
||||
)
|
||||
for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"):
|
||||
self.assertFalse(
|
||||
any(forbidden in str(value) for value in service_command),
|
||||
forbidden,
|
||||
)
|
||||
check_command = bootstrap.verify_command(
|
||||
bootstrap.venv_python(venv_dir)
|
||||
)
|
||||
self.assertTrue(
|
||||
any("import torch;" in str(value) for value in check_command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
"from ben2 import BEN_Base" in str(value)
|
||||
for value in check_command
|
||||
)
|
||||
)
|
||||
self.assertFalse(any("torch" in str(value) for value in check_command))
|
||||
self.assertTrue(
|
||||
any("import fastapi, uvicorn" in str(value) for value in check_command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any("onnxruntime" in str(value) for value in check_command)
|
||||
)
|
||||
self.assertFalse(
|
||||
any("import gsplat" in str(value) for value in check_command)
|
||||
)
|
||||
@@ -157,24 +111,24 @@ class InstallEnvTest(unittest.TestCase):
|
||||
if any(
|
||||
repository in command
|
||||
for repository in (
|
||||
"Ruicheng/moge-2-vitb-normal",
|
||||
"PramaLLC/BEN2",
|
||||
"Ruicheng/moge-2-vits-normal-onnx",
|
||||
"onnx-community/BEN2-ONNX",
|
||||
)
|
||||
)
|
||||
]
|
||||
self.assertEqual(len(default_downloads), 2)
|
||||
self.assertTrue(any("Ruicheng/moge-2-vitb-normal" in map(str, command) for command in default_downloads))
|
||||
self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads))
|
||||
ben2_download = next(
|
||||
command
|
||||
for command in default_downloads
|
||||
if "PramaLLC/BEN2" in command
|
||||
if "onnx-community/BEN2-ONNX" in command
|
||||
)
|
||||
patterns = [
|
||||
ben2_download[index + 1]
|
||||
for index, value in enumerate(ben2_download)
|
||||
if value == "--allow-pattern"
|
||||
]
|
||||
self.assertEqual(patterns, ["config.json", "model.safetensors"])
|
||||
self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
|
||||
self.assertNotIn("--stage", ben2_download)
|
||||
self.assertNotIn("--stage-label", ben2_download)
|
||||
self.assertEqual(
|
||||
@@ -182,7 +136,7 @@ class InstallEnvTest(unittest.TestCase):
|
||||
"BEN2 Base",
|
||||
)
|
||||
|
||||
def test_darwin_install_skips_cuda_index_and_xformers(self):
|
||||
def test_darwin_install_uses_default_onnx_runtime(self):
|
||||
bootstrap = load_install_env()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
@@ -224,15 +178,15 @@ class InstallEnvTest(unittest.TestCase):
|
||||
commands = [
|
||||
call.args[4] for call in mocked_run_install.call_args_list
|
||||
]
|
||||
torch_command = next(
|
||||
runtime_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any(
|
||||
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
|
||||
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value)
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
self.assertNotIn("--index-url", torch_command)
|
||||
self.assertNotIn("--index-url", runtime_command)
|
||||
self.assertFalse(
|
||||
any(
|
||||
"xformers" in str(value)
|
||||
@@ -241,30 +195,27 @@ class InstallEnvTest(unittest.TestCase):
|
||||
)
|
||||
)
|
||||
|
||||
def test_intel_mac_is_rejected_with_clear_error(self):
|
||||
def test_intel_mac_uses_portable_onnx_runtime(self):
|
||||
bootstrap = load_install_env()
|
||||
args = SimpleNamespace(
|
||||
uv="uv",
|
||||
venv_dir=".venv",
|
||||
models_dir="models",
|
||||
torch_variant="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status="status.json",
|
||||
)
|
||||
with patch.object(
|
||||
bootstrap.sys,
|
||||
"platform",
|
||||
"darwin",
|
||||
), patch.object(
|
||||
bootstrap.platform,
|
||||
"machine",
|
||||
return_value="x86_64",
|
||||
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
|
||||
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
|
||||
bootstrap.install(args)
|
||||
|
||||
mocked_rmtree.assert_not_called()
|
||||
):
|
||||
stages = bootstrap.install_stages(
|
||||
SimpleNamespace(
|
||||
models_dir="models",
|
||||
status="status.json",
|
||||
),
|
||||
".venv",
|
||||
"python",
|
||||
"uv",
|
||||
)
|
||||
runtime = stages[1][1]
|
||||
self.assertIn(
|
||||
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}",
|
||||
runtime,
|
||||
)
|
||||
|
||||
def test_run_stage_writes_stage_progress(self):
|
||||
bootstrap = load_install_env()
|
||||
@@ -414,10 +365,10 @@ class InstallEnvTest(unittest.TestCase):
|
||||
)
|
||||
self.assertEqual(ends[-1], 1.0)
|
||||
self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0])
|
||||
runtime_ranges = stage_ranges[:4]
|
||||
runtime_ranges = stage_ranges[:2]
|
||||
self.assertEqual(
|
||||
[item[0] for item in runtime_ranges],
|
||||
[1, 2, 3, 3],
|
||||
[1, 2],
|
||||
)
|
||||
self.assertTrue(
|
||||
all(item[0] is None for item in stage_ranges[4:])
|
||||
@@ -606,19 +557,19 @@ class InstallEnvTest(unittest.TestCase):
|
||||
self.assertEqual(bootstrap.format_elapsed(12), "12s")
|
||||
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
|
||||
|
||||
def test_ai_component_estimate_continues_across_internal_steps(self):
|
||||
def test_onnx_runtime_estimate_decreases_with_progress(self):
|
||||
bootstrap = load_install_env()
|
||||
|
||||
first = bootstrap.estimated_remaining(
|
||||
"AI components",
|
||||
bootstrap.XFORMERS_SLICE[0],
|
||||
)
|
||||
second = bootstrap.estimated_remaining(
|
||||
"AI components",
|
||||
"ONNX Runtime",
|
||||
bootstrap.AI_COMPONENTS_SLICE[0],
|
||||
)
|
||||
second = bootstrap.estimated_remaining(
|
||||
"ONNX Runtime",
|
||||
sum(bootstrap.AI_COMPONENTS_SLICE) / 2,
|
||||
)
|
||||
|
||||
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["AI components"])
|
||||
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["ONNX Runtime"])
|
||||
self.assertGreater(second, 0)
|
||||
self.assertLess(second, first)
|
||||
|
||||
@@ -727,7 +678,7 @@ class InstallEnvTest(unittest.TestCase):
|
||||
ready["packages"],
|
||||
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
||||
)
|
||||
self.assertEqual(ready["devices"], ["cuda", "cpu"])
|
||||
self.assertEqual(ready["devices"], ["directml", "cpu"])
|
||||
self.assertEqual(
|
||||
ready["manifest_revision"],
|
||||
bootstrap.MANIFEST_REVISION,
|
||||
@@ -814,7 +765,7 @@ class InstallEnvTest(unittest.TestCase):
|
||||
|
||||
def fake_run(command, environment=None):
|
||||
commands.append(command)
|
||||
if any("moge-2-vitb-normal" in str(value) for value in command):
|
||||
if any("moge-2-vits-normal" in str(value) for value in command):
|
||||
raise subprocess.CalledProcessError(1, command)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
@@ -852,7 +803,7 @@ class InstallEnvTest(unittest.TestCase):
|
||||
self.assertIn("Environment is installed", status["message"])
|
||||
self.assertTrue(
|
||||
any(
|
||||
"PramaLLC/BEN2" in map(str, command)
|
||||
"onnx-community/BEN2-ONNX" in map(str, command)
|
||||
for command in commands
|
||||
)
|
||||
)
|
||||
@@ -983,25 +934,13 @@ class InstallEnvTest(unittest.TestCase):
|
||||
for value in command
|
||||
)
|
||||
)
|
||||
ai_command = next(
|
||||
runtime_command = next(
|
||||
command
|
||||
for command in install_commands
|
||||
if any("depth-anything-3 @" in str(value) for value in command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
"depth-anything-3 @ file:" in str(value)
|
||||
and "snapform_vendor_" in str(value)
|
||||
for value in ai_command
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
any(
|
||||
"ben2 @ file:" in str(value)
|
||||
and "snapform_vendor_" in str(value)
|
||||
for value in ai_command
|
||||
)
|
||||
if any("onnxruntime" in str(value) for value in command)
|
||||
)
|
||||
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
|
||||
self.assertFalse(any("torch" in str(value) for value in runtime_command))
|
||||
|
||||
def test_install_reuses_existing_uv(self):
|
||||
bootstrap = load_install_env()
|
||||
@@ -1047,81 +986,12 @@ class InstallEnvTest(unittest.TestCase):
|
||||
|
||||
self.assertEqual(uv_path.read_bytes(), b"binary")
|
||||
|
||||
def test_copy_vendor_inputs_copies_missing_archives(self):
|
||||
bootstrap = load_install_env()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
bundled = root / "bundled"
|
||||
bundled.mkdir()
|
||||
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
|
||||
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
|
||||
(bundled / "moge.tar.gz").write_bytes(b"moge")
|
||||
vendor = root / "vendor"
|
||||
|
||||
bootstrap.copy_vendor_inputs(bundled, vendor)
|
||||
|
||||
self.assertEqual(
|
||||
(vendor / "depth-anything-3.tar.gz").read_bytes(),
|
||||
b"da3",
|
||||
)
|
||||
self.assertEqual(
|
||||
(vendor / "ben2.tar.gz").read_bytes(),
|
||||
b"ben2",
|
||||
)
|
||||
|
||||
def test_copy_vendor_inputs_replaces_same_size_stale_archive(self):
|
||||
bootstrap = load_install_env()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
bundled = root / "bundled"
|
||||
bundled.mkdir()
|
||||
(bundled / "depth-anything-3.tar.gz").write_bytes(b"new")
|
||||
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
|
||||
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
|
||||
(bundled / "moge.tar.gz").write_bytes(b"moge")
|
||||
vendor = root / "vendor"
|
||||
vendor.mkdir()
|
||||
(vendor / "depth-anything-3.tar.gz").write_bytes(b"old")
|
||||
|
||||
bootstrap.copy_vendor_inputs(bundled, vendor)
|
||||
|
||||
self.assertEqual(
|
||||
(vendor / "depth-anything-3.tar.gz").read_bytes(),
|
||||
b"new",
|
||||
)
|
||||
|
||||
def test_remove_legacy_vendor_removes_only_install_inputs(self):
|
||||
bootstrap = load_install_env()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
vendor = root / "vendor"
|
||||
vendor.mkdir()
|
||||
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||
(root / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||
(root / "ben2.tar.gz").write_bytes(b"ben2")
|
||||
models = root / "models"
|
||||
models.mkdir()
|
||||
|
||||
bootstrap.remove_legacy_vendor(root)
|
||||
|
||||
self.assertFalse(vendor.exists())
|
||||
self.assertFalse((root / "depth-anything-3.tar.gz").exists())
|
||||
self.assertFalse((root / "ben2.tar.gz").exists())
|
||||
self.assertTrue(models.is_dir())
|
||||
|
||||
def test_install_stages_uses_local_source_archives(self):
|
||||
def test_install_stages_use_only_runtime_packages(self):
|
||||
bootstrap = load_install_env()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
venv_dir = root / ".venv"
|
||||
venv_dir.mkdir()
|
||||
vendor = root / "vendor"
|
||||
vendor.mkdir()
|
||||
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
|
||||
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
|
||||
(vendor / "moge.tar.gz").write_bytes(b"moge")
|
||||
uv_path = root / "uv"
|
||||
uv_path.write_bytes(b"uv")
|
||||
args = SimpleNamespace(
|
||||
@@ -1141,16 +1011,14 @@ class InstallEnvTest(unittest.TestCase):
|
||||
uv_path,
|
||||
)
|
||||
|
||||
ai_command = next(
|
||||
runtime_command = next(
|
||||
command
|
||||
for _label, command, _environment, kind in stages
|
||||
if kind == "install"
|
||||
and any("depth-anything-3" in str(value) for value in command)
|
||||
and any("onnxruntime" in str(value) for value in command)
|
||||
)
|
||||
da3_local = (vendor / "depth-anything-3.tar.gz").as_uri()
|
||||
ben2_local = (vendor / "ben2.tar.gz").as_uri()
|
||||
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
|
||||
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
|
||||
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
|
||||
self.assertFalse(any("torch" in str(value) for value in runtime_command))
|
||||
|
||||
def test_install_stages_download_only_default_models(self):
|
||||
bootstrap = load_install_env()
|
||||
@@ -1189,13 +1057,13 @@ class InstallEnvTest(unittest.TestCase):
|
||||
if kind == "download"
|
||||
}
|
||||
self.assertIn(
|
||||
"moge2_vitb_normal",
|
||||
" ".join(map(str, downloads["MoGe-2 ViT-B Normal"])),
|
||||
"moge-2-vits-normal-onnx",
|
||||
" ".join(map(str, downloads["MoGe-2 ViT-S Normal"])),
|
||||
)
|
||||
self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"])))
|
||||
self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"])))
|
||||
self.assertEqual(
|
||||
set(downloads),
|
||||
{"MoGe-2 ViT-B Normal", "BEN2 Base"},
|
||||
{"MoGe-2 ViT-S Normal", "BEN2 Base"},
|
||||
)
|
||||
|
||||
def test_install_does_not_verify_env(self):
|
||||
|
||||
@@ -30,12 +30,6 @@ class ModelTest(unittest.TestCase):
|
||||
self.assertEqual(model.repository, "depth-anything/DA3-BASE")
|
||||
self.assertEqual(model.license_name, "Apache 2.0")
|
||||
|
||||
def test_large_model_exposes_noncommercial_license(self):
|
||||
model = self.models.get_model("LARGE_1_1")
|
||||
|
||||
self.assertEqual(model.license_name, "CC BY-NC 4.0")
|
||||
self.assertFalse(model.enabled)
|
||||
|
||||
def test_giant_models_expose_gaussian_head(self):
|
||||
for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"):
|
||||
with self.subTest(model_key=model_key):
|
||||
@@ -54,19 +48,20 @@ class ModelTest(unittest.TestCase):
|
||||
def test_ben2_download_only_includes_required_hugging_face_files(self):
|
||||
model = self.models.get_downloadable_model("BEN2_BASE")
|
||||
|
||||
self.assertEqual(model.repository, "PramaLLC/BEN2")
|
||||
self.assertEqual(model.license_name, "MIT")
|
||||
self.assertEqual(
|
||||
model.download_patterns,
|
||||
("config.json", "model.safetensors"),
|
||||
)
|
||||
self.assertEqual(model.repository, "onnx-community/BEN2-ONNX")
|
||||
self.assertEqual(model.license_name, "Apache 2.0")
|
||||
self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",))
|
||||
self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",))
|
||||
self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx")
|
||||
self.assertNotIn("BEN2_BASE", self.models.MODELS)
|
||||
|
||||
def test_da3_model_directories_share_the_da3_prefix(self):
|
||||
def test_da3_model_directories_use_repository_names(self):
|
||||
for model_key, model in self.models.MODELS.items():
|
||||
if not model.enabled:
|
||||
continue
|
||||
with self.subTest(model_key=model_key):
|
||||
self.assertTrue(
|
||||
model.directory_name.startswith("da3_"),
|
||||
model.directory_name.startswith("DA3-"),
|
||||
model.directory_name,
|
||||
)
|
||||
self.assertEqual(
|
||||
@@ -74,19 +69,36 @@ class ModelTest(unittest.TestCase):
|
||||
model.directory_name,
|
||||
)
|
||||
|
||||
def test_ben2_model_directory_uses_ben2_prefix(self):
|
||||
def test_ben2_model_directory_uses_repository_name(self):
|
||||
self.assertEqual(
|
||||
self.models.model_directory_name("BEN2_BASE"),
|
||||
"ben2_base",
|
||||
"BEN2-ONNX",
|
||||
)
|
||||
|
||||
def test_moge2_uses_single_checkpoint_and_mit_license(self):
|
||||
model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL")
|
||||
def test_enabled_da3_models_use_onnx_and_r2_checksums(self):
|
||||
for model in self.models.MODELS.values():
|
||||
if not model.enabled:
|
||||
continue
|
||||
with self.subTest(model=model.key):
|
||||
self.assertTrue(model.download_patterns[0].endswith(".onnx"))
|
||||
self.assertEqual(model.ready_patterns, ("model.onnx",))
|
||||
self.assertEqual(model.r2_files[0][0], "model.onnx")
|
||||
|
||||
self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal")
|
||||
def test_moge2_models_use_official_onnx_models(self):
|
||||
for key, variant in (("MOGE2_VITS_NORMAL", "vits"),
|
||||
("MOGE2_VITB_NORMAL", "vitb"),
|
||||
("MOGE2_VITL_NORMAL", "vitl")):
|
||||
with self.subTest(key=key):
|
||||
model = self.models.get_downloadable_model(key)
|
||||
self.assertEqual(
|
||||
model.repository,
|
||||
f"Ruicheng/moge-2-{variant}-normal-onnx",
|
||||
)
|
||||
self.assertEqual(model.license_name, "MIT")
|
||||
self.assertEqual(model.download_patterns, ("model.pt",))
|
||||
self.assertEqual(model.ready_patterns, ("model.pt",))
|
||||
self.assertEqual(model.download_patterns, ("model.onnx",))
|
||||
self.assertEqual(model.ready_patterns, ("model.onnx",))
|
||||
self.assertEqual(model.r2_directory, model.repository.split("/")[-1])
|
||||
self.assertEqual(model.r2_files[0][0], "model.onnx")
|
||||
self.assertFalse(model.supports_gs)
|
||||
|
||||
def test_moge2_models_match_uploaded_r2_variants(self):
|
||||
@@ -96,5 +108,9 @@ class ModelTest(unittest.TestCase):
|
||||
)
|
||||
self.assertEqual(
|
||||
[model.r2_directory for model in self.models.MOGE2_MODELS.values()],
|
||||
["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"],
|
||||
[
|
||||
"moge-2-vits-normal-onnx",
|
||||
"moge-2-vitb-normal-onnx",
|
||||
"moge-2-vitl-normal-onnx",
|
||||
],
|
||||
)
|
||||
|
||||
130
tests/test_moge_onnx.py
Normal file
130
tests/test_moge_onnx.py
Normal file
@@ -0,0 +1,130 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_module():
|
||||
specification = importlib.util.spec_from_file_location(
|
||||
"snapform_moge_onnx",
|
||||
PROJECT_ROOT / "src" / "snapform" / "moge_onnx.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(specification)
|
||||
specification.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class FakeInput:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, outputs):
|
||||
self.outputs = outputs
|
||||
self.calls = []
|
||||
|
||||
def run(self, names, inputs):
|
||||
self.calls.append((names, inputs))
|
||||
return [self.outputs[name] for name in names]
|
||||
|
||||
def get_outputs(self):
|
||||
return [FakeInput(name) for name in self.outputs]
|
||||
|
||||
|
||||
class MogeOnnxTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.moge = load_module()
|
||||
|
||||
def test_auto_provider_prefers_accelerated_runtime(self):
|
||||
providers = self.moge.select_providers(
|
||||
"auto",
|
||||
["CPUExecutionProvider", "DmlExecutionProvider"],
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
providers,
|
||||
["DmlExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
|
||||
def test_explicit_unavailable_provider_falls_back_to_cpu(self):
|
||||
providers = self.moge.select_providers(
|
||||
"cuda",
|
||||
["CPUExecutionProvider"],
|
||||
)
|
||||
|
||||
self.assertEqual(providers, ["CPUExecutionProvider"])
|
||||
|
||||
def test_infer_uses_dynamic_token_input_and_numpy_postprocess(self):
|
||||
height, width = 8, 12
|
||||
raw = {
|
||||
"points": np.zeros((1, height, width, 3), dtype=np.float32),
|
||||
"normal": np.zeros((1, height, width, 3), dtype=np.float32),
|
||||
"mask": np.ones((1, height, width), dtype=np.float32),
|
||||
"scale": np.ones((1,), dtype=np.float32),
|
||||
}
|
||||
raw["points"][..., 2] = 1.0
|
||||
raw["normal"][..., 2] = 1.0
|
||||
session = FakeSession(raw)
|
||||
expected = {
|
||||
"points": np.zeros((height, width, 3), dtype=np.float32),
|
||||
"depth": np.ones((height, width), dtype=np.float32),
|
||||
"normal": raw["normal"][0],
|
||||
"mask": np.ones((height, width), dtype=bool),
|
||||
"intrinsics": np.eye(3, dtype=np.float32),
|
||||
}
|
||||
|
||||
with patch.object(self.moge, "postprocess", return_value=expected):
|
||||
result = self.moge.infer(
|
||||
session,
|
||||
np.zeros((height, width, 3), dtype=np.float32),
|
||||
resolution_level=9,
|
||||
)
|
||||
|
||||
self.assertIs(result, expected)
|
||||
names, inputs = session.calls[0]
|
||||
self.assertEqual(names, ["points", "normal", "mask", "scale"])
|
||||
self.assertEqual(inputs["image"].shape, (1, 3, height, width))
|
||||
self.assertEqual(inputs["num_tokens"].dtype, np.int64)
|
||||
self.assertEqual(inputs["num_tokens"], 3600)
|
||||
|
||||
def test_postprocess_recovers_camera_projection_and_metric_scale(self):
|
||||
height, width = 64, 80
|
||||
focal = np.float32(0.9)
|
||||
shift = np.float32(0.3)
|
||||
metric_scale = np.float32(2.0)
|
||||
uv = self.moge.normalized_view_plane_uv(width, height, np.float32)
|
||||
depth = np.linspace(1.0, 2.4, height * width, dtype=np.float32).reshape(
|
||||
height,
|
||||
width,
|
||||
)
|
||||
raw_points = np.concatenate(
|
||||
(uv * depth[..., None] / focal, (depth - shift)[..., None]),
|
||||
axis=-1,
|
||||
)
|
||||
raw = {
|
||||
"points": raw_points[None],
|
||||
"normal": np.broadcast_to(
|
||||
np.array((0.0, 0.0, 1.0), dtype=np.float32),
|
||||
(1, height, width, 3),
|
||||
),
|
||||
"mask": np.ones((1, height, width), dtype=np.float32),
|
||||
"metric_scale": np.array((metric_scale,), dtype=np.float32),
|
||||
}
|
||||
|
||||
result = self.moge.postprocess(raw)
|
||||
|
||||
np.testing.assert_allclose(result["depth"], depth * metric_scale, rtol=1e-3)
|
||||
np.testing.assert_allclose(result["points"][..., 2], result["depth"], rtol=1e-5)
|
||||
self.assertTrue(result["mask"].all())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,7 +1,4 @@
|
||||
import ast
|
||||
import importlib.util
|
||||
import io
|
||||
import tarfile
|
||||
import tempfile
|
||||
import tomllib
|
||||
import unittest
|
||||
@@ -64,8 +61,6 @@ class PackTest(unittest.TestCase):
|
||||
self.assertIn("service.py", names)
|
||||
self.assertNotIn("snapform/__init__.py", names)
|
||||
self.assertNotIn("importer.py", names)
|
||||
self.assertIn("vendor/depth-anything-3.tar.gz", names)
|
||||
self.assertIn("vendor/ben2.tar.gz", names)
|
||||
self.assertIn(
|
||||
f"vendor/{pack.PLATFORMS[platform]}",
|
||||
names,
|
||||
@@ -108,95 +103,5 @@ class PackTest(unittest.TestCase):
|
||||
{"windows-x64", "macos-arm64", "linux-x64"},
|
||||
)
|
||||
|
||||
def test_bundled_depth_anything_3_has_no_xformers_dependency(self):
|
||||
pack = load_pack()
|
||||
with tarfile.open(
|
||||
fileobj=io.BytesIO(
|
||||
pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
|
||||
),
|
||||
mode="r:gz",
|
||||
) as archive:
|
||||
raw_pyproject = archive.extractfile(
|
||||
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||
"pyproject.toml"
|
||||
).read()
|
||||
pyproject = raw_pyproject.decode("utf-8")
|
||||
self.assertNotIn("xformers", pyproject)
|
||||
self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
|
||||
|
||||
def test_bundled_depth_anything_3_exports_are_lazy(self):
|
||||
pack = load_pack()
|
||||
with tarfile.open(
|
||||
fileobj=io.BytesIO(
|
||||
pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
|
||||
),
|
||||
mode="r:gz",
|
||||
) as archive:
|
||||
source = archive.extractfile(
|
||||
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||
"src/depth_anything_3/utils/export/__init__.py"
|
||||
).read().decode("utf-8")
|
||||
tree = ast.parse(source)
|
||||
module_imports = [
|
||||
ast.get_source_segment(source, node)
|
||||
for node in tree.body
|
||||
if isinstance(node, (ast.Import, ast.ImportFrom))
|
||||
]
|
||||
self.assertEqual(
|
||||
module_imports,
|
||||
["from depth_anything_3.specs import Prediction"],
|
||||
)
|
||||
|
||||
def test_official_depth_anything_3_baseline_is_pristine(self):
|
||||
pack = load_pack()
|
||||
archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz"
|
||||
with tarfile.open(archive_path, "r:gz") as archive:
|
||||
raw_pyproject = archive.extractfile(
|
||||
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||
"pyproject.toml"
|
||||
).read()
|
||||
export_source = archive.extractfile(
|
||||
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||
"src/depth_anything_3/utils/export/__init__.py"
|
||||
).read().decode("utf-8")
|
||||
pyproject = raw_pyproject.decode("utf-8")
|
||||
self.assertIn("xformers", pyproject)
|
||||
self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
|
||||
self.assertIn(
|
||||
"from depth_anything_3.utils.export.gs import",
|
||||
export_source,
|
||||
)
|
||||
|
||||
def test_bundled_depth_anything_3_qr_falls_back_on_mps(self):
|
||||
pack = load_pack()
|
||||
with tarfile.open(
|
||||
fileobj=io.BytesIO(
|
||||
pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
|
||||
),
|
||||
mode="r:gz",
|
||||
) as archive:
|
||||
source = archive.extractfile(
|
||||
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||
"src/depth_anything_3/utils/ray_utils.py"
|
||||
).read().decode("utf-8")
|
||||
self.assertIn('use_cpu = device.type == "mps"', source)
|
||||
self.assertIn("A_tilde = A_tilde.cpu()", source)
|
||||
self.assertIn("Q_tilde = Q_tilde.to(device=device)", source)
|
||||
|
||||
def test_official_depth_anything_3_ray_utils_is_pristine(self):
|
||||
pack = load_pack()
|
||||
archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz"
|
||||
with tarfile.open(archive_path, "r:gz") as archive:
|
||||
source = archive.extractfile(
|
||||
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
|
||||
"src/depth_anything_3/utils/ray_utils.py"
|
||||
).read().decode("utf-8")
|
||||
self.assertIn(
|
||||
"Q_tilde, R_tilde = torch.linalg.qr(A_tilde)",
|
||||
source,
|
||||
)
|
||||
self.assertNotIn("use_cpu", source)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -40,7 +40,6 @@ class ServiceTest(unittest.TestCase):
|
||||
|
||||
def create_cache(self):
|
||||
cache = self.service.ModelCache.__new__(self.service.ModelCache)
|
||||
cache.torch = SimpleNamespace(cuda=FakeCuda())
|
||||
cache.ben2_model = None
|
||||
cache.ben2_key = None
|
||||
cache.da3_model = None
|
||||
@@ -52,65 +51,41 @@ class ServiceTest(unittest.TestCase):
|
||||
def test_ben2_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
loads.append(model_dir)
|
||||
return cls()
|
||||
|
||||
def to(self, device):
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
ben2 = ModuleType("ben2")
|
||||
ben2.BEN_Base = FakeModel
|
||||
ben2 = ModuleType("ben2_onnx")
|
||||
ben2.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
cache = self.create_cache()
|
||||
device = SimpleNamespace(type="cuda")
|
||||
with patch.dict(sys.modules, {"ben2": ben2}):
|
||||
with patch.dict(sys.modules, {"ben2_onnx": ben2}):
|
||||
first, _load_ms = cache.get_ben2("model", device)
|
||||
second, cached_load_ms = cache.get_ben2("model", device)
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, ["model"])
|
||||
self.assertEqual(loads, [("model", "cuda")])
|
||||
self.assertEqual(cache.status()["ben2"], "model")
|
||||
self.assertEqual(cached_load_ms, 0.0)
|
||||
|
||||
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
loads.append(model_dir)
|
||||
return cls()
|
||||
|
||||
def to(self, device):
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
package = ModuleType("depth_anything_3")
|
||||
api = ModuleType("depth_anything_3.api")
|
||||
api.DepthAnything3 = FakeModel
|
||||
da3 = ModuleType("da3_onnx")
|
||||
da3.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
cache = self.create_cache()
|
||||
device = SimpleNamespace(type="cuda")
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"depth_anything_3": package,
|
||||
"depth_anything_3.api": api,
|
||||
"da3_onnx": da3,
|
||||
},
|
||||
):
|
||||
first = cache.get_da3("model", device)
|
||||
second = cache.get_da3("model", device)
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, ["model"])
|
||||
self.assertEqual(loads, [("model", "cuda")])
|
||||
self.assertEqual(cache.status()["da3"], "model")
|
||||
|
||||
def test_depth_model_cache_keeps_only_da3_or_moge(self):
|
||||
@@ -135,13 +110,10 @@ class ServiceTest(unittest.TestCase):
|
||||
|
||||
moge = object()
|
||||
|
||||
def load_moge(_model_dir, _device, key):
|
||||
cache.moge_model = moge
|
||||
cache.moge_key = key
|
||||
return moge
|
||||
|
||||
with patch.object(cache, "_load_moge", side_effect=load_moge):
|
||||
self.assertIs(cache.get_moge("moge", device), moge)
|
||||
fake_moge = ModuleType("moge_onnx")
|
||||
fake_moge.create_session = lambda _model_dir, _device: moge
|
||||
with patch.dict(sys.modules, {"moge_onnx": fake_moge}):
|
||||
self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge)
|
||||
|
||||
self.assertIsNone(cache.da3_model)
|
||||
self.assertIsNotNone(cache.ben2_model)
|
||||
|
||||
@@ -66,57 +66,6 @@ class WorkerTest(unittest.TestCase):
|
||||
|
||||
self.worker.check_cancelled(lambda: False)
|
||||
|
||||
def test_select_device_supports_mps_and_auto_fallback(self):
|
||||
mps_torch = SimpleNamespace(
|
||||
cuda=SimpleNamespace(is_available=lambda: False),
|
||||
backends=SimpleNamespace(
|
||||
mps=SimpleNamespace(is_available=lambda: True),
|
||||
),
|
||||
)
|
||||
mps_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
self.worker.select_device("mps", mps_torch).type,
|
||||
"mps",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("auto", mps_torch).type,
|
||||
"mps",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("cpu", mps_torch).type,
|
||||
"cpu",
|
||||
)
|
||||
|
||||
cuda_torch = SimpleNamespace(
|
||||
cuda=SimpleNamespace(is_available=lambda: True),
|
||||
)
|
||||
cuda_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("auto", cuda_torch).type,
|
||||
"cuda",
|
||||
)
|
||||
|
||||
cpu_only_torch = SimpleNamespace(
|
||||
cuda=SimpleNamespace(is_available=lambda: False),
|
||||
backends=SimpleNamespace(
|
||||
mps=SimpleNamespace(is_available=lambda: False),
|
||||
),
|
||||
)
|
||||
cpu_only_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type
|
||||
)
|
||||
self.assertEqual(
|
||||
self.worker.select_device("auto", cpu_only_torch).type,
|
||||
"cpu",
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
|
||||
self.worker.select_device("mps", cpu_only_torch)
|
||||
|
||||
def test_reconstruct_parser_accepts_background_model_directory(self):
|
||||
args = self.worker.parser().parse_args(
|
||||
[
|
||||
@@ -436,12 +385,6 @@ class WorkerTest(unittest.TestCase):
|
||||
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
|
||||
self.worker.prepare_input(input_path, 16, directory)
|
||||
|
||||
def test_moge_runtime_disables_unsupported_xformers_attention(self):
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
self.worker.prepare_moge_runtime()
|
||||
|
||||
self.assertEqual(os.environ["XFORMERS_DISABLED"], "1")
|
||||
|
||||
def test_download_model_uses_actual_request_without_preflight(self):
|
||||
calls = []
|
||||
fake_hugging_face = ModuleType("huggingface_hub")
|
||||
@@ -605,8 +548,11 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("ben2")
|
||||
fake_ben2.BEN_Base = FakeModel
|
||||
fake_ben2 = ModuleType("ben2_onnx")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image: (
|
||||
inference_options.append((image, {})) or FakeImage()
|
||||
)
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
||||
|
||||
@@ -614,7 +560,7 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"ben2": fake_ben2,
|
||||
"ben2_onnx": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
@@ -635,7 +581,7 @@ class WorkerTest(unittest.TestCase):
|
||||
|
||||
self.assertEqual(
|
||||
inference_options[0][1],
|
||||
{"refine_foreground": False},
|
||||
{},
|
||||
)
|
||||
self.assertEqual(
|
||||
save_options[0],
|
||||
@@ -690,8 +636,11 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("ben2")
|
||||
fake_ben2.BEN_Base = FakeModel
|
||||
fake_ben2 = ModuleType("ben2_onnx")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image: (
|
||||
inference_options.append((image, {})) or FakeImage()
|
||||
)
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
||||
|
||||
@@ -699,7 +648,7 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"ben2": fake_ben2,
|
||||
"ben2_onnx": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
@@ -734,7 +683,7 @@ class WorkerTest(unittest.TestCase):
|
||||
|
||||
self.assertEqual(len(inference_options), 2)
|
||||
for _image, options in inference_options:
|
||||
self.assertEqual(options, {"refine_foreground": False})
|
||||
self.assertEqual(options, {})
|
||||
|
||||
def test_remove_background_parser_accepts_video_frames(self):
|
||||
args = self.worker.parser().parse_args(
|
||||
@@ -813,64 +762,25 @@ class WorkerTest(unittest.TestCase):
|
||||
|
||||
def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
|
||||
inference_options = []
|
||||
export_calls = []
|
||||
mesh_calls = []
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
if model_dir != "model":
|
||||
raise AssertionError(f"Unexpected model directory: {model_dir}")
|
||||
return cls()
|
||||
|
||||
def to(self, device):
|
||||
if device.type != "cpu":
|
||||
raise AssertionError(f"Unexpected device: {device.type}")
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
def inference(self, images, **options):
|
||||
inference_options.append((images, options))
|
||||
return SimpleNamespace(
|
||||
prediction = SimpleNamespace(
|
||||
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
|
||||
conf=np.ones((1, 2, 2), dtype=np.float32),
|
||||
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
|
||||
intrinsics=np.eye(3, dtype=np.float32)[None],
|
||||
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
|
||||
)
|
||||
fake_da3 = ModuleType("da3_onnx")
|
||||
fake_da3.create_session = lambda _model_dir, _device: object()
|
||||
|
||||
def fake_export(prediction, export_format, export_dir, **options):
|
||||
export_calls.append((prediction, export_format, export_dir, options))
|
||||
output_directory = Path(export_dir)
|
||||
if export_format == "glb":
|
||||
(output_directory / "scene.glb").touch()
|
||||
elif export_format == "gs_ply":
|
||||
ply_directory = output_directory / "gs_ply"
|
||||
ply_directory.mkdir()
|
||||
(ply_directory / "0000.ply").touch()
|
||||
elif export_format == "depth_vis":
|
||||
depth_directory = output_directory / "depth_vis"
|
||||
depth_directory.mkdir()
|
||||
(depth_directory / "0000.jpg").touch()
|
||||
else:
|
||||
mini_directory = output_directory / "exports" / "mini_npz"
|
||||
mini_directory.mkdir(parents=True)
|
||||
np.savez_compressed(
|
||||
mini_directory / "results.npz",
|
||||
depth=np.zeros((14, 14)),
|
||||
)
|
||||
def fake_infer(_session, images, **options):
|
||||
inference_options.append((images, options))
|
||||
return prediction
|
||||
|
||||
fake_torch = ModuleType("torch")
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type,
|
||||
)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_api = ModuleType("depth_anything_3.api")
|
||||
fake_api.DepthAnything3 = FakeModel
|
||||
fake_export_module = ModuleType("depth_anything_3.utils.export")
|
||||
fake_export_module.export = fake_export
|
||||
fake_utils = ModuleType("depth_anything_3.utils")
|
||||
fake_utils.__path__ = []
|
||||
fake_package = ModuleType("depth_anything_3")
|
||||
fake_package.__path__ = []
|
||||
fake_da3.infer = fake_infer
|
||||
fake_runtime = ModuleType("onnx_runtime")
|
||||
fake_runtime.provider_label = lambda _session: "CPU"
|
||||
|
||||
def fake_mesh_export(
|
||||
prediction,
|
||||
@@ -893,6 +803,20 @@ class WorkerTest(unittest.TestCase):
|
||||
output_path.touch()
|
||||
return output_path
|
||||
|
||||
def fake_point_cloud(_prediction, output_directory, *_args):
|
||||
(Path(output_directory) / "scene.glb").touch()
|
||||
|
||||
def fake_depth_preview(_prediction, output_directory):
|
||||
target = Path(output_directory) / "depth_vis"
|
||||
target.mkdir()
|
||||
(target / "0000.jpg").touch()
|
||||
|
||||
def fake_mini_npz(_prediction, output_directory):
|
||||
np.savez_compressed(
|
||||
Path(output_directory) / "results.npz",
|
||||
depth=np.zeros((1, 2, 2)),
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
input_path = root / "input.png"
|
||||
@@ -900,20 +824,28 @@ class WorkerTest(unittest.TestCase):
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"depth_anything_3": fake_package,
|
||||
"depth_anything_3.api": fake_api,
|
||||
"depth_anything_3.utils": fake_utils,
|
||||
"depth_anything_3.utils.export": fake_export_module,
|
||||
"da3_onnx": fake_da3,
|
||||
"onnx_runtime": fake_runtime,
|
||||
},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"export_textured_mesh_glb",
|
||||
side_effect=fake_mesh_export,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"export_da3_point_cloud",
|
||||
side_effect=fake_point_cloud,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"export_da3_depth_preview",
|
||||
side_effect=fake_depth_preview,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"export_da3_mini_npz",
|
||||
side_effect=fake_mini_npz,
|
||||
):
|
||||
debug_formats = (
|
||||
"glb",
|
||||
"gs_ply",
|
||||
"depth_vis",
|
||||
"mini_npz",
|
||||
"mesh",
|
||||
@@ -945,15 +877,11 @@ class WorkerTest(unittest.TestCase):
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertIn(export_format, status["message"])
|
||||
result_directory = Path(args.output_dir)
|
||||
if export_format == "gs_ply":
|
||||
self.assertTrue((result_directory / "0000.ply").is_file())
|
||||
self.assertFalse((result_directory / "gs_ply").exists())
|
||||
if export_format == "depth_vis":
|
||||
self.assertTrue((result_directory / "0000.jpg").is_file())
|
||||
self.assertFalse((result_directory / "depth_vis").exists())
|
||||
|
||||
self.assertEqual(len(inference_options), 5)
|
||||
self.assertEqual(len(export_calls), 4)
|
||||
self.assertEqual(len(inference_options), 4)
|
||||
self.assertEqual(len(mesh_calls), 1)
|
||||
for images, options in inference_options:
|
||||
self.assertEqual(images, [str(input_path)])
|
||||
@@ -964,82 +892,21 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
self.assertTrue(options["use_ray_pose"])
|
||||
self.assertEqual(options["ref_view_strategy"], "middle")
|
||||
for (
|
||||
_prediction,
|
||||
export_format,
|
||||
_export_dir,
|
||||
options,
|
||||
) in export_calls:
|
||||
np.testing.assert_array_equal(
|
||||
_prediction.depth,
|
||||
np.full((1, 2, 2), 2.0, dtype=np.float32),
|
||||
)
|
||||
if export_format == "glb":
|
||||
self.assertEqual(
|
||||
options,
|
||||
{
|
||||
"glb": {
|
||||
"conf_thresh_percentile": 35.0,
|
||||
"num_max_points": 12345,
|
||||
"show_cameras": False,
|
||||
"export_depth_vis": False,
|
||||
}
|
||||
},
|
||||
)
|
||||
elif export_format == "gs_ply":
|
||||
self.assertEqual(
|
||||
options,
|
||||
{"gs_ply": {"gs_views_interval": 2}},
|
||||
)
|
||||
else:
|
||||
self.assertEqual(options, {})
|
||||
self.assertEqual(mesh_calls[0][1], [input_path])
|
||||
self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2))
|
||||
|
||||
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
|
||||
self.assertEqual(
|
||||
infer_gs_values,
|
||||
[False, True, False, False, False],
|
||||
)
|
||||
|
||||
def test_reconstruct_copies_frames_for_multi_frame_input(self):
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
return cls()
|
||||
|
||||
def to(self, device):
|
||||
return self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
def inference(self, images, **options):
|
||||
return SimpleNamespace(
|
||||
fake_da3 = ModuleType("da3_onnx")
|
||||
fake_da3.create_session = lambda _model_dir, _device: object()
|
||||
fake_da3.infer = lambda *_args, **_kwargs: SimpleNamespace(
|
||||
depth=np.full((2, 2, 2), 2.0, dtype=np.float32),
|
||||
conf=np.ones((2, 2, 2), dtype=np.float32),
|
||||
extrinsics=np.repeat(np.eye(4, dtype=np.float32)[None, :3], 2, axis=0),
|
||||
intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0),
|
||||
processed_images=np.zeros((2, 2, 2, 3), dtype=np.uint8),
|
||||
)
|
||||
|
||||
def fake_export(prediction, export_format, export_dir, **options):
|
||||
mini_directory = Path(export_dir) / "exports" / "mini_npz"
|
||||
mini_directory.mkdir(parents=True)
|
||||
np.savez_compressed(
|
||||
mini_directory / "results.npz",
|
||||
depth=np.zeros((2, 14, 14)),
|
||||
)
|
||||
|
||||
fake_torch = ModuleType("torch")
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(
|
||||
type=device_type,
|
||||
)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_api = ModuleType("depth_anything_3.api")
|
||||
fake_api.DepthAnything3 = FakeModel
|
||||
fake_export_module = ModuleType("depth_anything_3.utils.export")
|
||||
fake_export_module.export = fake_export
|
||||
fake_utils = ModuleType("depth_anything_3.utils")
|
||||
fake_utils.__path__ = []
|
||||
fake_package = ModuleType("depth_anything_3")
|
||||
fake_package.__path__ = []
|
||||
fake_runtime = ModuleType("onnx_runtime")
|
||||
fake_runtime.provider_label = lambda _session: "CPU"
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
@@ -1050,11 +917,8 @@ class WorkerTest(unittest.TestCase):
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"depth_anything_3": fake_package,
|
||||
"depth_anything_3.api": fake_api,
|
||||
"depth_anything_3.utils": fake_utils,
|
||||
"depth_anything_3.utils.export": fake_export_module,
|
||||
"da3_onnx": fake_da3,
|
||||
"onnx_runtime": fake_runtime,
|
||||
},
|
||||
):
|
||||
args = SimpleNamespace(
|
||||
|
||||
@@ -1,32 +1,5 @@
|
||||
schema_version = 1
|
||||
|
||||
[artifacts.depth-anything-3]
|
||||
revision = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
|
||||
filename = "depth-anything-3.tar.gz"
|
||||
url = "https://codeload.github.com/ByteDance-Seed/Depth-Anything-3/tar.gz/3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
|
||||
sha256 = "98aa2dd53ab44b96cef5190ae4841c6ae51797d099b784e08e12a1a55bd3a69b"
|
||||
patches = "depth-anything-3"
|
||||
|
||||
[artifacts.ben2]
|
||||
revision = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
|
||||
filename = "ben2.tar.gz"
|
||||
url = "https://codeload.github.com/PramaLLC/BEN2/tar.gz/2c99a5da477b5523585bfa5c893888a6e818a8f6"
|
||||
sha256 = "01fd26716f3ff3a3f18ad2336c768af5b759edfc7c38bb3bf733b7c86eee6282"
|
||||
|
||||
[artifacts.moge]
|
||||
revision = "925b8ed835a7a9cdb7578ba15c658a0afc969030"
|
||||
filename = "moge.tar.gz"
|
||||
url = "https://codeload.github.com/microsoft/MoGe/tar.gz/925b8ed835a7a9cdb7578ba15c658a0afc969030"
|
||||
sha256 = "c6d90b2564d643f0f3ae65b5ffe163abaea3b55d194c6e2ba9976d67f552bd26"
|
||||
patches = "moge"
|
||||
|
||||
[artifacts.utils3d]
|
||||
revision = "3fab839f0be9931dac7c8488eb0e1600c236e183"
|
||||
filename = "utils3d.tar.gz"
|
||||
url = "https://codeload.github.com/EasternJournalist/utils3d/tar.gz/3fab839f0be9931dac7c8488eb0e1600c236e183"
|
||||
sha256 = "f1f167b1f68809bea391a66dd0d2ad0e14f7382cbddf087b655181c31a285abb"
|
||||
patches = "utils3d"
|
||||
|
||||
[uv]
|
||||
version = "0.11.3"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user