From 835225d26451cc4e81bdb6a58fe95f74694d25f6 Mon Sep 17 00:00:00 2001 From: nan Date: Tue, 18 Aug 2026 23:27:10 +0800 Subject: [PATCH] feat: Add surface generation and media handling - Implemented surface generation functionality in `surface.py`, including geometry estimation and texture writing. - Introduced media input handling in `input.py`, supporting image and video frame extraction. - Added model management for geometry adapters in `models/__init__.py`. - Created BEN2 and DA3 model inference logic in `ben2.py` and `da3.py`, respectively. - Developed ONNX model inference utilities for BEN2 and DA3 in `onnx_ben2.py` and `onnx_da3.py`. - Added support for MoGe-2 model inference in `moge2.py` with necessary ONNX utilities. - Implemented output handling for depth and surface data in `depth.py` and `surface.py`. - Created debug output utilities for DA3 and MoGe-2 predictions in `debug_da3.py` and `debug_moge2.py`. - Enhanced ONNX runtime session management in `onnx_runtime.py`. - Updated `__init__.py` files to include new modules and functions for better organization. --- README.md | 18 +- docs/depth-plane.md | 6 +- docs/development.md | 49 +- docs/features.md | 17 +- docs/index.md | 2 +- docs/outputs.md | 39 +- docs/server.md | 14 +- docs/surface.md | 59 +- src/snapform/assets/O_SnapForm.blend | 4 +- src/snapform/blender/depth_plane.py | 15 + src/snapform/blender/surface.py | 156 ++--- src/snapform/operators.py | 55 +- src/snapform/panel.py | 35 +- src/snapform/properties.py | 98 --- src/snapform/server/__init__.py | 41 +- src/snapform/server/background.py | 170 ----- src/snapform/server/da3.py | 232 ------- src/snapform/server/geometry.py | 616 ------------------ src/snapform/server/geometry/__init__.py | 35 + src/snapform/server/geometry/grid_mesh.py | 202 ++++++ src/snapform/server/geometry/projection.py | 31 + src/snapform/server/geometry/surface.py | 133 ++++ src/snapform/server/jobs/__init__.py | 5 + src/snapform/server/jobs/background.py | 78 +++ src/snapform/server/jobs/common.py | 20 + src/snapform/server/jobs/debug_ben2.py | 26 + src/snapform/server/jobs/debug_da3.py | 21 + src/snapform/server/jobs/debug_moge2.py | 26 + src/snapform/server/jobs/depth.py | 49 ++ src/snapform/server/jobs/foreground.py | 51 ++ src/snapform/server/jobs/frames.py | 21 + src/snapform/server/jobs/surface.py | 75 +++ src/snapform/server/media.py | 306 --------- src/snapform/server/media/__init__.py | 26 + src/snapform/server/media/input.py | 158 +++++ src/snapform/server/model_manager.py | 6 +- src/snapform/server/models/__init__.py | 8 + src/snapform/server/models/ben2.py | 78 +++ src/snapform/server/models/da3.py | 94 +++ src/snapform/server/models/moge2.py | 119 ++++ src/snapform/server/{ => models}/onnx_ben2.py | 0 src/snapform/server/{ => models}/onnx_da3.py | 0 .../server/{ => models}/onnx_moge2.py | 0 .../server/{ => models}/onnx_runtime.py | 0 src/snapform/server/moge2.py | 304 --------- src/snapform/server/outputs/__init__.py | 11 + src/snapform/server/outputs/debug_da3.py | 24 + src/snapform/server/outputs/debug_moge2.py | 13 + src/snapform/server/outputs/depth.py | 64 ++ src/snapform/server/outputs/surface.py | 44 ++ tests/test_addon.py | 14 +- tests/test_ben2_onnx.py | 1 + tests/test_da3_onnx.py | 1 + tests/test_moge_onnx.py | 22 +- tests/test_pack.py | 27 +- tests/test_server.py | 30 +- tests/test_worker.py | 555 ++++++++-------- tools/build_geometry_nodes_asset.py | 24 +- tools/validate_depth_plane_asset.py | 7 +- tools/validate_geometry_nodes_asset.py | 23 +- tools/validate_moge_depth_plane.py | 4 + tools/validate_moge_surface.py | 82 ++- 62 files changed, 2126 insertions(+), 2318 deletions(-) delete mode 100644 src/snapform/server/background.py delete mode 100644 src/snapform/server/da3.py delete mode 100644 src/snapform/server/geometry.py create mode 100644 src/snapform/server/geometry/__init__.py create mode 100644 src/snapform/server/geometry/grid_mesh.py create mode 100644 src/snapform/server/geometry/projection.py create mode 100644 src/snapform/server/geometry/surface.py create mode 100644 src/snapform/server/jobs/__init__.py create mode 100644 src/snapform/server/jobs/background.py create mode 100644 src/snapform/server/jobs/common.py create mode 100644 src/snapform/server/jobs/debug_ben2.py create mode 100644 src/snapform/server/jobs/debug_da3.py create mode 100644 src/snapform/server/jobs/debug_moge2.py create mode 100644 src/snapform/server/jobs/depth.py create mode 100644 src/snapform/server/jobs/foreground.py create mode 100644 src/snapform/server/jobs/frames.py create mode 100644 src/snapform/server/jobs/surface.py delete mode 100644 src/snapform/server/media.py create mode 100644 src/snapform/server/media/__init__.py create mode 100644 src/snapform/server/media/input.py create mode 100644 src/snapform/server/models/__init__.py create mode 100644 src/snapform/server/models/ben2.py create mode 100644 src/snapform/server/models/da3.py create mode 100644 src/snapform/server/models/moge2.py rename src/snapform/server/{ => models}/onnx_ben2.py (100%) rename src/snapform/server/{ => models}/onnx_da3.py (100%) rename src/snapform/server/{ => models}/onnx_moge2.py (100%) rename src/snapform/server/{ => models}/onnx_runtime.py (100%) delete mode 100644 src/snapform/server/moge2.py create mode 100644 src/snapform/server/outputs/__init__.py create mode 100644 src/snapform/server/outputs/debug_da3.py create mode 100644 src/snapform/server/outputs/debug_moge2.py create mode 100644 src/snapform/server/outputs/depth.py create mode 100644 src/snapform/server/outputs/surface.py diff --git a/README.md b/README.md index 054b32d..dbe694b 100644 --- a/README.md +++ b/README.md @@ -53,10 +53,12 @@ Environment 未安装时，同一位置显示两个安装按钮。Environment 选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Convert to Surface`。可选择 MoGe-2 ViT-S、ViT-B、ViT-L 或 DA3 Small、Base，其中 MoGe-2 ViT-S 为默认模型。SnapForm 会按模型输出的 -相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。MoGe-2 使用内置 NumPy -网格构建与深度断层 mask 生成干净拓扑；DA3 保留 confidence 与边拉伸剔除。两者再由独立的 Geometry Nodes -修改器完成剔面、快速重拓扑、Relax 与 UV 重投。开启 Remove Background 时，背景三角面会在转换期间 -直接删除，不再交给 Geometry Nodes。匹配相机默认不创建；开启后会作为 Surface 的子级。 +相机内参与三维点建立透视网格，并匹配原 Empty 的可见尺寸、Offset 与 Transform。Server 通过统一的 +Surface 输出流程构建 DA3 与 MoGe-2 源网格：MoGe-2 使用深度断层 mask，DA3 保留 confidence 与边拉伸 +数据。两者再由独立的 Geometry Nodes +修改器完成剔面、快速重拓扑、Relax 与 UV 重投；顶层 `Iterations` 可设为 `0`。开启 Remove Background +时，BEN2 Alpha 只要大于 `0` 就视为前景；三角面有任一前景点即保留，只有全背景面会删除并且不再交给 +Geometry Nodes；Surface 材质使用 `.foreground`。匹配相机默认不创建；开启后会作为 Surface 的子级。 MoGe-2 的像素级 normal 会作为 Object Space normal 贴图接入 `O MoGe2 Surface Layer`，并与原图 Bump 合并；网格本身不写入 custom normal，因此重拓扑降低顶点数后仍能保留细腻的法线细节。 当前此功能仅支持单张图片。 @@ -70,14 +72,18 @@ MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base 与 BEN2 全部使用 ONNX 以及分割线下方的 `Keep Original`。 界面会显示计算所得的实际 Process Res。剔除和重拓扑参数直接在 Geometry Nodes 修改器中调整；`Unshaded` 可使用与 Depth Plane 相同的 Emission 材质。普通材质的 -Bump Scale 会在材质组内乘以 `0.001`，并按透视表面的深度空间缩放。Surface attributes 统一使用 `o_` 前缀；Server 准备过程显示 +Bump Scale 会在材质组内乘以 `0.001`，并按透视表面的深度空间缩放。Surface 与 Depth Plane 的 +Layer Color 同时驱动 Principled BSDF 的 Base Color、Subsurface Radius 和 Emission Color，Layer Normal +同时驱动 Normal 与 Coat Normal；Subsurface、Emission、Coat 的其他设置保持默认。Surface attributes +统一使用 `o_` 前缀；Server 准备过程显示 `Preparing surface data n/6`，可看到当前阶段与总步骤。 ### 视频与序列帧 SnapForm 支持视频（mp4、mkv、mov 等）和序列帧，会逐帧处理并生成逐帧对应的立体面片。处理范围默认 跟随 Empty 上设置的 `Image Frames` 与 `Offset`，不需要额外配置。DA3 的转换任务复用同一个 ONNX -Session 逐帧推理，显存峰值不随视频帧数平方增长；Debug 多帧输出仍使用多视图联合推理。 +Session 逐帧推理，显存峰值不随视频帧数平方增长。Debug 只保存模型 API 的直接返回：DA3 多帧输入 +保留联合推理，MoGe-2 与 BEN2 按各自的单图 API 逐帧保存。 ## 任务结果保存在哪里 diff --git a/docs/depth-plane.md b/docs/depth-plane.md index 37a3cfe..ba8b158 100644 --- a/docs/depth-plane.md +++ b/docs/depth-plane.md @@ -11,7 +11,7 @@ `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 进程内用内置 NumPy 网格构建生成 `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 则逐帧保存 API 直接返回的 points、depth、normal、mask 与 intrinsics，不生成网格或预览。 - `Original`：直接把原图作为深度贴图，不显示估算参数。仅在开启 `Remove Background` 时运行 BEN2 取得前景 Alpha；视频或序列帧按需同时抽取原始帧与前景序列，把原始帧作为 `.depth` 序列。 `MoGe-2 ViT-S` 是默认深度来源。选择 `DA3` 时，其 Model 与 Resolution Scale @@ -23,7 +23,9 @@ Image Empty 引用视频或序列帧时，DA3 与 MoGe-2 会按 Empty 的 Image Frames 数值逐帧估计 深度（未设置时为完整帧数，显式设为 1 时只处理 1 帧），起始帧跟随 Empty 的 Offset，与背景蒙版范围一致； DA3 会复用同一个 ONNX Session 对每帧独立推理，再合并为原有序列输出，避免多帧 Attention 的平方级显存占用。 - 该行为只用于 Depth Plane 转换；Surface 仅支持单张图片，Debug 多帧输出仍联合推理。 + DA3 与 MoGe-2 均先用未经抠图的原始 RGB 生成深度，再运行可选 BEN2；BEN2 Alpha 只参与最终输出过滤和材质 Mask。 + 该行为只用于 Depth Plane 转换；Surface 仅支持单张图片。Debug 中 DA3 多帧输入联合推理， + MoGe-2 与 BEN2 按单图 API 逐帧运行。 Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty 一致；Original 模式的视频或序列帧使用 `frame.0001.jpg` 等原图序列作为颜色与深度， `foreground.0001.png` 只提供前景 Alpha。 diff --git a/docs/development.md b/docs/development.md index 0eda2c1..2e35b11 100644 --- a/docs/development.md +++ b/docs/development.md @@ -24,9 +24,14 @@ flowchart LR - `blender/depth_plane.py` 与 `blender/surface.py` 实现 Image Empty 输入检查、几何构建和资产挂载 - `job_runtime.py` 以入口字符串声明独立 Server，并配置 Environment、Storage Root 与首次安装后的推荐模型下载 - `server/__init__.py` 创建 `JobServer`、注册公开 Job，并把 `ModelManager` 注册为 Resource -- `server/media.py`、`background.py`、`da3.py`、`moge2.py` 负责业务流程；`onnx_*.py` 只处理模型输入、推理和输出恢复 +- `server/jobs/` 按任务组织完整流程；DA3、MoGe-2、BEN2 Debug 各自拥有独立模块和 Job，不用模型分支总控函数 +- `server/models/` 包含三个模型适配器及其 ONNX 实现，只负责模型载入、推理和 Prediction；不接收 Job Context、不写文件 +- `server/geometry/` 包含投影、像素网格和唯一 Surface Builder，均为纯数据转换；不负责 Debug 或文件输出 +- `server/outputs/` 只负责 Depth、Surface 与各模型 Debug 的序列化；Debug writer 原样保存 API 字段，不构建派生几何 +- `server/media/input.py` 只处理图片、序列和视频输入；共享的可选前景流程保留在 `jobs/foreground.py` - 根目录 `model_catalog.py` 供 Blender UI 与本地就绪检查使用；`server/model_catalog.py` 独立声明下载与推理所需模型，二者的模型键、目录名和 `ready_patterns` 由测试约束一致 -- `server/geometry.py` 生成 NPZ、预览、点云和纹理网格；所有任务产物写入 Runtime 分配的 opaque-ID Job 目录 +- `server/__init__.py` 只注册 Job 与 Resource；`jobs/background.py`、`jobs/frames.py` 分别实现去背景和帧提取 +- 所有任务产物写入 Runtime 分配的 opaque-ID Job 目录 功能细节继续按主题维护：[Depth Plane](depth-plane.md)、[Surface](surface.md)、[输出格式](outputs.md) 和 [Job Server](server.md)。 @@ -56,18 +61,37 @@ flowchart LR │ ├── model_catalog.py │ ├── server/ │ │ ├── __init__.py -│ │ ├── background.py -│ │ ├── da3.py -│ │ ├── geometry.py -│ │ ├── media.py +│ │ ├── jobs/ +│ │ │ ├── background.py +│ │ │ ├── debug_ben2.py +│ │ │ ├── debug_da3.py +│ │ │ ├── debug_moge2.py +│ │ │ ├── depth.py +│ │ │ ├── foreground.py +│ │ │ ├── frames.py +│ │ │ └── surface.py +│ │ ├── geometry/ +│ │ │ ├── projection.py +│ │ │ ├── grid_mesh.py +│ │ │ └── surface.py +│ │ ├── media/ +│ │ │ └── input.py │ │ ├── model_download.py │ │ ├── model_catalog.py │ │ ├── model_manager.py -│ │ ├── moge2.py -│ │ ├── onnx_ben2.py -│ │ ├── onnx_da3.py -│ │ ├── onnx_moge2.py -│ │ └── onnx_runtime.py +│ │ ├── models/ +│ │ │ ├── ben2.py +│ │ │ ├── da3.py +│ │ │ ├── moge2.py +│ │ │ ├── onnx_ben2.py +│ │ │ ├── onnx_da3.py +│ │ │ ├── onnx_moge2.py +│ │ │ └── onnx_runtime.py +│ │ └── outputs/ +│ │ ├── debug_da3.py +│ │ ├── debug_moge2.py +│ │ ├── depth.py +│ │ └── surface.py │ └── wheels/ │ └── blendjob-0.1.13-py3-none-any.whl ├── tests/ @@ -106,6 +130,9 @@ BLENDER_BIN="/c/Program Files/Blender Foundation/Blender 4.5/blender.exe" "$BLENDER_BIN" --background --python tools/validate_moge_surface.py ``` +验证脚本会自动把 `src/snapform/wheels/` 中的 BlendJob wheel 加入 Blender Python 路径，因此可在 +`--factory-startup` 的干净环境中直接运行。 + 只改 Python、文档或打包逻辑时不需要重建 `.blend` 资产；修改节点组、材质组或其接口时，应运行相应资产验证脚本，并再次运行完整测试。 ## 上游与许可证 diff --git a/docs/features.md b/docs/features.md index 18a26ca..f0cbbef 100644 --- a/docs/features.md +++ b/docs/features.md @@ -5,7 +5,7 @@ `SnapForm` 是一套从图片快速构建空间结构的 Blender Extension。目前以 [Depth Anything 3](https://github.com/ByteDance-Seed/Depth-Anything-3) -作为首个 AI 能力，在独立 Environment 中运行 DA3，可将官方导出结果保存在独立 job 目录， +作为首个 AI 能力，在独立 Environment 中运行 DA3，可将模型 API Prediction 保存在独立 job 目录， 也可把 Image Empty 转换成带几何与材质置换的立体面片。MoGe-2 作为独立深度来源，保留其 metric depth、normal、mask 与 intrinsics 输出，不复用 DA3 Prediction 后处理。Debug 面板还可运行 [BEN2](https://github.com/PramaLLC/BEN2) 背景剔除功能测试。 @@ -26,14 +26,11 @@ - `Models` panel 默认折叠，展开后可逐项查看和下载所有可用模型 - Environment 已安装时不额外占用面板空间；模型列表每个模型只占一行 - `Debug` 默认折叠，内部的 `DA3`、`MoGe-2` 与 `BEN2` 也分别默认折叠 -- 使用一个输出类型下拉菜单和一个 `Generate Output` 按钮 -- DA3 与 MoGe-2 Debug 都先选择单一 Output Format，只显示当前格式的专属参数；GLB 与 PLY 不会在一次测试中同时生成 -- Debug 的模型选项按格式过滤；PLY 只提供两个带 Gaussian head 的 Giant 模型 +- DA3、MoGe-2 与 BEN2 分别使用独立 Debug Job，直接保存各自 API 返回，不共享按模型分支的总控函数 +- DA3 Debug 保存 Prediction NPZ；MoGe-2 Debug 保存逐帧原始预测 NPZ；BEN2 Debug 保存直接返回的透明前景 - Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态；未下载模型带 `(Missing)`，已下载模型保持原名 -- 分别调用 DA3 的 `glb`、`gs_ply`、`depth_vis` 和 `mini_npz` exporter -- `Textured Surface (GLB)` 直接使用 DA3 Prediction 中的深度、置信度和相机参数，把相邻像素 - 组成三角面并写入 `mesh.glb`；每个输入帧生成一张独立纹理网格，不做多视图融合 +- Debug 不运行点云、Surface、深度预览、置信度 percentile 或坐标变换等 SnapForm 后处理 - Image Empty 右键菜单可直接运行 DA3，根据 `mini_npz` 深度生成面片 - Image Empty 右键菜单可直接运行 BEN2，保留原图并换用新建的透明背景图片 - 缺少必需的 `MoGe-2 ViT-S Normal` 或 `BEN2 Base` 时，Image Empty 的 `SnapForm` @@ -43,13 +40,13 @@ - 菜单、模型面板和转换对话框直接检查 Storage Root 的 `models/` 必需文件，显示模型状态时不会启动 Job Server - Image Empty 支持视频与序列帧输入：`Convert to Depth Plane` 逐帧估计深度， `Remove Background` 逐帧抠图，入口保持不变。DA3 转换时每帧独立执行 ONNX 推理并复用模型 Session； - Debug 多帧任务继续联合推理，以保留跨视图相机与深度关系 + DA3 Debug 多帧任务继续联合推理；MoGe-2 与 BEN2 Debug 按单图 API 逐帧运行 - 视频生成的序列帧以外部文件引用保留在 job 目录（Blender 无法整体 Pack 序列帧）； Geometry Nodes 与材质 Image Texture 均可按场景帧播放序列帧 - 使用浮点深度贴图、`O DA3 Depth Plane` Geometry Nodes 和原图 Bump 构建立体面片 - 转换对话框选择推理模式、处理分辨率与背景移除；细分和置换方式在 Geometry Nodes 修改器中调整 - Debug 的 DA3 输出或 BEN2 背景剔除完成后只打开对应 job 文件夹，不修改当前场景 -- 完整保留每次任务生成的所有官方文件 +- Debug 完整保留模型 API 返回的数组字段，不混入 SnapForm 派生几何 - 安装和模型下载使用后台子进程；推理提交给常驻 Server，状态栏实时显示阶段、百分比进度条与取消按钮 - 后台任务运行时，Install、Download 和 Export 按钮保持原位置并变灰，任务完成后恢复 @@ -79,7 +76,7 @@ 官方安装固定使用 GitHub、uv 默认 Python 下载源和 PyPI；镜像安装固定使用 CNB uv Release、`python-build-standalone` 镜像和清华 PyPI。两种方式都不测速、不竞速、不自动切换，镜像配置只注入当前安装子进程。模型仍按自身的 Hugging Face 与 OmooLab R2 逻辑下载。 5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 6. 如需其他模型，展开 `Models` panel，找到模型并点击该行的 `Download`。 -7. 展开默认收起的 `Debug`，选择 `Output Type` 和本次执行使用的模型。 +7. 展开默认收起的 `Debug`，在 DA3、MoGe-2 或 BEN2 区域选择模型与输入，然后直接运行对应 API。 8. 选择一张图片或一个视频文件。视频可设置从头截取的帧数，然后点击 `Generate Output`。 9. 测试背景剔除时，确认 `Models` panel 中的 `BEN2 Base` 已就绪，然后在 `Debug` 的 `BEN2` 区域选择单张图片并点击 `Remove Background`。 diff --git a/docs/index.md b/docs/index.md index 85bb131..dd575e7 100644 --- a/docs/index.md +++ b/docs/index.md @@ -5,6 +5,6 @@ - [功能与使用](features.md)：能力概览、安装使用和 Image Empty 去背景 - [Depth Plane](depth-plane.md)：深度面片的数据、节点与材质流程 - [Surface](surface.md)：透视网格、Geometry Nodes 后处理与材质法线 -- [输出格式](outputs.md)：Debug 输出类型、参数和任务产物 +- [输出格式](outputs.md)：Debug 原始 API 返回与 Depth、Surface 任务产物 - [Job Server](server.md)：Environment、模型缓存、任务目录与运行机制 - [程序设计、开发与许可证](development.md)：模块边界、调用流程、项目结构、常用命令与许可证 diff --git a/docs/outputs.md b/docs/outputs.md index 97cc1c3..a1aa88b 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -1,29 +1,32 @@ # 输出格式 -## DA3 ONNX +## Debug -`Output Type` 提供以下结果： +Debug 用于直接比较 DA3、MoGe-2 与 BEN2 模型 API 的返回效果，不复用 Depth、Surface 或 Geometry +后处理。三个入口对应三个独立 Job： -- Point Cloud (GLB)：由 depth、confidence、相机参数和原图生成彩色点云 -- Depth Preview (Images)：保存深度着色预览 -- Depth Data (NPZ)：保存 depth、confidence、extrinsics、intrinsics 与处理后的图像 -- Textured Surface (GLB)：SnapForm 根据 Prediction 构建纹理网格 +- debug-da3：保存 prediction.npz，字段直接来自 DA3 Prediction，包括 depth、conf、 + extrinsics、intrinsics、processed_images 及模型实际返回的可选标量或数组 +- debug-moge2：每张输入保存一个 prediction.npz 或 prediction.0001.npz，完整保留 MoGe-2 API + 字典中的 points、depth、normal、mask 与 intrinsics +- debug-ben2：直接保存 BEN2 API 返回的 foreground.png 或 foreground.0001.png -DA3 Small 与 Base 的 ONNX 图不包含 Gaussian head，因此不提供 Gaussian Splats。默认推理参数为： +Debug 不生成点云、三角网格、深度预览、置信度 percentile 或坐标变换后的派生产物。输入为多帧时， +DA3 保留联合推理；MoGe-2 与 BEN2 按其单图 API 逐帧保存结果。 -- `process_res=504` -- `process_res_method="upper_bound_resize"` -- `ref_view_strategy="saddle_balanced"` -- `conf_thresh_percentile=40.0` -- `num_max_points=1_000_000` -- `show_cameras=True` +## Depth -相机外参与内参由模型的 camera decoder 输出恢复，不运行 ray-pose 求解器。 +generate-depth 输出 Blender Depth Plane 所需的 results.npz。DA3 保存 depth、conf、 +extrinsics 与 intrinsics；MoGe-2 保存 depth、normal、mask 与 intrinsics，并额外写出 +Depth Plane 导入所需的贴图。这里的文件转换属于正式 Depth 工作流，不属于 Debug。 -## MoGe-2 ONNX +## Surface -MoGe-2 Debug 可保存原始 NPZ、depth/normal/mask maps、纹理 GLB 或彩色 PLY。Surface 输出由模型 points、mask 与 depth edge 在内部构建三角网格。 +generate-surface 统一写出 mesh.npz 与 texture.png。DA3 与 MoGe-2 都使用 +geometry/surface.py 的同一套 Surface 构建算法；序列化位于 outputs/surface.py。 +MoGe-2 normal 另存为 normal.png 并由 Blender 材质加载，不写入几何 NPZ。 -## BEN2 ONNX +## BEN2 -BEN2 输出 `foreground.png`，RGB 保留原始输入，Alpha 来自 ONNX mask 并恢复到输入分辨率。 +正式 remove-background 工作流输出 foreground.png 或透明前景序列。BEN2 ONNX API 返回的 Alpha +恢复到输入分辨率，RGB 保留输入图像颜色。 diff --git a/docs/server.md b/docs/server.md index f5ab424..990a4f2 100644 --- a/docs/server.md +++ b/docs/server.md @@ -39,17 +39,19 @@ SnapForm 的业务 Operator 直接提交参数字典，Job Type 一一对应： - `download_required_models` → `download-required-models` - `remove_background` → `remove-background` - `extract_frames` → `extract-frames` -- `generate_da3_depth` → `generate-da3-depth` -- `generate_moge2_depth` → `generate-moge2-depth` -- `generate_da3_surface` → `generate-da3-surface` -- `generate_moge2_surface` → `generate-moge2-surface` +- `generate_depth` → `generate-depth` +- `generate_surface` → `generate-surface` - `debug_da3` → `debug-da3` - `debug_moge2` → `debug-moge2` - `debug_ben2` → `debug-ben2` DA3 与 MoGe-2 的生产任务不再合并成带 Backend 参数的通用请求。每个 Operator 的 `request(context)` 直接构造自己的有效参数，Debug 任务也与生产任务完全分开。 -Server 端没有 `execute_*`、`*_result` 或 CLI worker 中转。全部 `@server.job(...)` 都集中在 `server/__init__.py`：简单任务直接实现，复杂任务调用按领域拆分的 `media.py`、`background.py`、`moge2.py` 或 `da3.py`；共享的纯几何计算位于 `geometry.py`。模型下载由 `ModelManager` 单向调用 `model_download.py`，不存在循环 import。 +Server 端没有 `execute_*`、`*_result` 或 CLI worker 中转。全部 `@server.job(...)` 都集中在 +`server/__init__.py`，这里只负责注册 Job 与 Resource。完整工作流位于 `jobs/`；模型适配器和 +ONNX 实现位于 `models/`；纯投影、网格与唯一 Surface Builder 位于 `geometry/`；文件序列化位于 +`outputs/`；输入准备位于 `media/`。模型下载由 `ModelManager` 单向调用 +`model_download.py`，不存在循环 import。 ## ONNX-only Environment @@ -114,7 +116,7 @@ Server 由 Environment 通过 `python -m blendjob.runner` 启动，避免 BlendJ DA3 ONNX 的注意力是朴素实现（`MatMul → Softmax → MatMul`），峰值显存与 `(帧数 × 每帧 patch 数)²` 成正比。迁移到 ONNX 前的 PyTorch 版使用 `scaled_dot_product_attention`（`XFORMERS_DISABLED=1`），按块计算注意力分数、不落地完整注意力矩阵，显存近似线性，因此视频能处理几十上百帧；ONNX 导出把 SDPA 回退成朴素注意力，多帧在中高分辨率下会显存爆炸（约 12 帧 @ 1008 即需 ~90 GB）。 -Depth Plane 的视频、序列帧转换会复用同一个 ONNX Session 逐帧独立推理，再把结果合并到现有序列输出，因此峰值显存只取决于单帧分辨率。逐帧结果不保留 DA3 的跨视图尺度与相机位姿一致性，适用于当前逐帧动画结构。Surface 只支持单张图片；Debug 多帧输出仍使用联合推理，以保留多视图关系。 +Depth Plane 的视频、序列帧转换会复用同一个 ONNX Session 逐帧独立推理，再把结果合并到现有序列输出，因此峰值显存只取决于单帧分辨率。逐帧结果不保留 DA3 的跨视图尺度与相机位姿一致性，适用于当前逐帧动画结构。Surface 只支持单张图片。DA3 Debug 多帧输入仍联合推理；MoGe-2 与 BEN2 Debug 按单图 API 逐帧保存直接返回。 联合推理时，`onnx_da3.infer` 会在执行前按 `views × patches` 预估注意力张量，超过 4 GiB 直接报错并提示减少帧数或 Resolution Scale，避免 onnxruntime 报出难以定位的分配失败或 `UnicodeDecodeError`（DirectML 在中文 Windows 下报错信息为 GBK 编码，被 UTF-8 解码掩盖）。 diff --git a/docs/surface.md b/docs/surface.md index 49ecb30..390cbfd 100644 --- a/docs/surface.md +++ b/docs/surface.md @@ -7,15 +7,20 @@ `Convert to Depth Plane`，当前不处理视频或序列帧。 模型列表包含三个 MoGe-2 Normal 模型和两个 DA3 模型，默认使用 MoGe-2 ViT-S Normal。DA3 的 Resolution Scale 与 Depth Plane 使用相同设计；MoGe-2 改为显示共享的 Resolution Level。 -2. Server 请求 DA3 的 `depth`、`conf` 与 `intrinsics`。仅在开启 `Remove Background` 时运行 BEN2 取得 - 前景 Alpha，并在输出前删除包含背景顶点的三角面；关闭时不加载 BEN2。Depth 不做 smoothing，Server - 反投影深度网格，建立源三角面，计算 confidence percentile 与 3×3 相对深度不连续度，并写入 +2. Server 先用未经抠图的原始 RGB 请求 DA3 的 `depth`、`conf` 与 `intrinsics`，随后仅在开启 + `Remove Background` 时运行 BEN2 取得前景 Alpha；Alpha 经目标网格分辨率缩放与轻微模糊后，只要 + 大于 `0` 就把对应点视为前景；每个三角面只要有一个前景点就保留，三个点均为背景时才在输出前删除； + 关闭时不加载 BEN2。Depth 不做 smoothing，Server 反投影深度网格，建立源三角面， + 计算 confidence percentile 与 3×3 相对深度不连续度，并写入 基于原图像素坐标的源 `UVMap`。这些步骤只准备修改器的输入网格，不执行 Relax、拓扑修复、平滑、 小岛清理或最终 UV 重投影。运行时显示的六个 `Preparing surface data` 步骤均指这段 Server 端源数据准备。 + DA3 与 MoGe-2 最终都由 `server/geometry/surface.py` 构建同一份 Surface 数据；模型模块只提供各自的 + 深度、点、法线、Mask 与相机数据，不再自行拥有 Surface 网格输出实现。 3. Blender 把勾选 `Usage > Modifier` 的 `O DA3 Surface` 节点组直接作为 Geometry Nodes 修改器 添加到 Surface，不另设修改器名称。节点组内没有子节点组，依次完成： - - 修改器顶层依次显示 `Resolution`、`Relax Iterations`、`Smooth`；默认展开的 `Culling` Panel - 依次包含 `Depth Continuity` 与 `Confidence`。前者直接读取面域 3×3 相对深度差，并使用指数映射： + - 修改器顶层依次显示 `Iterations`、`Resolution`、`Smooth`；`Iterations` 默认 `2`、最小 `0`， + 默认展开的 `Culling` Panel 依次包含 `Depth Continuity` 与 `Confidence`。前者直接读取面域 3×3 + 相对深度差，并使用指数映射： 默认 `50%` 对应官方 `rtol=0.04`，数值越大，内部阈值越小，剔除的断层面越多； - 每轮 Remesh Iteration 依次执行 Connected Merge、循环拓扑修复与一次 Smooth，使新拓扑在进入下一轮 Merge 前先稳定；拓扑修复删除非流形面、单面拓扑和小面积 mesh islands； @@ -37,18 +42,19 @@ 修改器关闭 `show_group_selector`。`Fix Topology` 位于默认折叠的 `Options` Panel。 节点图按 Culling、Remesh、Topology Fix、Island Cleanup 与 UV 分区摆放；各区就近使用独立的 Group Input，并隐藏没有连线的接口，避免跨越整个节点图的输入连线。 - `O DA3 Surface` 与 `O MoGe2 Surface` 的 `Resolution` 默认值均为 `0.1`，所有输入 socket 的 - Structure Type 均固定为 `Single`。 -5. Server 准备过程固定显示 6 个编号阶段：验证 DA3 数据、反投影、准备 DA3 与可选 BEN2 数据、构建原始 - 源三角面、源 UV/attributes、写出网格。状态栏显示 `Preparing surface data n/6: ...`。 + `O DA3 Surface` 与 `O MoGe2 Surface` 的 `Iterations` 默认值均为 `2`、最小值为 `0`， + `Resolution` 默认值均为 `0.1`，所有输入 socket 的 Structure Type 均固定为 `Single`。 +5. Surface Job 统一按输入准备、模型推理、可选 BEN2、构建 SurfaceMesh、写出产物五个阶段更新进度； + Surface Builder 是不接收 Job Context 或 progress callback 的纯数据转换。 6. Server 在 job 根目录写出 `mesh.npz` 与 `texture.png`。NPZ 保存 camera-space `vertices`、 `faces`、`uv`、`intrinsics`、`image_size` 与有效深度中位数 `reference_depth`。Blender 读取这些 - 数据后直接创建 Surface 与原图材质，不需要 Open3D 或 OpenCV 进入 Blender Python 环境。 + 数据后直接创建 Surface；材质默认使用原图，开启 Remove Background 时改用 `.foreground`，不需要 Open3D 或 OpenCV 进入 Blender Python 环境。 DA3 Surface 只保留 Geometry Nodes 后处理实际使用的 attributes：点域 `o_depth`、面域 - `o_face_confidence` 与面域 `o_depth_discontinuity`。Remove Background 已在 Server + `o_confidence` 与面域 `o_depth_discontinuity`。Remove Background 已在 Server 建面前直接删除背景三角面，不写入任何 background attribute。MoGe-2 的导入网格保留点域 `o_depth`， 以及面域 `o_mask`、`o_depth_discontinuity`，不伪造 confidence、stretch 或 background 属性， - 也不写入 custom normal。完整分辨率的 `normal_map` 保存在 NPZ 中并转换为材质贴图。 + 也不写入 custom normal。MoGe-2 推理返回的 normal 直接编码为独立的 `normal.png`，Blender 将它 + 作为 Non-Color 贴图导入现有材质；normal 不进入 `mesh.npz`。 所有 MoGe-2 Image datablock 都以源图片名而非对象名为基名：depth、normal 与 mask 分别使用 `图片名.depth`、`图片名.normal`、`图片名.mask`，不附加 `.moge2`。 7. 网格以 `reference_depth` 截面匹配 Image Empty 的显示矩形：该截面的投影四角与 Empty 的可见 @@ -59,6 +65,7 @@ `Remove Background`、`Unshaded`，以及分割线下方的 `Keep Original`。背景移除默认关闭；开启后 Server 使用 BEN2，并在写出 Surface 前直接删除背景三角面，不再由 Geometry Nodes 的 Background 开关处理。 该预删除同时适用于 DA3 与 MoGe-2；因此 MoGe-2 转换后即使关闭 `Mask`，BEN2 已删除的背景也不会恢复。 + 开启 Remove Background 时，Surface 材质的 Color Image Texture 使用打包后的 `.foreground`，原 Empty 仍保持原图。 `Unshaded` 使用与 Depth Plane 相同的 Emission/Transparent 材质路径；普通材质的 Bump Scale 按 `abs(object.scale.z) × camera scale_z` 设置，并在材质组内乘以固定基准 `0.001`， 与透视表面实际深度尺度一致；它允许负数以反转凹凸方向。 @@ -76,12 +83,6 @@ Process Res 的设计依据 DA3 官方实现：API 与 CLI 默认使用 `process [输入预处理实现](https://github.com/ByteDance-Seed/Depth-Anything-3/blob/main/src/depth_anything_3/utils/io/input_processor.py) 与[技术报告](https://depth-anything-3.github.io/assets/da3_tech_report_2025.pdf)。 -Debug 的 Textured Surface GLB 会把 depth、confidence 与 confidence percentile 写入自定义顶点 -attributes。其他 DA3 Debug 三维输出会同时写出未压缩的 `prediction.attributes.npz`，保存完整 -depth、confidence、intrinsics、extrinsics 与 confidence percentile。当前 trimesh 的 glTF 点图元 -导出器只写 Position/Color，无法可靠把任意 attribute 嵌入点云 GLB，因此点云先以 companion NPZ -保留完整结构；后续若改为 SnapForm 自有 glTF writer，可再把这些字段直接嵌入点图元。 - 所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的进度条内部，取消按钮 位于进度条右侧。状态字段支持阶段、当前操作与预计剩余秒数，界面统一组合为紧凑文案并保证 进度单调递增。任务进程与 Server 请求在后台启动，按钮点击后立即进入进度状态。Image Depth 的阶段包括 Environment 启动、输入准备、模型加载、深度估算、缓存 @@ -130,9 +131,10 @@ Depth Plane 使用 `O DA3 Depth Plane Layer` 材质组，并把背景剔除、 材质组、Principled BSDF 与 Material Output，方便继续调整。 Depth Plane 的 Bump Height 直接使用材质组的 `Color` 输入，只受独立的 `Bump Scale` 影响，不使用深度贴图， 也不经过 `Displacement` 模式开关、`Depth Scale` 或 `Background Depth`；材质 Displacement 仍使用这些参数处理后的深度。 -无论选择哪种置换方式，原始彩色贴图都会连接 -Principled BSDF 的 `Base Color`，并通过 -`Bump.Height → Bump.Normal → Principled.Normal` 提供表面细节。Bump 的基础 Distance 为 +无论选择哪种置换方式，Layer 的 `Color` 都会连接 Principled BSDF 的 `Base Color`、 +`Subsurface Radius` 与 `Emission Color`；Layer 的 `Normal` 同时连接 `Normal` 与 `Coat Normal`， +通过 `Bump.Height → Bump.Normal` 提供表面细节。Subsurface、Emission 与 Coat 的强度等其他输入 +保持 Principled BSDF 默认值。Bump 的基础 Distance 为 `0.001`，会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到 Principled BSDF 的 `Roughness`，曲线两端输出分别为 `0.5` 和 `0.1`；材质 IOR 保持 Principled BSDF 默认值 `1.5`。 @@ -162,23 +164,20 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN `depth.0001.exr`、`frame.0001.jpg`（`.frame`）或 `foreground.0001.png`，并在每帧保留原始 浮点深度；保存 `.blend` 后再次打开仍需对应 job 目录存在。 -选择 PLY 后，模型菜单只会提供 `DA3 Giant 1.1` 和 -`DA3 Nested Giant Large 1.1`。Small 和 Base 没有 Gaussian head，无法由 -DA3 生成 `gs_ply`。 - 每个 Server Job 被接受时，Job Runtime 在 Runtime Storage Root 的 `jobs/` 下创建一个独立文件夹。文件夹名是 32 位 opaque Job ID，不包含任务类型或对象名称。执行逻辑从 `JobContext.directory` 获取路径，输出产物直接写入该目录。 Blender 为 Packed Image 或序列帧生成的输入副本先保存在系统临时目录，Server 读取完成后由 Blender 清理，不归档到 Job 目录。普通磁盘文件和视频仍直接读取原文件。Server 日志统一写入 `server.log`，进度只保存在 Server 内存中。Environment Bootstrap 使用 Storage Root 下的 `install.log` 与临时的 `install-status.json`，不创建 Job 目录。 -与几何输出相关的 Job Type 为 `generate-da3-depth`、`generate-moge2-depth`、`generate-da3-surface`、`generate-moge2-surface`、`debug-da3` 与 `debug-moge2`。每种 Job 只接收对应模型与输出所需的参数。 +正式几何输出由 `generate-depth` 与 `generate-surface` 负责。Debug 使用彼此独立的 +`debug-da3`、`debug-moge2` 与 `debug-ben2`，只保存模型 API 的直接返回，不进入 Geometry 或 +Surface 输出流程。 -DA3 ONNX exporter 产生的 `depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺， -不会保留只含一层产物的子目录： +Job 目录中的典型正式产物如下： ```text / -├── scene.glb、*.ply、*.jpg 或 results.npz -└── 其它对应 Debug 格式的产物 +├── results.npz、mesh.npz、texture.png 或 normal.png +└── prediction.npz Debug 原始 API 返回 / ├── foreground.png 单图结果 diff --git a/src/snapform/assets/O_SnapForm.blend b/src/snapform/assets/O_SnapForm.blend index d5c45b9..5bf51a2 100644 --- a/src/snapform/assets/O_SnapForm.blend +++ b/src/snapform/assets/O_SnapForm.blend @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2ed29453e0bed988cccf001eddc8d566eea521c3fb0131f050fdcb7cf608c8e4 -size 1337030 +oid sha256:c317b23d520a1eba5fad92bddfe64043805f65b871fe7fd59e4d4827a22de272 +size 1337018 diff --git a/src/snapform/blender/depth_plane.py b/src/snapform/blender/depth_plane.py index b9631bb..a01d874 100644 --- a/src/snapform/blender/depth_plane.py +++ b/src/snapform/blender/depth_plane.py @@ -1055,9 +1055,12 @@ def create_moge_material( return material shader = nodes.new("ShaderNodeBsdfPrincipled") links.new(layer.outputs["Color"], shader.inputs["Base Color"]) + links.new(layer.outputs["Color"], shader.inputs["Subsurface Radius"]) + links.new(layer.outputs["Color"], shader.inputs["Emission Color"]) links.new(layer.outputs["Alpha"], shader.inputs["Alpha"]) links.new(layer.outputs["Roughness"], shader.inputs["Roughness"]) links.new(layer.outputs["Normal"], shader.inputs["Normal"]) + links.new(layer.outputs["Normal"], shader.inputs["Coat Normal"]) links.new(shader.outputs["BSDF"], output.inputs["Surface"]) links.new(layer.outputs["Displacement"], output.inputs["Displacement"]) set_displacement_only(material) @@ -1426,8 +1429,20 @@ def create_material( material_inputs.outputs["Color"], shader.inputs["Base Color"], ) + links.new( + material_inputs.outputs["Color"], + shader.inputs["Subsurface Radius"], + ) + links.new( + material_inputs.outputs["Color"], + shader.inputs["Emission Color"], + ) links.new(material_inputs.outputs["Alpha"], shader.inputs["Alpha"]) links.new(material_inputs.outputs["Normal"], shader.inputs["Normal"]) + links.new( + material_inputs.outputs["Normal"], + shader.inputs["Coat Normal"], + ) links.new( material_inputs.outputs["Roughness"], shader.inputs["Roughness"], diff --git a/src/snapform/blender/surface.py b/src/snapform/blender/surface.py index 4cdd79c..8753a8e 100644 --- a/src/snapform/blender/surface.py +++ b/src/snapform/blender/surface.py @@ -22,9 +22,9 @@ POINT_ATTRIBUTES = { "o_depth": ("FLOAT", "point_depth"), } FACE_ATTRIBUTES = { - "o_face_confidence": ( + "o_confidence": ( "FLOAT", - "face_confidence_percentile", + "face_confidence", ), } OPTIONAL_FACE_ATTRIBUTES = { @@ -57,10 +57,10 @@ def load_image_mesh_data(npz_path, moge2=False): "point_depth", ] if moge2: - required.extend(("face_mask", "face_depth_discontinuity", "normal_map")) + required.extend(("face_mask", "face_depth_discontinuity")) else: required.extend( - ("face_confidence_percentile", "face_depth_discontinuity") + ("face_confidence", "face_depth_discontinuity") ) missing = [key for key in required if key not in archive.files] if missing: @@ -89,11 +89,6 @@ def load_image_mesh_data(npz_path, moge2=False): for _name, (_data_type, key) in OPTIONAL_FACE_ATTRIBUTES.items(): if key in archive.files: attributes[key] = np.asarray(archive[key]) - if "normal_map" in archive.files: - attributes["normal_map"] = np.asarray( - archive["normal_map"], - dtype=np.float32, - ) except OSError as error: raise ValueError(f"Unable to read surface data: {npz_path}") from error if vertices.ndim != 2 or vertices.shape[1] != 3: @@ -108,14 +103,6 @@ def load_image_mesh_data(npz_path, moge2=False): raise ValueError(f"Unsupported image size: {image_size}") if not reference_depth > 0.0: raise ValueError("Surface reference depth must be positive") - if moge2 and attributes["normal_map"].shape != ( - image_size[1], - image_size[0], - 3, - ): - raise ValueError("Surface normal map and image sizes do not match") - if moge2 and not np.isfinite(attributes["normal_map"]).all(): - raise ValueError("Surface normal map contains invalid values") for _name, (_data_type, key) in POINT_ATTRIBUTES.items(): if key in attributes and len(attributes[key]) != len(vertices): raise ValueError(f"Point attribute has the wrong length: {key}") @@ -209,17 +196,6 @@ def assign_mesh_attribute(mesh, name, data_type, domain, values): def assign_snapform_attributes(mesh, attributes, moge2=False): - if moge2: - assign_mesh_attribute(mesh, "o_depth", "FLOAT", "POINT", attributes["point_depth"]) - assign_mesh_attribute(mesh, "o_mask", "BOOLEAN", "FACE", attributes["face_mask"]) - assign_mesh_attribute( - mesh, - "o_depth_discontinuity", - "FLOAT", - "FACE", - attributes["face_depth_discontinuity"], - ) - return for name, (data_type, key) in POINT_ATTRIBUTES.items(): assign_mesh_attribute( mesh, @@ -229,6 +205,8 @@ def assign_snapform_attributes(mesh, attributes, moge2=False): attributes[key], ) for name, (data_type, key) in FACE_ATTRIBUTES.items(): + if key not in attributes: + continue assign_mesh_attribute( mesh, name, @@ -254,87 +232,42 @@ def assign_snapform_attributes(mesh, attributes, moge2=False): ) -def prepare_moge_surface_archive(results_path, output_path): - import numpy as np - - with np.load(results_path, allow_pickle=False) as archive: - required = {"depth", "mask", "normal", "points", "intrinsics"} - missing = required.difference(archive.files) - if missing: - raise ValueError(f"MoGe-2 surface data is missing: {', '.join(sorted(missing))}") - depth = np.asarray(archive["depth"], dtype=np.float32) - mask = np.asarray(archive["mask"], dtype=bool) - points = np.asarray(archive["points"], dtype=np.float32) - intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32) - normal = np.asarray(archive["normal"], dtype=np.float32) - if depth.shape[0] != 1 or points.shape[:3] != (*depth.shape,): - raise ValueError("Convert to Surface requires one matching MoGe-2 point grid") - raw_depth = depth[0] - finite = np.isfinite(raw_depth) & (raw_depth > 0.0) - finite &= np.isfinite(points[0]).all(axis=-1) - height, width = raw_depth.shape - indices = np.arange(height * width, dtype=np.int32).reshape(height, width) - quads = np.stack( - (indices[:-1, :-1], indices[1:, :-1], indices[1:, 1:], indices[:-1, 1:]), - axis=-1, - ).reshape(-1, 4) - faces = np.concatenate((quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), axis=0) - faces = faces[finite.reshape(-1)[faces].all(axis=1)] - if not len(faces): - raise ValueError("MoGe-2 output does not contain a valid surface") - flat_points = points[0].reshape(-1, 3) - face_mask = mask[0].reshape(-1)[faces].all(axis=1) - depth_edge_ratio = depth_discontinuity_ratio(raw_depth) - face_depth_discontinuity = depth_edge_ratio.reshape(-1)[faces].max(axis=1) - used = np.unique(faces) - index_map = np.full(height * width, -1, dtype=np.int32) - index_map[used] = np.arange(len(used), dtype=np.int32) - u = np.linspace(0.0, 1.0, width, dtype=np.float32) - v = np.linspace(1.0, 0.0, height, dtype=np.float32) - uv = np.stack(np.meshgrid(u, v), axis=-1).reshape(-1, 2)[used] - camera_intrinsics = intrinsics[0].copy() - principal_x = float(camera_intrinsics[0, 2]) - principal_y = float(camera_intrinsics[1, 2]) - if 0.0 <= principal_x <= 1.0 and 0.0 <= principal_y <= 1.0: - camera_intrinsics[0] *= width - camera_intrinsics[1] *= height - camera_intrinsics[2] = (0.0, 0.0, 1.0) - point_depth = raw_depth.reshape(-1)[used] - masked_depth = raw_depth[finite & mask[0]] - reference_depth = float(np.median(masked_depth if len(masked_depth) else point_depth)) - np.savez( - output_path, - vertices=flat_points[used], - faces=index_map[faces], - uv=uv, - intrinsics=camera_intrinsics, - image_size=np.array((width, height), dtype=np.int32), - reference_depth=np.float32(reference_depth), - point_depth=point_depth, - normal_map=normal[0], - face_mask=face_mask, - face_depth_discontinuity=face_depth_discontinuity, +def load_surface_normal_image(path, name): + image = bpy.data.images.load( + str(path), + check_existing=False, ) - return output_path - - -def depth_discontinuity_ratio(depth): - import numpy as np - - padded = np.pad(depth, 1, mode="edge") - windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3)) - difference = np.nanmax(windows, axis=(-2, -1)) - difference -= np.nanmin(windows, axis=(-2, -1)) - return difference / np.maximum(depth, 1e-6) + image.name = name + image.colorspace_settings.name = "Non-Color" + image.pack() + return image def node_asset_path(): return Path(__file__).parent.parent / "assets" / NODE_ASSET_NAME +def configure_surface_node_group_interface(node_group): + inputs = [ + item + for item in node_group.interface.items_tree + if item.item_type == "SOCKET" and item.in_out == "INPUT" + ] + iterations = next( + (item for item in inputs if item.name in {"Iterations", "Relax Iterations"}), + None, + ) + if iterations is None: + raise RuntimeError("Geometry Nodes input is missing: Iterations") + iterations.name = "Iterations" + iterations.min_value = 0 + node_group.interface.move(iterations, 2) + + def load_geometry_node_group(group_name=NODE_GROUP_NAME): existing = bpy.data.node_groups.get(group_name) if existing is not None: + configure_surface_node_group_interface(existing) return existing asset_path = node_asset_path() if not asset_path.is_file(): @@ -355,6 +288,7 @@ def load_geometry_node_group(group_name=NODE_GROUP_NAME): if node_group is None: raise RuntimeError("Appended Geometry Nodes asset is unavailable") node_group.use_fake_user = True + configure_surface_node_group_interface(node_group) return node_group with bpy.data.libraries.load(str(asset_path), link=use_pack) as ( source, @@ -372,6 +306,7 @@ def load_geometry_node_group(group_name=NODE_GROUP_NAME): if linked_group != node_group and linked_group.users == 0: bpy.data.node_groups.remove(linked_group) node_group.use_fake_user = True + configure_surface_node_group_interface(node_group) return node_group @@ -469,6 +404,8 @@ def build_image_mesh( create_camera, unshaded, moge2=False, + color_image=None, + normal_path=None, ): import numpy as np @@ -500,24 +437,17 @@ def build_image_mesh( context.collection.objects.link(mesh_object) mesh_object.matrix_world = source_object.matrix_world.copy() depth_scale = material_depth_scale(mesh_object, camera_parameters) - normal_image = None - if moge2 and attributes.get("normal_map", np.empty(0)).size: - from .depth_plane import create_moge_normal_image - - normal_map = attributes["normal_map"].copy() - normal_map[..., 0] /= camera_parameters["scale_x"] - normal_map[..., 1] /= -camera_parameters["scale_y"] - normal_map[..., 2] /= -camera_parameters["scale_z"] - lengths = np.linalg.norm(normal_map, axis=-1, keepdims=True) - normal_map /= np.maximum(lengths, 1e-6) - normal_image = create_moge_normal_image( + normal_image = ( + load_surface_normal_image( + normal_path, normal_data_name(source_object), - normal_map, - convert_coordinates=False, ) + if normal_path is not None + else None + ) material = create_material( source_name, - source_object.data, + color_image or source_object.data, depth_scale, unshaded=unshaded, surface_group_name=( diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 3bd4c0b..08d3e83 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -262,12 +262,25 @@ def build_surface_response( unshaded, moge2, ): - from .blender.depth_plane import is_image_empty + from .blender.depth_plane import ( + create_removed_background_image, + is_image_empty, + ) from .blender.surface import build_image_mesh source_object = bpy.data.objects.get(source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source Image Empty no longer exists") + foreground_path = result.directory / "foreground.png" + color_image = ( + create_removed_background_image( + source_object, + foreground_path, + assign_to_source=False, + ) + if foreground_path.is_file() + else None + ) build_image_mesh( context, source_object, @@ -276,6 +289,12 @@ def build_surface_response( create_camera, unshaded, moge2, + color_image=color_image, + normal_path=( + result.directory / "normal.png" + if moge2 + else None + ), ) push_undo("Convert to Surface") @@ -303,10 +322,6 @@ def da3_parameters(settings, path, model, process_res, start_frame): "model": model, "process_res": process_res, "process_res_method": "upper_bound_resize", - "conf_thresh_percentile": settings.conf_thresh_percentile, - "num_max_points": settings.num_max_points, - "show_cameras": 1 if settings.show_cameras else 0, - "gs_views_interval": settings.gs_views_interval, "video_frames": settings.video_frames, "start_frame": start_frame, "device": production_device().lower(), @@ -359,16 +374,13 @@ class DebugDa3(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_da3" bl_label = "Debug DA3" job_type = "debug-da3" + starting_message = "Starting DA3 API inference" def request(self, context): require_environment() settings = context.scene.snapform_settings - model = settings.gs_model if settings.export_format == "gs_ply" else settings.model - spec = require_model(model) - if settings.export_format == "gs_ply" and not spec.get("supports_gs"): - raise RuntimeError("PLY export requires a DA3 model with a Gaussian head") - self.export_format = settings.export_format - self.starting_message = f"Starting {self.export_format} export" + model = settings.model + require_model(model) return { **da3_parameters( settings, @@ -378,16 +390,13 @@ class DebugDa3(JobOperatorBase, bpy.types.Operator): 1, ), "device": settings.device.lower(), - "export_format": settings.export_format, "inference_mode": "joint", "remove_background_model": "", - "mesh_filter_edges": settings.mesh_filter_edges, - "mesh_max_edge_length": settings.mesh_max_edge_length, } def response(self, _context, result): open_job_directory(result.directory) - return f"Exported {self.export_format} to {result.directory}" + return f"Saved DA3 API prediction to {result.directory}" class DebugMoge2(JobOperatorBase, bpy.types.Operator): @@ -409,8 +418,6 @@ class DebugMoge2(JobOperatorBase, bpy.types.Operator): 1, ), "device": settings.device.lower(), - "export_format": settings.moge_export_format, - "include_points": True, } def response(self, _context, result): @@ -444,7 +451,7 @@ class DebugBen2(JobOperatorBase, bpy.types.Operator): class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_da3_depth" bl_label = "Generate DA3 Depth" - job_type = "generate-da3-depth" + job_type = "generate-depth" starting_message = "Starting depth estimation" input_path: bpy.props.StringProperty(options=HIDDEN) @@ -471,7 +478,6 @@ class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): self.process_res, source_start_frame(self.source_object_name), ), - "export_format": "mini_npz", "inference_mode": "per_frame", "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" @@ -497,7 +503,7 @@ class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_moge2_depth" bl_label = "Generate MoGe-2 Depth" - job_type = "generate-moge2-depth" + job_type = "generate-depth" starting_message = "Starting MoGe-2 inference" input_path: bpy.props.StringProperty(options=HIDDEN) @@ -546,7 +552,7 @@ class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator): class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_da3_surface" bl_label = "Generate DA3 Surface" - job_type = "generate-da3-surface" + job_type = "generate-surface" starting_message = "Starting perspective surface generation" input_path: bpy.props.StringProperty(options=HIDDEN) @@ -573,13 +579,9 @@ class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): self.process_res, source_start_frame(self.source_object_name), ), - "export_format": "mesh_data", - "inference_mode": "per_frame", "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" ), - "mesh_filter_edges": settings.mesh_filter_edges, - "mesh_max_edge_length": settings.mesh_max_edge_length, } def response(self, context, result): @@ -601,7 +603,7 @@ class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_moge2_surface" bl_label = "Generate MoGe-2 Surface" - job_type = "generate-moge2-surface" + job_type = "generate-surface" starting_message = "Starting MoGe-2 inference" input_path: bpy.props.StringProperty(options=HIDDEN) @@ -631,7 +633,6 @@ class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator): "remove_background_model": ( BEN2_MODEL_KEY if self.remove_background else "" ), - "surface_mesh": True, } def response(self, context, result): diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 54b21da..4d1adf3 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -3,11 +3,6 @@ import bpy from .job_runtime import runtime from .properties import BEN2_MODEL_KEY, model_catalog, model_ready - -def selected_model_key(settings): - return settings.gs_model if settings.export_format == "gs_ply" else settings.model - - class ServerPanel(bpy.types.Panel): bl_label = "Job Server" bl_idname = "SNAPFORM_PT_server" @@ -215,36 +210,15 @@ class Da3OutputDebugPanel(bpy.types.Panel): layout = self.layout settings = context.scene.snapform_settings - layout.prop(settings, "export_format") - model_property = ( - "gs_model" if settings.export_format == "gs_ply" else "model" - ) - draw_da3_input_settings(layout, settings, model_property) - - if settings.export_format == "glb": - glb_box = layout.box() - glb_box.label(text="GLB Settings") - glb_box.prop(settings, "conf_thresh_percentile") - glb_box.prop(settings, "num_max_points") - glb_box.prop(settings, "show_cameras") - elif settings.export_format == "gs_ply": - ply_box = layout.box() - ply_box.label(text="PLY (3DGS) Settings") - ply_box.prop(settings, "gs_views_interval") - elif settings.export_format == "mesh": - mesh_box = layout.box() - mesh_box.label(text="Surface Settings") - mesh_box.prop(settings, "conf_thresh_percentile") - mesh_box.prop(settings, "mesh_filter_edges") - mesh_box.prop(settings, "mesh_max_edge_length") + draw_da3_input_settings(layout, settings, "model") export_row = layout.row() export_row.enabled = ( not runtime.server_busy() - and model_ready(selected_model_key(settings)) + and model_ready(settings.model) ) export_row.operator( "snapform.debug_da3", - text="Generate Output", + text="Run DA3 API", icon="EXPORT", ) @@ -262,7 +236,6 @@ class Moge2ApiDebugPanel(bpy.types.Panel): def draw(self, context): layout = self.layout settings = context.scene.snapform_settings - layout.prop(settings, "moge_export_format") layout.prop(settings, "input_path") layout.prop(settings, "moge_model") layout.prop(settings, "moge_resolution_level") @@ -273,7 +246,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel): ) run_row.operator( "snapform.debug_moge2", - text="Generate MoGe-2 Output", + text="Run MoGe-2 API", icon="EXPORT", ) diff --git a/src/snapform/properties.py b/src/snapform/properties.py index df9d510..c9980ae 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -106,16 +106,6 @@ def model_enum_items(_owner, _context): ) -def gs_model_enum_items(_owner, _context): - return _model_items( - model - for model in model_catalog() - if model.get("family") == "da3" - and model.get("enabled") - and model.get("supports_gs") - ) - - def moge2_model_enum_items(_owner, _context): return _model_items( model for model in model_catalog() if model.get("family") == "moge2" @@ -170,11 +160,6 @@ class SnapFormSettings(bpy.types.PropertyGroup): description="Depth Anything 3 model", items=model_enum_items, ) - gs_model: bpy.props.EnumProperty( - name="Model", - description="Depth Anything 3 model with a Gaussian head", - items=gs_model_enum_items, - ) device: bpy.props.EnumProperty( name="Debug Device", description=( @@ -221,36 +206,6 @@ class SnapFormSettings(bpy.types.PropertyGroup): items=moge2_model_enum_items, default=0, ) - moge_export_format: bpy.props.EnumProperty( - name="Output Format", - description="Additional MoGe-2 debug output generated with results.npz", - items=( - ("maps", "Maps", "Export depth, normal, and mask previews"), - ("glb", "Textured Surface (GLB)", "Export the official-style textured mesh"), - ("ply", "Point Cloud (PLY)", "Export the official-style colored point cloud"), - ), - default="maps", - ) - conf_thresh_percentile: bpy.props.FloatProperty( - name="Confidence Threshold", - description="Remove the lowest-confidence percentile of points or surface faces", - default=40.0, - min=0.0, - max=100.0, - subtype="PERCENTAGE", - ) - num_max_points: bpy.props.IntProperty( - name="Num Max Points", - description="DA3 GLB maximum point count", - default=1_000_000, - min=1, - soft_max=5_000_000, - ) - show_cameras: bpy.props.BoolProperty( - name="Show Cameras", - description="Include DA3 camera wireframes in the GLB export", - default=True, - ) ref_view_strategy: bpy.props.EnumProperty( name="Ref View Strategy", description="DA3 reference view selection strategy", @@ -270,56 +225,3 @@ class SnapFormSettings(bpy.types.PropertyGroup): ), default="saddle_balanced", ) - export_format: bpy.props.EnumProperty( - name="Output Type", - description="Result generated by the DA3 debug task", - items=( - ( - "glb", - "Point Cloud (GLB)", - "Generate a colored point cloud and optional cameras", - ), - ( - "gs_ply", - "Gaussian Splats (PLY)", - "Generate a 3D Gaussian Splatting PLY", - ), - ( - "depth_vis", - "Depth Preview (Images)", - "Generate image and colored-depth preview pairs", - ), - ( - "mini_npz", - "Depth Data (NPZ)", - "Generate compact depth, confidence, and camera data", - ), - ( - "mesh", - "Textured Surface (GLB)", - "Generate one filtered textured depth surface per input image", - ), - ), - default="glb", - ) - mesh_filter_edges: bpy.props.BoolProperty( - name="Filter Depth Edges", - description="Remove faces around sharp depth transitions", - default=True, - ) - mesh_max_edge_length: bpy.props.FloatProperty( - name="Max Edge Length", - description=( - "Remove faces with an edge longer than this scene distance; " - "zero disables the test" - ), - default=0.1, - min=0.0, - soft_max=2.0, - ) - gs_views_interval: bpy.props.IntProperty( - name="GS Views Interval", - description="Export 3D Gaussians every N input views", - default=1, - min=1, - ) diff --git a/src/snapform/server/__init__.py b/src/snapform/server/__init__.py index cc84811..08a03e8 100644 --- a/src/snapform/server/__init__.py +++ b/src/snapform/server/__init__.py @@ -2,11 +2,14 @@ from blendjob import JobServer +from .jobs import generate_depth as _generate_depth +from .jobs import generate_surface as _generate_surface +from .jobs.background import run as _remove_background +from .jobs.debug_ben2 import run as _debug_ben2 +from .jobs.debug_da3 import run as _debug_da3 +from .jobs.debug_moge2 import run as _debug_moge2 +from .jobs.frames import run as _extract_frames from .model_manager import ModelManager, ModelSessionCache -from .background import remove_background as _remove_background -from .da3 import generate_da3 as _generate_da3 -from .media import extract_frames as _extract_frames -from .moge2 import generate_moge2 as _generate_moge2 server = JobServer("SnapForm Job Server") @@ -29,7 +32,6 @@ def download_required_models(context, parameters): return manager.download_missing(context, parameters["models"]) -@server.job("debug-ben2") @server.job("remove-background") def remove_background(context, parameters): return _remove_background(context, parameters) @@ -40,18 +42,29 @@ def extract_frames(context, parameters): return _extract_frames(context, parameters) -@server.job("debug-moge2") -@server.job("generate-moge2-surface") -@server.job("generate-moge2-depth") -def generate_moge2(context, parameters): - return _generate_moge2(context, parameters) +@server.job("generate-depth") +def generate_depth(context, parameters): + return _generate_depth(context, parameters) + + +@server.job("generate-surface") +def generate_surface(context, parameters): + return _generate_surface(context, parameters) @server.job("debug-da3") -@server.job("generate-da3-surface") -@server.job("generate-da3-depth") -def generate_da3(context, parameters): - return _generate_da3(context, parameters) +def debug_da3(context, parameters): + return _debug_da3(context, parameters) + + +@server.job("debug-moge2") +def debug_moge2(context, parameters): + return _debug_moge2(context, parameters) + + +@server.job("debug-ben2") +def debug_ben2(context, parameters): + return _debug_ben2(context, parameters) __all__ = ("ModelManager", "ModelSessionCache", "server") diff --git a/src/snapform/server/background.py b/src/snapform/server/background.py deleted file mode 100644 index ad67270..0000000 --- a/src/snapform/server/background.py +++ /dev/null @@ -1,170 +0,0 @@ -import shutil -import time -from pathlib import Path - -from .media import ( - IMAGE_SUFFIXES, - check_cancelled, - copy_frames, - prepare_background_input, -) - - -def infer_foreground( - input_path, - model_dir, - device, - model_cache=None, - cancel_check=None, -): - from PIL import Image - - input_path = Path(input_path) - check_cancelled(cancel_check) - if input_path.suffix.lower() not in IMAGE_SUFFIXES: - raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") - from .onnx_ben2 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 = 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) - - check_cancelled(cancel_check) - with Image.open(input_path) as opened_image: - image = opened_image.convert("RGB") - inference_started = time.perf_counter() - foreground = infer(model, image) - foreground = preserve_source_rgb(image, foreground) - inference_ms = (time.perf_counter() - inference_started) * 1000.0 - check_cancelled(cancel_check) - return foreground, load_ms, inference_ms - - -def preserve_source_rgb(source, foreground): - """Combine the untouched source RGB with BEN2's foreground alpha.""" - source_rgba = source.convert("RGBA") - alpha = foreground.convert("RGBA").getchannel("A") - if alpha.size != source_rgba.size: - from PIL import Image - - alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS) - source_rgba.putalpha(alpha) - return source_rgba - - -def infer_foreground_frames( - images, - model_dir, - device, - output_directory, - progress_callback, - start_progress, - progress_span, - model_cache=None, - cancel_check=None, -): - """Run BEN2 on every frame and write foreground.0001.png etc.""" - total_load_ms = 0.0 - total_inference_ms = 0.0 - for index, image_path in enumerate(images): - check_cancelled(cancel_check) - progress_callback( - start_progress + progress_span * index / len(images), - f"Removing background ({index + 1}/{len(images)})", - ) - foreground, load_ms, inference_ms = infer_foreground( - image_path, - model_dir, - device, - model_cache=model_cache, - cancel_check=cancel_check, - ) - output_path = output_directory / f"foreground.{index + 1:04d}.png" - foreground.convert("RGBA").save( - output_path, - format="PNG", - compress_level=1, - ) - if not output_path.is_file(): - raise RuntimeError(f"BEN2 did not write: {output_path}") - total_load_ms += load_ms - total_inference_ms += inference_ms - return total_load_ms, total_inference_ms - - -def remove_background(context, parameters): - """Run BEN2 and write transparent foreground images into the Job directory.""" - total_started = time.perf_counter() - context.progress(0.02, "Starting Environment") - context.check_cancelled() - - output_directory = context.directory - output_directory.mkdir(parents=True, exist_ok=True) - device = parameters["device"] - model_manager = context.resource("model_manager") - model_directory = model_manager.directory(parameters["model"]) - - video_frames = parameters.get("video_frames", 16) - start_frame = parameters.get("start_frame", 1) - images, temporary_directory = prepare_background_input( - parameters["input"], - video_frames, - output_directory, - start_frame=start_frame, - cancel_check=context.check_cancelled, - ) - try: - context.progress(0.12, f"Loading BEN2 for {device.upper()}") - if len(images) == 1: - context.progress(0.35, "Removing background") - foreground, load_ms, inference_ms = infer_foreground( - images[0], - model_directory, - device, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - context.check_cancelled() - context.progress(0.92, "Writing transparent PNG") - output_path = output_directory / "foreground.png" - foreground.convert("RGBA").save( - output_path, - format="PNG", - compress_level=1, - ) - if not output_path.is_file(): - raise RuntimeError("BEN2 did not write a foreground image") - else: - load_ms, inference_ms = infer_foreground_frames( - images, - model_directory, - device, - output_directory, - context.progress, - 0.12, - 0.7, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - copy_frames(images, output_directory) - finally: - if temporary_directory is not None: - shutil.rmtree(temporary_directory, ignore_errors=True) - - context.check_cancelled() - total_ms = (time.perf_counter() - total_started) * 1000.0 - frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)" - message = ( - f"Background removed{frame_text} in {inference_ms:.0f} ms " - f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" - ) - print(message) - context.progress(1.0, message) - if tuple(output_directory.glob("foreground.*.png")): - return {"foreground": "."} - return {"foreground": "foreground.png"} diff --git a/src/snapform/server/da3.py b/src/snapform/server/da3.py deleted file mode 100644 index d112d2d..0000000 --- a/src/snapform/server/da3.py +++ /dev/null @@ -1,232 +0,0 @@ -import shutil - -from .background import infer_foreground, infer_foreground_frames -from .geometry import ( - export_da3_depth_preview, - export_da3_mini_npz, - export_da3_point_cloud, - export_image_mesh_data, - export_prediction_attributes, - export_textured_mesh_glb, -) -from .media import ( - copy_frames, - flatten_exported_paths, - prepare_input, - wait_for_exported_paths, -) - - -def merge_da3_frame_predictions(predictions): - """Combine independent single-frame DA3 results for existing exporters.""" - import numpy as np - - if not predictions: - raise ValueError("DA3 returned no frame predictions") - merged = predictions[0] - for name in ( - "depth", - "conf", - "extrinsics", - "intrinsics", - "processed_images", - ): - values = [getattr(prediction, name) for prediction in predictions] - setattr(merged, name, np.concatenate(values, axis=0)) - return merged - - -def generate_da3(context, parameters): - """Run DA3 inference and return the artifact used by the requesting Job.""" - context.progress(0.02, "Starting Environment") - context.check_cancelled() - from .onnx_da3 import infer - from .onnx_runtime import provider_label - - output_directory = context.directory - output_directory.mkdir(parents=True, exist_ok=True) - context.progress( 0.08, "Preparing input") - model_manager = context.resource("model_manager") - model_directory = model_manager.directory(parameters["model"]) - device = parameters["device"] - export_format = parameters["export_format"] - images, temporary_directory = prepare_input( - parameters["input"], - parameters.get("video_frames", 16), - output_directory, - start_frame=parameters.get("start_frame", 1), - cancel_check=context.check_cancelled, - ) - - try: - context.check_cancelled() - remove_background_model = parameters.get("remove_background_model", "") - remove_background_model_dir = ( - model_manager.directory(remove_background_model) - if remove_background_model - else "" - ) - foreground_images = [] - if remove_background_model_dir: - context.progress( 0.12, "Removing image background") - if len(images) == 1: - foreground, load_ms, inference_ms = infer_foreground( - images[0], - remove_background_model_dir, - device, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - context.check_cancelled() - foreground.convert("RGBA").save( - output_directory / "foreground.png", - format="PNG", - compress_level=1, - ) - if not (output_directory / "foreground.png").is_file(): - raise RuntimeError("BEN2 did not write a foreground image") - foreground_images = [output_directory / "foreground.png"] - else: - load_ms, inference_ms = infer_foreground_frames( - images, - remove_background_model_dir, - device, - output_directory, - context.progress, - 0.12, - 0.05, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - foreground_images = sorted( - output_directory.glob("foreground.*.png") - ) - if remove_background_model_dir: - print( - f"BEN2 inference {inference_ms:.0f} ms; " - f"model load {load_ms:.0f} ms" - ) - context.progress( 0.15, "Loading DA3 ONNX model") - model = model_manager.get_da3(model_directory, device) - context.progress( - 0.20, - f"Loaded DA3 on {provider_label(model)}", - ) - - context.check_cancelled() - context.progress( - 0.25, - f"Estimating depth for {len(images)} frame(s)", - ) - inference_options = { - "process_res": parameters["process_res"], - "process_res_method": parameters["process_res_method"], - "ref_view_strategy": parameters["ref_view_strategy"], - } - try: - if parameters.get("inference_mode", "joint") == "per_frame": - predictions = [] - for index, image in enumerate(images): - context.check_cancelled() - context.progress( - 0.25 + 0.60 * index / len(images), - f"Estimating DA3 depth ({index + 1}/{len(images)})", - ) - predictions.append( - infer(model, [str(image)], **inference_options) - ) - prediction = merge_da3_frame_predictions(predictions) - else: - prediction = infer( - model, - [str(path) for path in images], - **inference_options, - ) - except UnicodeDecodeError as error: - # DirectML reports failures with a locale-specific message that - # ONNX Runtime cannot always decode, masking the real error. - if provider_label(model) == "DirectML": - raise RuntimeError( - "DA3 inference failed on DirectML. Switch the device to " - "CPU in SnapForm Preferences and retry. " - f"Original error: {error}" - ) from error - raise - context.check_cancelled() - context.progress( 0.88, f"Exporting {export_format}") - if export_format == "mesh": - export_textured_mesh_glb( - prediction, - images, - output_directory, - parameters.get("conf_thresh_percentile", 40.0), - bool(parameters.get("mesh_filter_edges", True)), - float(parameters.get("mesh_max_edge_length", 0.1)), - ) - elif export_format == "mesh_data": - def mesh_progress(step, total, label): - context.progress( - 0.88 + 0.09 * step / total, - f"Preparing surface data {step}/{total}: {label}", - ) - - export_image_mesh_data( - prediction, - images[0], - output_directory, - foreground_path=( - foreground_images[0] if foreground_images else None - ), - progress_callback=mesh_progress, - ) - elif export_format == "mini_npz": - export_da3_mini_npz(prediction, output_directory) - elif export_format == "depth_vis": - export_da3_depth_preview(prediction, output_directory) - elif export_format == "glb": - export_da3_point_cloud( - prediction, - output_directory, - parameters["conf_thresh_percentile"], - parameters["num_max_points"], - ) - else: - raise ValueError( - f"DA3 ONNX does not support export format: {export_format}" - ) - if export_format in {"glb", "gs_ply", "depth_vis", "mesh"}: - context.progress( - 0.97, - "Writing structured debug attributes", - ) - export_prediction_attributes(prediction, output_directory) - context.check_cancelled() - - result_paths = wait_for_exported_paths( - output_directory, - export_format, - cancel_check=context.check_cancelled, - ) - if not result_paths: - raise RuntimeError( - f"No {export_format} files were exported" - ) - result_paths = flatten_exported_paths( - output_directory, - result_paths, - cancel_check=context.check_cancelled, - ) - context.check_cancelled() - copy_frames(images, output_directory) - context.progress( - 1.0, - f"Exported {len(result_paths)} {export_format} file(s)", - ) - if context.job_type == "generate-da3-surface": - return {"surface": "mesh.npz"} - if context.job_type == "generate-da3-depth": - return {"depth": "results.npz"} - return {"directory": "."} - finally: - if temporary_directory is not None: - shutil.rmtree(temporary_directory, ignore_errors=True) diff --git a/src/snapform/server/geometry.py b/src/snapform/server/geometry.py deleted file mode 100644 index b0fdbbb..0000000 --- a/src/snapform/server/geometry.py +++ /dev/null @@ -1,616 +0,0 @@ -from pathlib import Path - - -SCENE_FILENAME = "scene.glb" - - -def unproject_depth_frame(depth, intrinsics, extrinsics): - import numpy as np - - height, width = depth.shape - x, y = np.meshgrid( - np.arange(width, dtype=np.float32), - np.arange(height, dtype=np.float32), - ) - pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) - rays = (np.linalg.inv(intrinsics) @ pixels.T).T - camera_points = rays * depth.reshape(-1, 1) - camera_points = np.column_stack( - (camera_points, np.ones(camera_points.shape[0], dtype=np.float32)) - ) - world_to_camera = np.eye(4, dtype=np.float32) - world_to_camera[: extrinsics.shape[0], :4] = extrinsics - world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T - return world_points[:, :3].reshape(height, width, 3) - - -def unproject_camera_depth(depth, intrinsics): - import numpy as np - - height, width = depth.shape - x, y = np.meshgrid( - np.arange(width, dtype=np.float32), - np.arange(height, dtype=np.float32), - ) - pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) - rays = (np.linalg.inv(intrinsics) @ pixels.T).T - return (rays * depth.reshape(-1, 1)).reshape(height, width, 3) - - -def compact_mesh(vertices, faces): - import numpy as np - - vertices = np.asarray(vertices, dtype=np.float32) - faces = np.asarray(faces, dtype=np.int32) - used_vertices = np.unique(faces) - index_map = np.full(len(vertices), -1, dtype=np.int32) - index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32) - vertices = vertices[used_vertices] - faces = index_map[faces] - return vertices, faces - - -def project_camera_uv(vertices, intrinsics, image_size): - import numpy as np - - width, height = image_size - z = vertices[:, 2] - if np.any(~np.isfinite(z)) or np.any(z <= 0.0): - raise ValueError("Simplified mesh contains invalid camera depth") - u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2] - v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2] - denominator_x = max(width - 1, 1) - denominator_y = max(height - 1, 1) - return np.column_stack( - ( - np.clip(u / denominator_x, 0.0, 1.0), - np.clip(1.0 - v / denominator_y, 0.0, 1.0), - ) - ).astype(np.float32) - - -def project_camera_pixels(vertices, intrinsics): - import numpy as np - - z = vertices[:, 2] - if np.any(~np.isfinite(z)) or np.any(z <= 0.0): - raise ValueError("Mesh contains invalid camera depth") - return np.column_stack( - ( - intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2], - intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2], - ) - ) - - -def sample_image_nearest(values, pixels): - import numpy as np - - height, width = values.shape - x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1) - y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1) - return np.asarray(values)[y, x] - - -def confidence_percentile_map(depth, confidence): - import numpy as np - - valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) - result = np.zeros(depth.shape, dtype=np.float32) - if not valid.any(): - return result - sorted_values = np.sort(np.asarray(confidence)[valid]) - ranks = np.searchsorted( - sorted_values, - np.asarray(confidence)[valid], - side="right", - ) - result[valid] = ranks.astype(np.float32) / len(sorted_values) - return result - - -def depth_discontinuity_ratio(depth): - import numpy as np - - padded = np.pad(depth, 1, mode="edge") - windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3)) - difference = np.nanmax(windows, axis=(-2, -1)) - difference -= np.nanmin(windows, axis=(-2, -1)) - return difference / np.maximum(depth, 1e-6) - - -def triangle_edge_stretch_values(points, triangles): - import numpy as np - - flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3) - triangle_points = flat_points[triangles] - triangle_depth = triangle_points[..., 2] - safe_depth = np.where( - np.abs(triangle_depth) > 1e-8, - triangle_depth, - np.nan, - ) - triangle_rays = triangle_points / safe_depth[..., None] - next_points = np.roll(triangle_points, -1, axis=1) - next_rays = np.roll(triangle_rays, -1, axis=1) - next_depth = np.roll(triangle_depth, -1, axis=1) - local_depth = (triangle_depth + next_depth) * 0.5 - flat_edges = ( - triangle_rays * local_depth[..., None] - - next_rays * local_depth[..., None] - ) - flat_lengths = np.linalg.norm(flat_edges, axis=-1) - actual_lengths = np.linalg.norm( - triangle_points - next_points, - axis=-1, - ) - stretch = actual_lengths / np.maximum(flat_lengths, 1e-8) - stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0 - return stretch - - -def image_mesh_attributes( - vertices, - faces, - intrinsics, - confidence_percentiles, - face_depth_discontinuity, -): - import numpy as np - - point_percentile = sample_image_nearest( - confidence_percentiles, - project_camera_pixels(vertices, intrinsics), - ).astype(np.float32) - face_confidence = point_percentile[faces].min(axis=1) - return { - "point_depth": vertices[:, 2].astype(np.float32), - "face_confidence_percentile": face_confidence.astype(np.float32), - "face_depth_discontinuity": np.asarray( - face_depth_discontinuity, - dtype=np.float32, - ), - } - - -def export_prediction_attributes(prediction, output_directory): - import numpy as np - - values = {} - for name in ("depth", "conf", "intrinsics", "extrinsics"): - value = getattr(prediction, name, None) - if value is not None: - values[name] = np.asarray(value) - if "depth" in values and "conf" in values: - values["confidence_percentile"] = np.stack( - [ - confidence_percentile_map(depth, confidence) - for depth, confidence in zip( - values["depth"], - values["conf"], - ) - ] - ) - output_path = Path(output_directory) / "prediction.attributes.npz" - np.savez(output_path, **values) - 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) - return faces[foreground_points[faces].all(axis=1)] - - -def export_image_mesh_data( - prediction, - image_path, - output_directory, - foreground_path=None, - progress_callback=None, -): - import numpy as np - from PIL import Image - - total_steps = 6 - - def progress(step, label): - if progress_callback is not None: - progress_callback(step, total_steps, label) - - progress(1, "Validating DA3 data") - required = ("depth", "conf", "intrinsics") - missing = [ - name - for name in required - if getattr(prediction, name, None) is None - ] - if missing: - raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") - depth_frames = np.asarray(prediction.depth, dtype=np.float32) - confidence_frames = np.asarray(prediction.conf, dtype=np.float32) - if depth_frames.shape[0] != 1: - raise ValueError("Convert to Surface currently supports one image") - raw_depth = depth_frames[0] - confidence = confidence_frames[0] - intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32) - progress(2, "Unprojecting the depth grid") - points = unproject_camera_depth(raw_depth, intrinsics) - progress(3, "Preparing confidence and background data") - foreground_mask = np.ones(raw_depth.shape, dtype=bool) - if foreground_path is not None: - foreground_mask = foreground_mesh_mask( - foreground_path, - raw_depth.shape, - ) - confidence_percentiles = confidence_percentile_map( - raw_depth, - confidence, - ) - progress(4, "Building source triangles") - faces = mesh_triangle_faces( - points, - raw_depth, - confidence, - None, - False, - 0.0, - ) - if foreground_path is not None: - faces = remove_background_faces(faces, foreground_mask) - if not len(faces): - raise ValueError( - "DA3 output does not contain any valid surface faces" - ) - progress(5, "Preparing source UV and attributes") - depth_discontinuity = depth_discontinuity_ratio(raw_depth) - face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1) - vertices, faces = compact_mesh( - points.reshape(-1, 3), - faces, - ) - height, width = raw_depth.shape - uv = project_camera_uv(vertices, intrinsics, (width, height)) - attributes = image_mesh_attributes( - vertices, - faces, - intrinsics, - confidence_percentiles, - face_depth_discontinuity, - ) - reference_depth = float(np.median(vertices[:, 2])) - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - progress(6, "Writing source surface data") - result_path = output_directory / "mesh.npz" - np.savez_compressed( - result_path, - vertices=vertices, - faces=faces, - uv=uv, - intrinsics=intrinsics, - image_size=np.asarray((width, height), dtype=np.int32), - reference_depth=np.asarray(reference_depth, dtype=np.float32), - **attributes, - ) - texture_path = output_directory / "texture.png" - texture_source = foreground_path or image_path - with Image.open(texture_source) as image: - image.convert("RGBA").save(texture_path, format="PNG", compress_level=1) - return result_path, texture_path - - -def depth_edge_mask(depth, threshold=12.0 / 255.0): - import numpy as np - - values = np.asarray(depth, dtype=np.float32) - edges = np.zeros(values.shape, dtype=bool) - for axis in (0, 1): - first = np.take(values, range(values.shape[axis] - 1), axis=axis) - second = np.take(values, range(1, values.shape[axis]), axis=axis) - valid = ( - np.isfinite(first) - & np.isfinite(second) - & (first > 0.0) - & (second > 0.0) - ) - scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) - discontinuity = valid & (np.abs(first - second) / scale > threshold) - lower = [slice(None), slice(None)] - upper = [slice(None), slice(None)] - lower[axis] = slice(0, -1) - upper[axis] = slice(1, None) - edges[tuple(lower)] |= discontinuity - edges[tuple(upper)] |= discontinuity - return edges - - -def edge_culling_stretch_limit(culling): - factor = min(max(float(culling) / 100.0, 0.0001), 1.0) - return 1.0 / factor - - -def foreground_mesh_mask(foreground_path, depth_shape): - import numpy as np - from PIL import Image, ImageFilter - - height, width = depth_shape - with Image.open(foreground_path) as source: - alpha = source.convert("RGBA").getchannel("A") - alpha = alpha.resize( - (width, height), - resample=Image.Resampling.BILINEAR, - ) - alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75)) - return np.asarray(alpha) >= 128 - - -def confidence_percentile_threshold(depth, confidence, percentile): - import numpy as np - - if float(percentile) <= 0.0: - return None - valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) - if not valid.any(): - raise ValueError("DA3 returned no valid depth confidence values") - return float(np.percentile(confidence[valid], percentile)) - - -def mesh_triangle_faces( - points, - depth, - confidence, - confidence_threshold, - filter_edges, - max_edge_length, - edge_threshold=12.0 / 255.0, - edge_stretch_limit=None, - vertex_mask=None, -): - import numpy as np - - height, width = depth.shape - indices = np.arange(height * width, dtype=np.int32).reshape(height, width) - quads = np.stack( - ( - indices[:-1, :-1], - indices[1:, :-1], - indices[1:, 1:], - indices[:-1, 1:], - ), - axis=-1, - ).reshape(-1, 4) - valid = ( - np.isfinite(points).all(axis=-1) - & np.isfinite(depth) - & (depth > 0.0) - ) - if confidence_threshold is not None: - valid &= np.isfinite(confidence) & ( - confidence >= confidence_threshold - ) - if filter_edges: - valid &= ~depth_edge_mask(depth, edge_threshold) - if vertex_mask is not None: - if np.shape(vertex_mask) != depth.shape: - raise ValueError("Foreground mask and depth dimensions do not match") - valid &= np.asarray(vertex_mask, dtype=bool) - triangles = np.concatenate( - (quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), - axis=0, - ) - face_valid = valid.reshape(-1)[triangles].all(axis=1) - if edge_stretch_limit is not None: - face_valid &= triangle_edge_stretch_mask( - points, - triangles, - edge_stretch_limit, - ) - if max_edge_length > 0.0: - flat_points = points.reshape(-1, 3) - triangle_points = flat_points[triangles] - next_points = np.roll(triangle_points, -1, axis=1) - edge_lengths = np.linalg.norm( - triangle_points - next_points, - axis=-1, - ) - face_valid &= (edge_lengths <= max_edge_length).all(axis=1) - triangles = triangles[face_valid] - if not len(triangles): - return np.empty((0, 3), dtype=np.int32) - return triangles - - -def triangle_edge_stretch_mask(points, triangles, stretch_limit): - stretch = triangle_edge_stretch_values(points, triangles) - tolerance = float(stretch_limit) * (1.0 + 1e-5) - return (stretch <= tolerance).all(axis=1) - - -def export_textured_mesh_glb( - prediction, - image_paths, - output_directory, - confidence_percentile, - filter_edges, - max_edge_length, -): - import numpy as np - import trimesh - from PIL import Image - - required = ("depth", "conf", "intrinsics", "extrinsics") - missing = [name for name in required if getattr(prediction, name) is None] - if missing: - raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") - depth = np.asarray(prediction.depth, dtype=np.float32) - confidence = np.asarray(prediction.conf, dtype=np.float32) - threshold = confidence_percentile_threshold( - depth, - confidence, - confidence_percentile, - ) - meshes = [] - referenced_points = [] - for index in range(depth.shape[0]): - points = unproject_depth_frame( - depth[index], - prediction.intrinsics[index], - prediction.extrinsics[index], - ) - faces = mesh_triangle_faces( - points, - depth[index], - confidence[index], - threshold, - filter_edges, - max_edge_length, - ) - if not len(faces): - continue - height, width = depth[index].shape - 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) - uv = np.stack((uu, vv), axis=-1).reshape(-1, 2) - with Image.open(image_paths[index]) as source_image: - texture_image = source_image.convert("RGB").copy() - visual = trimesh.visual.texture.TextureVisuals( - uv=uv, - image=texture_image, - ) - mesh = trimesh.Trimesh( - vertices=points.reshape(-1, 3), - faces=faces, - visual=visual, - process=False, - ) - mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1) - mesh.vertex_attributes["o_confidence"] = confidence[ - index - ].reshape(-1) - mesh.vertex_attributes["o_confidence_percentile"] = ( - confidence_percentile_map( - depth[index], - confidence[index], - ).reshape(-1) - ) - mesh.remove_unreferenced_vertices() - meshes.append(mesh) - referenced_points.append(mesh.vertices) - if not meshes: - raise ValueError( - "Surface filters removed every face; lower Confidence Threshold or disable filters" - ) - - world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64) - if world_to_camera.shape == (3, 4): - homogeneous = np.eye(4, dtype=np.float64) - homogeneous[:3, :4] = world_to_camera - world_to_camera = homogeneous - cv_to_gltf = np.eye(4, dtype=np.float64) - cv_to_gltf[1, 1] = -1.0 - cv_to_gltf[2, 2] = -1.0 - alignment = cv_to_gltf @ world_to_camera - aligned_points = trimesh.transform_points( - np.concatenate(referenced_points, axis=0), - alignment, - ) - center = np.median(aligned_points, axis=0) - alignment[:3, 3] -= center - - scene = trimesh.Scene() - for index, mesh in enumerate(meshes): - mesh.apply_transform(alignment) - scene.add_geometry(mesh, node_name=f"DepthMesh.{index + 1:04d}") - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - output_path = output_directory / "mesh.glb" - scene.export(output_path) - return output_path diff --git a/src/snapform/server/geometry/__init__.py b/src/snapform/server/geometry/__init__.py new file mode 100644 index 0000000..c063b52 --- /dev/null +++ b/src/snapform/server/geometry/__init__.py @@ -0,0 +1,35 @@ +from .grid_mesh import ( + compact_mesh, + confidence_percentile_map, + confidence_percentile_threshold, + depth_edge_mask, + depth_discontinuity_ratio, + edge_culling_stretch_limit, + foreground_mesh_mask, + mesh_triangle_faces, + remove_background_faces, + triangle_edge_stretch_values, + triangle_edge_stretch_mask, +) +from .projection import unproject_camera_depth, unproject_depth_frame +from .surface import GeometryFrame, SurfaceMesh, build_surface + + +__all__ = ( + "GeometryFrame", + "SurfaceMesh", + "build_surface", + "compact_mesh", + "confidence_percentile_map", + "confidence_percentile_threshold", + "depth_edge_mask", + "depth_discontinuity_ratio", + "edge_culling_stretch_limit", + "foreground_mesh_mask", + "mesh_triangle_faces", + "remove_background_faces", + "triangle_edge_stretch_values", + "triangle_edge_stretch_mask", + "unproject_camera_depth", + "unproject_depth_frame", +) diff --git a/src/snapform/server/geometry/grid_mesh.py b/src/snapform/server/geometry/grid_mesh.py new file mode 100644 index 0000000..ff0e5b5 --- /dev/null +++ b/src/snapform/server/geometry/grid_mesh.py @@ -0,0 +1,202 @@ +def compact_mesh(vertices, faces): + import numpy as np + + vertices = np.asarray(vertices, dtype=np.float32) + faces = np.asarray(faces, dtype=np.int32) + used_vertices = np.unique(faces) + index_map = np.full(len(vertices), -1, dtype=np.int32) + index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32) + vertices = vertices[used_vertices] + faces = index_map[faces] + return vertices, faces + + +def confidence_percentile_map(depth, confidence): + import numpy as np + + valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) + result = np.zeros(depth.shape, dtype=np.float32) + if not valid.any(): + return result + sorted_values = np.sort(np.asarray(confidence)[valid]) + ranks = np.searchsorted( + sorted_values, + np.asarray(confidence)[valid], + side="right", + ) + result[valid] = ranks.astype(np.float32) / len(sorted_values) + return result + + +def depth_discontinuity_ratio(depth): + import numpy as np + + padded = np.pad(depth, 1, mode="edge") + windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3)) + difference = np.nanmax(windows, axis=(-2, -1)) + difference -= np.nanmin(windows, axis=(-2, -1)) + return difference / np.maximum(depth, 1e-6) + + +def triangle_edge_stretch_values(points, triangles): + import numpy as np + + flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3) + triangle_points = flat_points[triangles] + triangle_depth = triangle_points[..., 2] + safe_depth = np.where( + np.abs(triangle_depth) > 1e-8, + triangle_depth, + np.nan, + ) + triangle_rays = triangle_points / safe_depth[..., None] + next_points = np.roll(triangle_points, -1, axis=1) + next_rays = np.roll(triangle_rays, -1, axis=1) + next_depth = np.roll(triangle_depth, -1, axis=1) + local_depth = (triangle_depth + next_depth) * 0.5 + flat_edges = ( + triangle_rays * local_depth[..., None] + - next_rays * local_depth[..., None] + ) + flat_lengths = np.linalg.norm(flat_edges, axis=-1) + actual_lengths = np.linalg.norm( + triangle_points - next_points, + axis=-1, + ) + stretch = actual_lengths / np.maximum(flat_lengths, 1e-8) + stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0 + return stretch + + +def remove_background_faces(faces, foreground_mask): + """Keep faces with at least one foreground vertex.""" + foreground_points = foreground_mask.reshape(-1) + return faces[foreground_points[faces].any(axis=1)] + + +def depth_edge_mask(depth, threshold=12.0 / 255.0): + import numpy as np + + values = np.asarray(depth, dtype=np.float32) + edges = np.zeros(values.shape, dtype=bool) + for axis in (0, 1): + first = np.take(values, range(values.shape[axis] - 1), axis=axis) + second = np.take(values, range(1, values.shape[axis]), axis=axis) + valid = ( + np.isfinite(first) + & np.isfinite(second) + & (first > 0.0) + & (second > 0.0) + ) + scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) + discontinuity = valid & (np.abs(first - second) / scale > threshold) + lower = [slice(None), slice(None)] + upper = [slice(None), slice(None)] + lower[axis] = slice(0, -1) + upper[axis] = slice(1, None) + edges[tuple(lower)] |= discontinuity + edges[tuple(upper)] |= discontinuity + return edges + + +def edge_culling_stretch_limit(culling): + factor = min(max(float(culling) / 100.0, 0.0001), 1.0) + return 1.0 / factor + + +def foreground_mesh_mask(foreground_path, depth_shape): + import numpy as np + from PIL import Image, ImageFilter + + height, width = depth_shape + with Image.open(foreground_path) as source: + alpha = source.convert("RGBA").getchannel("A") + alpha = alpha.resize( + (width, height), + resample=Image.Resampling.BILINEAR, + ) + alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75)) + return np.asarray(alpha) > 0 + + +def confidence_percentile_threshold(depth, confidence, percentile): + import numpy as np + + if float(percentile) <= 0.0: + return None + valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) + if not valid.any(): + raise ValueError("DA3 returned no valid depth confidence values") + return float(np.percentile(confidence[valid], percentile)) + + +def mesh_triangle_faces( + points, + depth, + confidence, + confidence_threshold, + filter_edges, + max_edge_length, + edge_threshold=12.0 / 255.0, + edge_stretch_limit=None, + vertex_mask=None, +): + import numpy as np + + height, width = depth.shape + indices = np.arange(height * width, dtype=np.int32).reshape(height, width) + quads = np.stack( + ( + indices[:-1, :-1], + indices[1:, :-1], + indices[1:, 1:], + indices[:-1, 1:], + ), + axis=-1, + ).reshape(-1, 4) + valid = ( + np.isfinite(points).all(axis=-1) + & np.isfinite(depth) + & (depth > 0.0) + ) + if confidence_threshold is not None: + valid &= np.isfinite(confidence) & ( + confidence >= confidence_threshold + ) + if filter_edges: + valid &= ~depth_edge_mask(depth, edge_threshold) + if vertex_mask is not None: + if np.shape(vertex_mask) != depth.shape: + raise ValueError("Foreground mask and depth dimensions do not match") + valid &= np.asarray(vertex_mask, dtype=bool) + triangles = np.concatenate( + (quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), + axis=0, + ) + face_valid = valid.reshape(-1)[triangles].all(axis=1) + if edge_stretch_limit is not None: + face_valid &= triangle_edge_stretch_mask( + points, + triangles, + edge_stretch_limit, + ) + if max_edge_length > 0.0: + flat_points = points.reshape(-1, 3) + triangle_points = flat_points[triangles] + next_points = np.roll(triangle_points, -1, axis=1) + edge_lengths = np.linalg.norm( + triangle_points - next_points, + axis=-1, + ) + face_valid &= (edge_lengths <= max_edge_length).all(axis=1) + triangles = triangles[face_valid] + if not len(triangles): + return np.empty((0, 3), dtype=np.int32) + return triangles + + +def triangle_edge_stretch_mask(points, triangles, stretch_limit): + stretch = triangle_edge_stretch_values(points, triangles) + tolerance = float(stretch_limit) * (1.0 + 1e-5) + return (stretch <= tolerance).all(axis=1) + diff --git a/src/snapform/server/geometry/projection.py b/src/snapform/server/geometry/projection.py new file mode 100644 index 0000000..fd2e087 --- /dev/null +++ b/src/snapform/server/geometry/projection.py @@ -0,0 +1,31 @@ +def unproject_depth_frame(depth, intrinsics, extrinsics): + import numpy as np + + height, width = depth.shape + x, y = np.meshgrid( + np.arange(width, dtype=np.float32), + np.arange(height, dtype=np.float32), + ) + pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) + rays = (np.linalg.inv(intrinsics) @ pixels.T).T + camera_points = rays * depth.reshape(-1, 1) + camera_points = np.column_stack( + (camera_points, np.ones(camera_points.shape[0], dtype=np.float32)) + ) + world_to_camera = np.eye(4, dtype=np.float32) + world_to_camera[: extrinsics.shape[0], :4] = extrinsics + world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T + return world_points[:, :3].reshape(height, width, 3) + + +def unproject_camera_depth(depth, intrinsics): + import numpy as np + + height, width = depth.shape + x, y = np.meshgrid( + np.arange(width, dtype=np.float32), + np.arange(height, dtype=np.float32), + ) + pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) + rays = (np.linalg.inv(intrinsics) @ pixels.T).T + return (rays * depth.reshape(-1, 1)).reshape(height, width, 3) diff --git a/src/snapform/server/geometry/surface.py b/src/snapform/server/geometry/surface.py new file mode 100644 index 0000000..999a2a3 --- /dev/null +++ b/src/snapform/server/geometry/surface.py @@ -0,0 +1,133 @@ +from dataclasses import dataclass, field +from .grid_mesh import ( + confidence_percentile_map, + depth_discontinuity_ratio, + mesh_triangle_faces, + remove_background_faces, +) + + +@dataclass +class GeometryFrame: + depth: object + points: object + intrinsics: object + valid_mask: object + confidence: object | None = None + model_mask: object | None = None + + +@dataclass +class SurfaceMesh: + vertices: object + faces: object + uv: object + intrinsics: object + image_size: tuple[int, int] + reference_depth: float + point_attributes: dict = field(default_factory=dict) + face_attributes: dict = field(default_factory=dict) + + +def _validate_frame(frame): + import numpy as np + + depth = np.asarray(frame.depth, dtype=np.float32) + points = np.asarray(frame.points, dtype=np.float32) + intrinsics = np.asarray(frame.intrinsics, dtype=np.float32) + valid_mask = np.asarray(frame.valid_mask, dtype=bool) + if depth.ndim != 2: + raise ValueError(f"Unsupported depth shape: {depth.shape}") + if points.shape != (*depth.shape, 3): + raise ValueError("Point grid and depth dimensions do not match") + if intrinsics.shape != (3, 3): + raise ValueError(f"Unsupported intrinsics shape: {intrinsics.shape}") + if valid_mask.shape != depth.shape: + raise ValueError("Valid mask and depth dimensions do not match") + if frame.confidence is not None: + confidence = np.asarray(frame.confidence) + if confidence.shape != depth.shape: + raise ValueError("Confidence and depth dimensions do not match") + if frame.model_mask is not None: + model_mask = np.asarray(frame.model_mask) + if model_mask.shape != depth.shape: + raise ValueError("Model mask and depth dimensions do not match") + return depth, points, intrinsics, valid_mask + + +def _compact_grid(points, faces, depth): + import numpy as np + + height, width = depth.shape + used = np.unique(faces) + index_map = np.full(height * width, -1, dtype=np.int32) + index_map[used] = np.arange(len(used), dtype=np.int32) + u = np.linspace(0.0, 1.0, width, dtype=np.float32) + v = np.linspace(1.0, 0.0, height, dtype=np.float32) + uv = np.stack(np.meshgrid(u, v), axis=-1).reshape(-1, 2)[used] + return ( + points.reshape(-1, 3)[used].astype(np.float32), + index_map[faces], + uv, + used, + ) + + +def build_surface(frame, foreground_mask=None): + import numpy as np + + depth, points, intrinsics, valid_mask = _validate_frame(frame) + faces = mesh_triangle_faces( + points, + depth, + np.ones(depth.shape, dtype=np.float32), + None, + False, + 0.0, + vertex_mask=valid_mask, + ) + if foreground_mask is not None: + foreground_mask = np.asarray(foreground_mask, dtype=bool) + if foreground_mask.shape != depth.shape: + raise ValueError("Foreground mask and depth dimensions do not match") + faces = remove_background_faces(faces, foreground_mask) + if not len(faces): + raise ValueError("Model output does not contain any valid surface faces") + + face_attributes = { + "face_depth_discontinuity": ( + depth_discontinuity_ratio(depth).reshape(-1)[faces].max(axis=1) + ).astype(np.float32), + } + if frame.confidence is not None: + confidence = confidence_percentile_map( + depth, + np.asarray(frame.confidence, dtype=np.float32), + ) + face_attributes["face_confidence"] = ( + confidence.reshape(-1)[faces].min(axis=1).astype(np.float32) + ) + if frame.model_mask is not None: + model_mask = np.asarray(frame.model_mask, dtype=bool) + face_attributes["face_mask"] = model_mask.reshape(-1)[faces].all(axis=1) + + vertices, compact_faces, uv, used = _compact_grid(points, faces, depth) + point_depth = depth.reshape(-1)[used].astype(np.float32) + reference_values = point_depth + if frame.model_mask is not None: + masked = point_depth[ + np.asarray(frame.model_mask, dtype=bool).reshape(-1)[used] + ] + if len(masked): + reference_values = masked + return SurfaceMesh( + vertices=vertices, + faces=compact_faces.astype(np.int32), + uv=uv.astype(np.float32), + intrinsics=intrinsics.astype(np.float32), + image_size=(depth.shape[1], depth.shape[0]), + reference_depth=float(np.median(reference_values)), + point_attributes={"point_depth": point_depth}, + face_attributes=face_attributes, + ) + diff --git a/src/snapform/server/jobs/__init__.py b/src/snapform/server/jobs/__init__.py new file mode 100644 index 0000000..9a97649 --- /dev/null +++ b/src/snapform/server/jobs/__init__.py @@ -0,0 +1,5 @@ +from .depth import run as generate_depth +from .surface import run as generate_surface + + +__all__ = ("generate_depth", "generate_surface") diff --git a/src/snapform/server/jobs/background.py b/src/snapform/server/jobs/background.py new file mode 100644 index 0000000..9bb3888 --- /dev/null +++ b/src/snapform/server/jobs/background.py @@ -0,0 +1,78 @@ +import shutil +import time + +from ..media import copy_frames, prepare_background_input +from ..models.ben2 import infer_foreground, infer_foreground_frames + + +def run(context, parameters): + """Run BEN2 and write transparent foreground images into the Job directory.""" + total_started = time.perf_counter() + context.progress(0.02, "Starting Environment") + context.check_cancelled() + + output_directory = context.directory + output_directory.mkdir(parents=True, exist_ok=True) + device = parameters["device"] + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(parameters["model"]) + + video_frames = parameters.get("video_frames", 16) + start_frame = parameters.get("start_frame", 1) + images, temporary_directory = prepare_background_input( + parameters["input"], + video_frames, + output_directory, + start_frame=start_frame, + cancel_check=context.check_cancelled, + ) + try: + context.progress(0.12, f"Loading BEN2 for {device.upper()}") + if len(images) == 1: + context.progress(0.35, "Removing background") + foreground, load_ms, inference_ms = infer_foreground( + images[0], + model_directory, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + context.progress(0.92, "Writing transparent PNG") + output_path = output_directory / "foreground.png" + foreground.convert("RGBA").save( + output_path, + format="PNG", + compress_level=1, + ) + if not output_path.is_file(): + raise RuntimeError("BEN2 did not write a foreground image") + else: + load_ms, inference_ms = infer_foreground_frames( + images, + model_directory, + device, + output_directory, + context.progress, + 0.12, + 0.7, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + copy_frames(images, output_directory) + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) + + context.check_cancelled() + total_ms = (time.perf_counter() - total_started) * 1000.0 + frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)" + message = ( + f"Background removed{frame_text} in {inference_ms:.0f} ms " + f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" + ) + print(message) + context.progress(1.0, message) + if tuple(output_directory.glob("foreground.*.png")): + return {"foreground": "."} + return {"foreground": "foreground.png"} diff --git a/src/snapform/server/jobs/common.py b/src/snapform/server/jobs/common.py new file mode 100644 index 0000000..c7693ea --- /dev/null +++ b/src/snapform/server/jobs/common.py @@ -0,0 +1,20 @@ +import shutil +from contextlib import contextmanager + +from ..media import prepare_input + + +@contextmanager +def input_frames(context, parameters): + images, temporary_directory = prepare_input( + parameters["input"], + parameters.get("video_frames", 16), + context.directory, + start_frame=parameters.get("start_frame", 1), + cancel_check=context.check_cancelled, + ) + try: + yield images + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) diff --git a/src/snapform/server/jobs/debug_ben2.py b/src/snapform/server/jobs/debug_ben2.py new file mode 100644 index 0000000..aa65f4f --- /dev/null +++ b/src/snapform/server/jobs/debug_ben2.py @@ -0,0 +1,26 @@ +from ..media import copy_frames +from ..models.ben2 import infer_foreground +from .common import input_frames + + +def run(context, parameters): + context.progress(0.02, "Preparing BEN2 API input") + context.directory.mkdir(parents=True, exist_ok=True) + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(parameters["model"]) + with input_frames(context, parameters) as images: + for index, image in enumerate(images): + context.check_cancelled() + foreground, _load_ms, _inference_ms = infer_foreground( + image, + model_directory, + parameters["device"], + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + suffix = "" if len(images) == 1 else f".{index + 1:04d}" + foreground.save(context.directory / f"foreground{suffix}.png") + copy_frames(images, context.directory) + context.check_cancelled() + context.progress(1.0, "Saved raw BEN2 API foreground") + return {"directory": "."} diff --git a/src/snapform/server/jobs/debug_da3.py b/src/snapform/server/jobs/debug_da3.py new file mode 100644 index 0000000..67a496e --- /dev/null +++ b/src/snapform/server/jobs/debug_da3.py @@ -0,0 +1,21 @@ +from ..media import copy_frames +from ..models import da3 +from ..outputs.debug_da3 import write_prediction +from .common import input_frames + + +def run(context, parameters): + context.progress(0.02, "Preparing DA3 API input") + context.directory.mkdir(parents=True, exist_ok=True) + with input_frames(context, parameters) as images: + prediction = da3.infer( + context.resource("model_manager"), + parameters, + images, + cancel_check=context.check_cancelled, + ) + write_prediction(prediction, context.directory / "prediction.npz") + copy_frames(images, context.directory) + context.check_cancelled() + context.progress(1.0, "Saved raw DA3 API prediction") + return {"directory": "."} diff --git a/src/snapform/server/jobs/debug_moge2.py b/src/snapform/server/jobs/debug_moge2.py new file mode 100644 index 0000000..6d2d2c7 --- /dev/null +++ b/src/snapform/server/jobs/debug_moge2.py @@ -0,0 +1,26 @@ +from ..media import copy_frames +from ..models.moge2 import infer_one +from ..outputs.debug_moge2 import write_prediction +from .common import input_frames + + +def run(context, parameters): + context.progress(0.02, "Preparing MoGe-2 API input") + context.directory.mkdir(parents=True, exist_ok=True) + with input_frames(context, parameters) as images: + for index, image in enumerate(images): + context.check_cancelled() + prediction = infer_one( + context.resource("model_manager"), + parameters, + image, + ) + suffix = "" if len(images) == 1 else f".{index + 1:04d}" + write_prediction( + prediction, + context.directory / f"prediction{suffix}.npz", + ) + copy_frames(images, context.directory) + context.check_cancelled() + context.progress(1.0, "Saved raw MoGe-2 API prediction") + return {"directory": "."} diff --git a/src/snapform/server/jobs/depth.py b/src/snapform/server/jobs/depth.py new file mode 100644 index 0000000..b5dc567 --- /dev/null +++ b/src/snapform/server/jobs/depth.py @@ -0,0 +1,49 @@ +from ..media import copy_frames +from ..models import da3, geometry_adapter, moge2 +from ..outputs.depth import write_da3_depth, write_moge_depth +from .common import input_frames +from .foreground import write_optional + + +def run(context, parameters): + context.progress(0.02, "Preparing depth input") + context.directory.mkdir(parents=True, exist_ok=True) + model_manager = context.resource("model_manager") + adapter = geometry_adapter(parameters["model"]) + with input_frames(context, parameters) as images: + if adapter is da3: + prediction = adapter.infer( + model_manager, + parameters, + images, + progress_callback=lambda index, total: context.progress( + 0.15 + 0.65 * index / total, + f"Estimating DA3 depth ({index + 1}/{total})", + ), + cancel_check=context.check_cancelled, + ) + write_da3_depth(prediction, context.directory / "results.npz") + else: + prediction = adapter.infer( + model_manager, + parameters, + images, + include_points=False, + progress_callback=lambda index, total: context.progress( + 0.15 + 0.65 * index / total, + f"Estimating MoGe-2 depth ({index + 1}/{total})", + ), + cancel_check=context.check_cancelled, + ) + write_moge_depth(context.directory, prediction) + write_optional( + context, + parameters, + images, + progress_start=0.84, + progress_span=0.08, + ) + copy_frames(images, context.directory) + context.check_cancelled() + context.progress(1.0, f"Exported depth for {len(images)} frame(s)") + return {"depth": "results.npz"} diff --git a/src/snapform/server/jobs/foreground.py b/src/snapform/server/jobs/foreground.py new file mode 100644 index 0000000..4fb81de --- /dev/null +++ b/src/snapform/server/jobs/foreground.py @@ -0,0 +1,51 @@ +from ..models.ben2 import infer_foreground, infer_foreground_frames + + +def write_optional( + context, + parameters, + images, + *, + progress_start, + progress_span, +): + model_key = parameters.get("remove_background_model", "") + if not model_key: + return [] + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(model_key) + device = parameters["device"] + if len(images) == 1: + context.progress(progress_start, "Removing image background") + foreground, load_ms, inference_ms = infer_foreground( + images[0], + model_directory, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + path = context.directory / "foreground.png" + foreground.convert("RGBA").save( + path, + format="PNG", + compress_level=1, + ) + paths = [path] + else: + load_ms, inference_ms = infer_foreground_frames( + images, + model_directory, + device, + context.directory, + context.progress, + progress_start, + progress_span, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + paths = sorted(context.directory.glob("foreground.*.png")) + print( + f"BEN2 inference {inference_ms:.0f} ms; " + f"model load {load_ms:.0f} ms" + ) + return paths diff --git a/src/snapform/server/jobs/frames.py b/src/snapform/server/jobs/frames.py new file mode 100644 index 0000000..9f3d5c4 --- /dev/null +++ b/src/snapform/server/jobs/frames.py @@ -0,0 +1,21 @@ +import shutil + +from ..media import copy_frames, prepare_input + + +def run(context, parameters): + context.directory.mkdir(parents=True, exist_ok=True) + images, temporary_directory = prepare_input( + parameters["input"], + parameters.get("video_frames", 16), + context.directory, + start_frame=parameters.get("start_frame", 1), + cancel_check=context.check_cancelled, + ) + try: + copy_frames(images, context.directory) + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) + context.progress(1.0, f"Extracted {len(images)} frame(s)") + return {"frames": len(images)} diff --git a/src/snapform/server/jobs/surface.py b/src/snapform/server/jobs/surface.py new file mode 100644 index 0000000..af55ea4 --- /dev/null +++ b/src/snapform/server/jobs/surface.py @@ -0,0 +1,75 @@ +from ..geometry import foreground_mesh_mask +from ..models import geometry_adapter, moge2 +from ..geometry import build_surface +from ..outputs.surface import ( + write_normal_texture, + write_surface_archive, + write_surface_texture, +) +from .common import input_frames +from .foreground import write_optional + + +def run(context, parameters): + context.progress(0.02, "Preparing surface input") + context.directory.mkdir(parents=True, exist_ok=True) + model_manager = context.resource("model_manager") + adapter = geometry_adapter(parameters["model"]) + with input_frames(context, parameters) as images: + if len(images) != 1: + raise ValueError("Convert to Surface currently supports one image") + context.progress(0.12, "Estimating surface geometry") + if adapter is moge2: + prediction = adapter.infer( + model_manager, + parameters, + images, + include_points=True, + progress_callback=lambda index, total: context.progress( + 0.12 + 0.58 * index / total, + f"Estimating MoGe-2 geometry ({index + 1}/{total})", + ), + cancel_check=context.check_cancelled, + ) + else: + prediction = adapter.infer( + model_manager, + parameters, + images, + progress_callback=lambda index, total: context.progress( + 0.12 + 0.58 * index / total, + f"Estimating DA3 geometry ({index + 1}/{total})", + ), + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + frame = adapter.surface_frame(prediction) + foreground_paths = write_optional( + context, + parameters, + images, + progress_start=0.72, + progress_span=0.08, + ) + foreground_mask = None + if foreground_paths: + foreground_mask = foreground_mesh_mask( + foreground_paths[0], + frame.depth.shape, + ) + context.progress(0.82, "Building surface") + mesh = build_surface(frame, foreground_mask=foreground_mask) + context.progress(0.92, "Writing surface") + write_surface_archive(mesh, context.directory / "mesh.npz") + write_surface_texture( + foreground_paths[0] if foreground_paths else images[0], + context.directory / "texture.png", + ) + if adapter is moge2: + write_normal_texture( + prediction.normal[0], + context.directory / "normal.png", + ) + context.check_cancelled() + context.progress(1.0, "Surface is ready") + return {"surface": "mesh.npz"} diff --git a/src/snapform/server/media.py b/src/snapform/server/media.py deleted file mode 100644 index aa751fa..0000000 --- a/src/snapform/server/media.py +++ /dev/null @@ -1,306 +0,0 @@ -import shutil -import tempfile -import time -from pathlib import Path - - -IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} -VIDEO_SUFFIXES = { - ".mp4", - ".mov", - ".avi", - ".mkv", - ".webm", - ".m4v", - ".mts", - ".mpg", - ".mpeg", -} -SCENE_FILENAME = "scene.glb" - - -class TaskCancelled(RuntimeError): - pass - - -def check_cancelled(cancel_check): - if cancel_check is not None and cancel_check(): - raise TaskCancelled("Task cancelled") - - -def collect_images(directory): - directory = Path(directory) - return sorted( - path - for path in directory.iterdir() - if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES - ) - - -def first_frames(images, maximum_frames, start_frame=1): - """Return up to maximum_frames image paths in order from start_frame.""" - if maximum_frames <= 0: - return images - start_index = max(start_frame - 1, 0) - return images[start_index:start_index + maximum_frames] - - -def extract_video_frames( - video_path, - output_directory, - maximum_frames, - av, - start_frame=1, - cancel_check=None, -): - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - try: - container = av.open(str(video_path)) - except Exception as error: - raise RuntimeError(f"Unable to open video: {video_path}") from error - image_paths = [] - skipped = 0 - try: - stream = container.streams.video[0] - for frame in container.decode(stream): - check_cancelled(cancel_check) - if skipped < start_frame - 1: - skipped += 1 - continue - image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" - frame.to_image().save(image_path, format="JPEG") - image_paths.append(image_path) - if len(image_paths) >= maximum_frames: - break - finally: - container.close() - - if not image_paths: - raise RuntimeError(f"Unable to read any video frame from: {video_path}") - return image_paths - - -def prepare_input( - input_path, - video_frames, - temporary_root, - start_frame=1, - cancel_check=None, -): - check_cancelled(cancel_check) - input_path = Path(input_path) - if input_path.is_dir(): - images = collect_images(input_path) - if not images: - raise RuntimeError(f"No image frames found in: {input_path}") - return first_frames(images, video_frames, start_frame), None - suffix = input_path.suffix.lower() - if suffix in IMAGE_SUFFIXES: - return [input_path], None - if suffix not in VIDEO_SUFFIXES: - supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES)) - raise RuntimeError(f"Unsupported input type; expected one of: {supported}") - - import av - - temporary_directory = Path( - tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root) - ) - try: - images = extract_video_frames( - input_path, - temporary_directory, - video_frames, - av, - start_frame=start_frame, - cancel_check=cancel_check, - ) - except Exception: - shutil.rmtree(temporary_directory, ignore_errors=True) - raise - return images, temporary_directory - - -def prepare_background_input( - input_path, - video_frames, - temporary_root, - start_frame=1, - cancel_check=None, -): - """Resolve a BEN2 input to frame images; returns (images, temporary_directory).""" - input_path = Path(input_path) - if input_path.is_dir(): - images = collect_images(input_path) - if not images: - raise RuntimeError(f"No image frames found in: {input_path}") - return first_frames(images, video_frames, start_frame), None - if input_path.suffix.lower() in IMAGE_SUFFIXES: - return [input_path], None - return prepare_input( - input_path, - video_frames, - temporary_root, - start_frame=start_frame, - cancel_check=cancel_check, - ) - - -def copy_frames(images, output_directory): - """Copy a multi-frame input as numbered frame files for the Blender sequence.""" - if len(images) <= 1: - return - for index, image_path in enumerate(images): - suffix = Path(image_path).suffix.lower() or ".jpg" - destination = ( - Path(output_directory) / f"frame.{index + 1:04d}{suffix}" - ) - if Path(image_path).resolve() == destination.resolve(): - continue - shutil.copy2(image_path, destination) - - -def extract_frames(context, parameters): - """Copy selected source frames into the current job directory.""" - context.directory.mkdir(parents=True, exist_ok=True) - images, temporary_directory = prepare_input( - parameters["input"], - parameters.get("video_frames", 16), - context.directory, - start_frame=parameters.get("start_frame", 1), - cancel_check=context.check_cancelled, - ) - try: - copy_frames(images, context.directory) - finally: - if temporary_directory is not None: - shutil.rmtree(temporary_directory, ignore_errors=True) - message = f"Extracted {len(images)} frame(s)" - context.progress(1.0, message) - return {"frames": len(images)} - - -def exported_paths(output_directory, export_format): - output_directory = Path(output_directory) - if export_format == "glb": - scene_path = output_directory / SCENE_FILENAME - return [scene_path] if scene_path.is_file() else [] - if export_format == "gs_ply": - return sorted((output_directory / "gs_ply").glob("*.ply")) - if export_format == "depth_vis": - return sorted((output_directory / "depth_vis").glob("*.jpg")) - if export_format == "mesh": - result_path = output_directory / "mesh.glb" - return [result_path] if result_path.is_file() else [] - if export_format == "mesh_data": - result_path = output_directory / "mesh.npz" - return [result_path] if result_path.is_file() else [] - if export_format == "mini_npz": - result_path = output_directory / "results.npz" - if result_path.is_file(): - return [result_path] - result_path = ( - output_directory / "exports" / "mini_npz" / "results.npz" - ) - return [result_path] if result_path.is_file() else [] - return [] - - -def npz_is_complete(path): - """Return True only when the npz archive can be fully read.""" - try: - import numpy as np - - with np.load(path, allow_pickle=False) as archive: - return "depth" in archive.files - except Exception: - return False - - -def wait_for_exported_paths( - output_directory, - export_format, - timeout=30.0, - cancel_check=None, -): - # DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整， - # 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取 - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - check_cancelled(cancel_check) - result_paths = exported_paths(output_directory, export_format) - if result_paths: - if ( - export_format in ("mini_npz", "npz") - and not npz_is_complete(result_paths[0]) - ): - time.sleep(0.1) - continue - return result_paths - time.sleep(0.05) - return [] - - -def move_exported_file( - source, - destination, - timeout=10.0, - cancel_check=None, -): - """Move an async export after Windows releases its file handle.""" - source = Path(source) - destination = Path(destination) - deadline = time.monotonic() + timeout - while True: - check_cancelled(cancel_check) - try: - source.replace(destination) - return destination - except OSError as error: - is_file_lock = isinstance(error, PermissionError) or getattr( - error, - "winerror", - None, - ) in {32, 33} - if not is_file_lock or time.monotonic() >= deadline: - raise - time.sleep(0.1) - - -def flatten_exported_paths( - output_directory, - result_paths, - cancel_check=None, -): - """Move files produced in exporter subdirectories to the job root.""" - output_directory = Path(output_directory) - flattened_paths = [] - for source in result_paths: - source = Path(source) - check_cancelled(cancel_check) - if source.parent == output_directory: - flattened_paths.append(source) - continue - destination = output_directory / source.name - if destination.exists(): - raise RuntimeError(f"Export filename already exists: {destination.name}") - flattened_paths.append( - move_exported_file( - source, - destination, - cancel_check=cancel_check, - ) - ) - for directory in sorted( - {path.parent for path in map(Path, result_paths)}, - key=lambda path: len(path.parts), - reverse=True, - ): - while directory != output_directory: - try: - directory.rmdir() - except OSError: - break - directory = directory.parent - return flattened_paths diff --git a/src/snapform/server/media/__init__.py b/src/snapform/server/media/__init__.py new file mode 100644 index 0000000..7bb9282 --- /dev/null +++ b/src/snapform/server/media/__init__.py @@ -0,0 +1,26 @@ +from .input import ( + IMAGE_SUFFIXES, + VIDEO_SUFFIXES, + TaskCancelled, + check_cancelled, + collect_images, + copy_frames, + extract_video_frames, + first_frames, + prepare_background_input, + prepare_input, +) + + +__all__ = ( + "IMAGE_SUFFIXES", + "VIDEO_SUFFIXES", + "TaskCancelled", + "check_cancelled", + "collect_images", + "copy_frames", + "extract_video_frames", + "first_frames", + "prepare_background_input", + "prepare_input", +) diff --git a/src/snapform/server/media/input.py b/src/snapform/server/media/input.py new file mode 100644 index 0000000..bfc7195 --- /dev/null +++ b/src/snapform/server/media/input.py @@ -0,0 +1,158 @@ +import shutil +import tempfile +from pathlib import Path + + +IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} +VIDEO_SUFFIXES = { + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", +} +class TaskCancelled(RuntimeError): + pass + + +def check_cancelled(cancel_check): + if cancel_check is not None and cancel_check(): + raise TaskCancelled("Task cancelled") + + +def collect_images(directory): + directory = Path(directory) + return sorted( + path + for path in directory.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + + +def first_frames(images, maximum_frames, start_frame=1): + """Return up to maximum_frames image paths in order from start_frame.""" + if maximum_frames <= 0: + return images + start_index = max(start_frame - 1, 0) + return images[start_index:start_index + maximum_frames] + + +def extract_video_frames( + video_path, + output_directory, + maximum_frames, + av, + start_frame=1, + cancel_check=None, +): + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + try: + container = av.open(str(video_path)) + except Exception as error: + raise RuntimeError(f"Unable to open video: {video_path}") from error + image_paths = [] + skipped = 0 + try: + stream = container.streams.video[0] + for frame in container.decode(stream): + check_cancelled(cancel_check) + if skipped < start_frame - 1: + skipped += 1 + continue + image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" + frame.to_image().save(image_path, format="JPEG") + image_paths.append(image_path) + if len(image_paths) >= maximum_frames: + break + finally: + container.close() + + if not image_paths: + raise RuntimeError(f"Unable to read any video frame from: {video_path}") + return image_paths + + +def prepare_input( + input_path, + video_frames, + temporary_root, + start_frame=1, + cancel_check=None, +): + check_cancelled(cancel_check) + input_path = Path(input_path) + if input_path.is_dir(): + images = collect_images(input_path) + if not images: + raise RuntimeError(f"No image frames found in: {input_path}") + return first_frames(images, video_frames, start_frame), None + suffix = input_path.suffix.lower() + if suffix in IMAGE_SUFFIXES: + return [input_path], None + if suffix not in VIDEO_SUFFIXES: + supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES)) + raise RuntimeError(f"Unsupported input type; expected one of: {supported}") + + import av + + temporary_directory = Path( + tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root) + ) + try: + images = extract_video_frames( + input_path, + temporary_directory, + video_frames, + av, + start_frame=start_frame, + cancel_check=cancel_check, + ) + except Exception: + shutil.rmtree(temporary_directory, ignore_errors=True) + raise + return images, temporary_directory + + +def prepare_background_input( + input_path, + video_frames, + temporary_root, + start_frame=1, + cancel_check=None, +): + """Resolve a BEN2 input to frame images; returns (images, temporary_directory).""" + input_path = Path(input_path) + if input_path.is_dir(): + images = collect_images(input_path) + if not images: + raise RuntimeError(f"No image frames found in: {input_path}") + return first_frames(images, video_frames, start_frame), None + if input_path.suffix.lower() in IMAGE_SUFFIXES: + return [input_path], None + return prepare_input( + input_path, + video_frames, + temporary_root, + start_frame=start_frame, + cancel_check=cancel_check, + ) + + +def copy_frames(images, output_directory): + """Copy a multi-frame input as numbered frame files for the Blender sequence.""" + if len(images) <= 1: + return + for index, image_path in enumerate(images): + suffix = Path(image_path).suffix.lower() or ".jpg" + destination = ( + Path(output_directory) / f"frame.{index + 1:04d}{suffix}" + ) + if Path(image_path).resolve() == destination.resolve(): + continue + shutil.copy2(image_path, destination) + diff --git a/src/snapform/server/model_manager.py b/src/snapform/server/model_manager.py index 6b43439..53aa533 100644 --- a/src/snapform/server/model_manager.py +++ b/src/snapform/server/model_manager.py @@ -27,7 +27,7 @@ class ModelSessionCache: return self._load_ben2(model_dir, requested_device, key) def _load_ben2(self, model_dir, requested_device, key): - from .onnx_ben2 import create_session + from .models.onnx_ben2 import create_session started = time.perf_counter() model = create_session(model_dir, requested_device) @@ -47,7 +47,7 @@ class ModelSessionCache: return self._load_da3(model_dir, requested_device, key) def _load_da3(self, model_dir, requested_device, key): - from .onnx_da3 import create_session + from .models.onnx_da3 import create_session model = create_session(model_dir, requested_device) self.da3_model = model @@ -55,7 +55,7 @@ class ModelSessionCache: return model def get_moge2(self, model_dir, requested_device): - from .onnx_moge2 import create_session + from .models.onnx_moge2 import create_session key = (str(Path(model_dir).resolve()), requested_device, "onnx") if self.da3_model is not None: diff --git a/src/snapform/server/models/__init__.py b/src/snapform/server/models/__init__.py new file mode 100644 index 0000000..0cc87e6 --- /dev/null +++ b/src/snapform/server/models/__init__.py @@ -0,0 +1,8 @@ +from . import ben2, da3, moge2 + + +def geometry_adapter(model_key): + return moge2 if model_key.startswith("MOGE2_") else da3 + + +__all__ = ("ben2", "da3", "moge2", "geometry_adapter") diff --git a/src/snapform/server/models/ben2.py b/src/snapform/server/models/ben2.py new file mode 100644 index 0000000..a93203c --- /dev/null +++ b/src/snapform/server/models/ben2.py @@ -0,0 +1,78 @@ +import time +from pathlib import Path + +from ..media import IMAGE_SUFFIXES, check_cancelled + + +def infer_foreground( + input_path, + model_dir, + device, + model_cache=None, + cancel_check=None, +): + from PIL import Image + + input_path = Path(input_path) + check_cancelled(cancel_check) + if input_path.suffix.lower() not in IMAGE_SUFFIXES: + raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") + from .onnx_ben2 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 = 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) + + check_cancelled(cancel_check) + with Image.open(input_path) as opened_image: + image = opened_image.convert("RGB") + inference_started = time.perf_counter() + foreground = infer(model, image) + inference_ms = (time.perf_counter() - inference_started) * 1000.0 + check_cancelled(cancel_check) + return foreground, load_ms, inference_ms + + +def infer_foreground_frames( + images, + model_dir, + device, + output_directory, + progress_callback, + start_progress, + progress_span, + model_cache=None, + cancel_check=None, +): + """Run BEN2 on every frame and write foreground.0001.png etc.""" + total_load_ms = 0.0 + total_inference_ms = 0.0 + for index, image_path in enumerate(images): + check_cancelled(cancel_check) + progress_callback( + start_progress + progress_span * index / len(images), + f"Removing background ({index + 1}/{len(images)})", + ) + foreground, load_ms, inference_ms = infer_foreground( + image_path, + model_dir, + device, + model_cache=model_cache, + cancel_check=cancel_check, + ) + output_path = output_directory / f"foreground.{index + 1:04d}.png" + foreground.convert("RGBA").save( + output_path, + format="PNG", + compress_level=1, + ) + if not output_path.is_file(): + raise RuntimeError(f"BEN2 did not write: {output_path}") + total_load_ms += load_ms + total_inference_ms += inference_ms + return total_load_ms, total_inference_ms diff --git a/src/snapform/server/models/da3.py b/src/snapform/server/models/da3.py new file mode 100644 index 0000000..c1c7c47 --- /dev/null +++ b/src/snapform/server/models/da3.py @@ -0,0 +1,94 @@ +from ..geometry import unproject_camera_depth +from ..geometry import GeometryFrame + + +def merge_frame_predictions(predictions): + import numpy as np + + if not predictions: + raise ValueError("DA3 returned no frame predictions") + merged = predictions[0] + for name in ( + "depth", + "conf", + "extrinsics", + "intrinsics", + "processed_images", + ): + setattr( + merged, + name, + np.concatenate( + [getattr(prediction, name) for prediction in predictions], + axis=0, + ), + ) + return merged + + +def infer( + model_manager, + parameters, + images, + *, + progress_callback=None, + cancel_check=None, +): + from .onnx_da3 import infer as infer_onnx + from .onnx_runtime import provider_label + + model_directory = model_manager.directory(parameters["model"]) + device = parameters["device"] + model = model_manager.get_da3(model_directory, device) + options = { + "process_res": parameters["process_res"], + "process_res_method": parameters["process_res_method"], + "ref_view_strategy": parameters["ref_view_strategy"], + } + + def run(paths): + try: + return infer_onnx(model, paths, **options) + except UnicodeDecodeError as error: + if provider_label(model) == "DirectML": + raise RuntimeError( + "DA3 inference failed on DirectML. Switch the device to " + "CPU in SnapForm Preferences and retry. " + f"Original error: {error}" + ) from error + raise + + if parameters.get("inference_mode", "joint") != "per_frame": + return run([str(path) for path in images]) + + predictions = [] + for index, image in enumerate(images): + if cancel_check is not None: + cancel_check() + if progress_callback is not None: + progress_callback(index, len(images)) + predictions.append(run([str(image)])) + return merge_frame_predictions(predictions) + + +def surface_frame(prediction): + import numpy as np + + depth = np.asarray(prediction.depth, dtype=np.float32) + confidence = np.asarray(prediction.conf, dtype=np.float32) + intrinsics = np.asarray(prediction.intrinsics, dtype=np.float32) + if depth.shape[0] != 1: + raise ValueError("Convert to Surface currently supports one image") + points = unproject_camera_depth(depth[0], intrinsics[0]) + valid = ( + np.isfinite(depth[0]) + & (depth[0] > 0.0) + & np.isfinite(points).all(axis=-1) + ) + return GeometryFrame( + depth=depth[0], + points=points, + intrinsics=intrinsics[0], + valid_mask=valid, + confidence=confidence[0], + ) diff --git a/src/snapform/server/models/moge2.py b/src/snapform/server/models/moge2.py new file mode 100644 index 0000000..7eb81e2 --- /dev/null +++ b/src/snapform/server/models/moge2.py @@ -0,0 +1,119 @@ +from dataclasses import dataclass + +from ..geometry import GeometryFrame + + +@dataclass +class Prediction: + depth: object + normal: object + mask: object + intrinsics: object + points: object | None + + def archive_values(self): + values = { + "depth": self.depth, + "normal": self.normal, + "mask": self.mask, + "intrinsics": self.intrinsics, + } + if self.points is not None: + values["points"] = self.points + return values + + +def infer( + model_manager, + parameters, + images, + *, + include_points, + progress_callback=None, + cancel_check=None, +): + import numpy as np + depths = [] + normals = [] + masks = [] + intrinsics = [] + points = [] + for index, image_path in enumerate(images): + if cancel_check is not None: + cancel_check() + if progress_callback is not None: + progress_callback(index, len(images)) + prediction = infer_one( + model_manager, + parameters, + image_path, + ) + depths.append(prediction["depth"]) + normals.append(prediction["normal"]) + masks.append(prediction["mask"].astype(bool)) + intrinsics.append(prediction["intrinsics"]) + if include_points: + points.append(prediction["points"]) + if len({array.shape for array in depths}) != 1: + raise RuntimeError("MoGe-2 frame sizes must match") + return Prediction( + depth=np.stack(depths).astype(np.float32), + normal=np.stack(normals).astype(np.float32), + mask=np.stack(masks), + intrinsics=np.stack(intrinsics).astype(np.float32), + points=( + np.stack(points).astype(np.float32) + if include_points + else None + ), + ) + + +def infer_one(model_manager, parameters, image_path): + import numpy as np + from PIL import Image + + from .onnx_moge2 import infer as infer_onnx + + model_directory = model_manager.directory(parameters["model"]) + model = model_manager.get_moge2( + model_directory, + parameters["device"], + ) + with Image.open(image_path) as opened: + image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 + return infer_onnx( + model, + image, + int(parameters["resolution_level"]), + ) + + +def surface_frame(prediction): + import numpy as np + + if prediction.depth.shape[0] != 1 or prediction.points is None: + raise ValueError("Convert to Surface requires one MoGe-2 point grid") + depth = prediction.depth[0] + points = prediction.points[0] + intrinsics = prediction.intrinsics[0].copy() + height, width = depth.shape + if ( + 0.0 <= float(intrinsics[0, 2]) <= 1.0 + and 0.0 <= float(intrinsics[1, 2]) <= 1.0 + ): + intrinsics[0] *= width + intrinsics[1] *= height + intrinsics[2] = (0.0, 0.0, 1.0) + valid = ( + np.isfinite(depth) + & (depth > 0.0) + & np.isfinite(points).all(axis=-1) + ) + return GeometryFrame( + depth=depth, + points=points, + intrinsics=intrinsics, + valid_mask=valid, + model_mask=prediction.mask[0], + ) diff --git a/src/snapform/server/onnx_ben2.py b/src/snapform/server/models/onnx_ben2.py similarity index 100% rename from src/snapform/server/onnx_ben2.py rename to src/snapform/server/models/onnx_ben2.py diff --git a/src/snapform/server/onnx_da3.py b/src/snapform/server/models/onnx_da3.py similarity index 100% rename from src/snapform/server/onnx_da3.py rename to src/snapform/server/models/onnx_da3.py diff --git a/src/snapform/server/onnx_moge2.py b/src/snapform/server/models/onnx_moge2.py similarity index 100% rename from src/snapform/server/onnx_moge2.py rename to src/snapform/server/models/onnx_moge2.py diff --git a/src/snapform/server/onnx_runtime.py b/src/snapform/server/models/onnx_runtime.py similarity index 100% rename from src/snapform/server/onnx_runtime.py rename to src/snapform/server/models/onnx_runtime.py diff --git a/src/snapform/server/moge2.py b/src/snapform/server/moge2.py deleted file mode 100644 index 723281b..0000000 --- a/src/snapform/server/moge2.py +++ /dev/null @@ -1,304 +0,0 @@ -import shutil -from pathlib import Path - -from .background import infer_foreground, infer_foreground_frames -from .geometry import ( - depth_discontinuity_ratio, - depth_edge_mask, - foreground_mesh_mask, - mesh_triangle_faces, -) -from .media import copy_frames, prepare_input - - -def _save_moge_preview(output_directory, index, depth, normal, mask): - import numpy as np - from PIL import Image - - suffix = "" if index is None else f".{index + 1:04d}" - valid = mask & np.isfinite(depth) & (depth > 0) - preview = np.zeros(depth.shape, dtype=np.uint8) - if np.any(valid): - low, high = np.percentile(depth[valid], (2.0, 98.0)) - span = max(float(high - low), 1e-6) - preview[valid] = np.clip( - (depth[valid] - low) / span * 255.0, - 0.0, - 255.0, - ).astype(np.uint8) - Image.fromarray(preview, mode="L").save( - output_directory / f"depth{suffix}.png" - ) - normal_rgb = np.clip( - (normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0) - * 127.5, - 0.0, - 255.0, - ).astype(np.uint8) - Image.fromarray(normal_rgb, mode="RGB").save( - output_directory / f"normal{suffix}.png" - ) - Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save( - output_directory / f"mask{suffix}.png" - ) - - -def export_moge_debug_mesh( - output_directory, - points, - depth, - mask, - normal, - image, - export_format, -): - import numpy as np - import trimesh - - clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04) - faces = mesh_triangle_faces( - points, - depth, - np.ones(depth.shape, dtype=np.float32), - None, - False, - 0.0, - vertex_mask=clean_mask, - ) - 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_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32) - 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( - output_path, - points, - depth, - mask, - normal, - intrinsics, - foreground_mask=None, -): - import numpy as np - - height, width = depth.shape - finite_mask = np.isfinite(depth) & (depth > 0.0) - finite_mask &= np.isfinite(points).all(axis=-1) - if foreground_mask is not None: - finite_mask &= foreground_mask - depth_discontinuity = depth_discontinuity_ratio(depth) - faces = mesh_triangle_faces( - points, - depth, - np.ones(depth.shape, dtype=np.float32), - None, - False, - 0.0, - vertex_mask=finite_mask, - ) - 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 = 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, - vertices=vertices.astype(np.float32), - faces=faces.astype(np.int32), - uv=vertex_uvs.astype(np.float32), - intrinsics=pixel_intrinsics.astype(np.float32), - image_size=np.array((width, height), dtype=np.int32), - reference_depth=np.float32(reference_depth), - point_depth=point_depth, - normal_map=normal.astype(np.float32), - face_mask=face_mask, - face_depth_discontinuity=face_depth_discontinuity, - ) - - -def generate_moge2(context, parameters): - """Run MoGe-2 inference and return the artifact used by the requesting Job.""" - import numpy as np - from PIL import Image - from .onnx_moge2 import ( - infer as infer_moge2, - provider_label, - ) - - context.progress(0.02, "Starting MoGe-2") - output_directory = context.directory - output_directory.mkdir(parents=True, exist_ok=True) - model_manager = context.resource("model_manager") - model_directory = model_manager.directory(parameters["model"]) - device = parameters["device"] - images, temporary_directory = prepare_input( - parameters["input"], - parameters.get("video_frames", 16), - output_directory, - start_frame=parameters.get("start_frame", 1), - cancel_check=context.check_cancelled, - ) - try: - model = model_manager.get_moge2(model_directory, device) - context.progress( - 0.08, - f"Loading MoGe-2 on {provider_label(model)}", - ) - - depths = [] - normals = [] - masks = [] - points = [] - intrinsics = [] - source_images = [] - for index, image_path in enumerate(images): - context.check_cancelled() - context.progress( - 0.15 + 0.68 * index / len(images), - f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})", - ) - with Image.open(image_path) as opened: - image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 - source_images.append(image) - prediction = infer_moge2( - model, - image, - int(parameters["resolution_level"]), - ) - 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) - intrinsics.append(intrinsic) - if bool( - parameters.get("include_points", False) - or parameters.get("surface_mesh", False) - ): - points.append(point) - debug_format = parameters.get("export_format", "") - if not debug_format or debug_format == "maps": - _save_moge_preview( - output_directory, - None if len(images) == 1 else index, - depth, - normal, - mask, - ) - - shapes = {array.shape for array in depths} - if len(shapes) != 1: - raise RuntimeError("MoGe-2 frame sizes must match") - result = { - "depth": np.stack(depths).astype(np.float32), - "normal": np.stack(normals).astype(np.float32), - "mask": np.stack(masks), - "intrinsics": np.stack(intrinsics).astype(np.float32), - } - if points: - result["points"] = np.stack(points).astype(np.float32) - debug_format = parameters.get("export_format", "") - if debug_format in {"glb", "ply"} and len(images) == 1: - context.progress( - 0.84, - f"Exporting MoGe-2 {debug_format.upper()}", - ) - export_moge_debug_mesh( - output_directory, - points[0], - depths[0], - masks[0], - normals[0], - source_images[0], - debug_format, - ) - - foreground_mask = None - remove_background_model = parameters.get("remove_background_model", "") - if remove_background_model: - remove_background_model_dir = model_manager.directory( - remove_background_model - ) - context.progress( 0.87, "Removing image background") - if len(images) == 1: - foreground, _, _ = infer_foreground( - images[0], - remove_background_model_dir, - device, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - foreground_path = output_directory / "foreground.png" - foreground.save(foreground_path) - if points: - foreground_mask = foreground_mesh_mask( - foreground_path, - result["depth"][0].shape, - ) - else: - infer_foreground_frames( - images, - remove_background_model_dir, - device, - output_directory, - context.progress, - 0.87, - 0.08, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - if bool(parameters.get("surface_mesh", False)): - if len(images) != 1 or not points: - raise RuntimeError("MoGe-2 Surface requires one point map") - context.progress( 0.94, "Building official MoGe-2 surface") - export_moge_surface_archive( - output_directory / "mesh.npz", - points[0], - result["depth"][0], - result["mask"][0], - result["normal"][0], - result["intrinsics"][0], - foreground_mask=( - foreground_mask - if remove_background_model - else None - ), - ) - np.savez(output_directory / "results.npz", **result) - copy_frames(images, output_directory) - context.progress( - 1.0, - f"Exported MoGe-2 data for {len(images)} frame(s)", - ) - if context.job_type == "generate-moge2-surface": - return {"surface": "mesh.npz"} - if context.job_type == "generate-moge2-depth": - return {"depth": "results.npz"} - return {"directory": "."} - finally: - if temporary_directory is not None: - shutil.rmtree(temporary_directory, ignore_errors=True) diff --git a/src/snapform/server/outputs/__init__.py b/src/snapform/server/outputs/__init__.py new file mode 100644 index 0000000..a8dc7ad --- /dev/null +++ b/src/snapform/server/outputs/__init__.py @@ -0,0 +1,11 @@ +from .depth import write_da3_depth, write_moge_depth +from .surface import write_normal_texture, write_surface_archive, write_surface_texture + + +__all__ = ( + "write_da3_depth", + "write_moge_depth", + "write_normal_texture", + "write_surface_archive", + "write_surface_texture", +) diff --git a/src/snapform/server/outputs/debug_da3.py b/src/snapform/server/outputs/debug_da3.py new file mode 100644 index 0000000..cef82d0 --- /dev/null +++ b/src/snapform/server/outputs/debug_da3.py @@ -0,0 +1,24 @@ +from pathlib import Path + + +def write_prediction(prediction, output_path): + import numpy as np + + values = {} + for name in ( + "depth", + "conf", + "extrinsics", + "intrinsics", + "processed_images", + "is_metric", + "sky", + "scale_factor", + ): + value = getattr(prediction, name, None) + if value is not None: + values[name] = np.asarray(value) + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez(output_path, **values) + return output_path diff --git a/src/snapform/server/outputs/debug_moge2.py b/src/snapform/server/outputs/debug_moge2.py new file mode 100644 index 0000000..d34d7df --- /dev/null +++ b/src/snapform/server/outputs/debug_moge2.py @@ -0,0 +1,13 @@ +from pathlib import Path + + +def write_prediction(prediction, output_path): + import numpy as np + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + output_path, + **{name: np.asarray(value) for name, value in prediction.items()}, + ) + return output_path diff --git a/src/snapform/server/outputs/depth.py b/src/snapform/server/outputs/depth.py new file mode 100644 index 0000000..8f7b194 --- /dev/null +++ b/src/snapform/server/outputs/depth.py @@ -0,0 +1,64 @@ +from pathlib import Path + +from .surface import write_normal_texture + + +def write_da3_depth(prediction, output_path): + 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_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(output_path, **values) + return output_path + + +def write_moge_depth(output_directory, prediction): + import numpy as np + + output_directory = Path(output_directory) + np.savez( + output_directory / "results.npz", + **prediction.archive_values(), + ) + for index, (depth, normal, mask) in enumerate( + zip(prediction.depth, prediction.normal, prediction.mask) + ): + _write_moge_maps( + output_directory, + None if len(prediction.depth) == 1 else index, + depth, + normal, + mask, + ) + + +def _write_moge_maps(output_directory, index, depth, normal, mask): + import numpy as np + from PIL import Image + + suffix = "" if index is None else f".{index + 1:04d}" + valid = mask & np.isfinite(depth) & (depth > 0) + preview = np.zeros(depth.shape, dtype=np.uint8) + if np.any(valid): + low, high = np.percentile(depth[valid], (2.0, 98.0)) + span = max(float(high - low), 1e-6) + preview[valid] = np.clip( + (depth[valid] - low) / span * 255.0, + 0.0, + 255.0, + ).astype(np.uint8) + Image.fromarray(preview, mode="L").save( + output_directory / f"depth{suffix}.png" + ) + write_normal_texture( + normal, + output_directory / f"normal{suffix}.png", + ) + Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save( + output_directory / f"mask{suffix}.png" + ) diff --git a/src/snapform/server/outputs/surface.py b/src/snapform/server/outputs/surface.py new file mode 100644 index 0000000..0583b6d --- /dev/null +++ b/src/snapform/server/outputs/surface.py @@ -0,0 +1,44 @@ +from pathlib import Path + + +def write_surface_archive(mesh, output_path): + import numpy as np + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + output_path, + vertices=mesh.vertices, + faces=mesh.faces, + uv=mesh.uv, + intrinsics=mesh.intrinsics, + image_size=np.asarray(mesh.image_size, dtype=np.int32), + reference_depth=np.asarray(mesh.reference_depth, dtype=np.float32), + **mesh.point_attributes, + **mesh.face_attributes, + ) + return output_path + + +def write_surface_texture(image_path, output_path): + from PIL import Image + + output_path = Path(output_path) + with Image.open(image_path) as image: + image.convert("RGBA").save(output_path, format="PNG", compress_level=1) + return output_path + + +def write_normal_texture(normal, output_path): + import numpy as np + from PIL import Image + + normal_rgb = np.clip( + (normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0) + * 127.5, + 0.0, + 255.0, + ).astype(np.uint8) + output_path = Path(output_path) + Image.fromarray(normal_rgb, mode="RGB").save(output_path) + return output_path diff --git a/tests/test_addon.py b/tests/test_addon.py index 2977eaf..939b8b6 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -71,6 +71,12 @@ class AddonTest(unittest.TestCase): sys.path.insert(0, str(SOURCE_ROOT)) cls.paths_before_import = tuple(sys.path) cls.addon = importlib.import_module("snapform") + cls.image_depth = importlib.import_module( + "snapform.blender.depth_plane" + ) + cls.image_mesh = importlib.import_module( + "snapform.blender.surface" + ) cls.paths_after_import = tuple(sys.path) cls.fake_bpy = fake_bpy @@ -494,9 +500,9 @@ class AddonTest(unittest.TestCase): object=SimpleNamespace(data=SimpleNamespace()), ) operators = ( - (self.addon.image_depth, self.addon.ConvertToDepthPlane), - (self.addon.image_mesh, self.addon.ConvertToSurface), - (self.addon.image_depth, self.addon.RemoveImageBackground), + (self.image_depth, self.addon.ConvertToDepthPlane), + (self.image_mesh, self.addon.ConvertToSurface), + (self.image_depth, self.addon.RemoveImageBackground), ) states = ( @@ -723,7 +729,7 @@ class AddonTest(unittest.TestCase): ) with patch.object( - self.addon.image_depth, + self.image_depth, "prepare_image_input", return_value=(Path(__file__), False), ): diff --git a/tests/test_ben2_onnx.py b/tests/test_ben2_onnx.py index c1a7680..76d46c2 100644 --- a/tests/test_ben2_onnx.py +++ b/tests/test_ben2_onnx.py @@ -12,6 +12,7 @@ SOURCE_DIR = ( / "src" / "snapform" / "server" + / "models" ) sys.path.insert(0, str(SOURCE_DIR)) diff --git a/tests/test_da3_onnx.py b/tests/test_da3_onnx.py index 6a007c3..a67d756 100644 --- a/tests/test_da3_onnx.py +++ b/tests/test_da3_onnx.py @@ -13,6 +13,7 @@ SOURCE_DIR = ( / "src" / "snapform" / "server" + / "models" ) sys.path.insert(0, str(SOURCE_DIR)) diff --git a/tests/test_moge_onnx.py b/tests/test_moge_onnx.py index 3124bc3..b1d28a5 100644 --- a/tests/test_moge_onnx.py +++ b/tests/test_moge_onnx.py @@ -16,12 +16,28 @@ def load_module(): snapform.__path__ = [str(PROJECT_ROOT / "src" / "snapform")] server = ModuleType("snapform.server") server.__path__ = [str(PROJECT_ROOT / "src" / "snapform" / "server")] + models = ModuleType("snapform.server.models") + models.__path__ = [ + str(PROJECT_ROOT / "src" / "snapform" / "server" / "models") + ] specification = importlib.util.spec_from_file_location( - "snapform.server.onnx_moge2", - PROJECT_ROOT / "src" / "snapform" / "server" / "onnx_moge2.py", + "snapform.server.models.onnx_moge2", + PROJECT_ROOT + / "src" + / "snapform" + / "server" + / "models" + / "onnx_moge2.py", ) module = importlib.util.module_from_spec(specification) - with patch.dict(sys.modules, {"snapform": snapform, "snapform.server": server}): + with patch.dict( + sys.modules, + { + "snapform": snapform, + "snapform.server": server, + "snapform.server.models": models, + }, + ): specification.loader.exec_module(module) return module diff --git a/tests/test_pack.py b/tests/test_pack.py index e193c38..ede9abf 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -58,19 +58,28 @@ class PackTest(unittest.TestCase): "assets/O_SnapForm.blend", names, ) - self.assertIn("server/media.py", names) - self.assertIn("server/geometry.py", names) - self.assertIn("server/background.py", names) - self.assertIn("server/moge2.py", names) - self.assertIn("server/da3.py", names) + self.assertIn("server/media/input.py", names) + self.assertIn("server/geometry/projection.py", names) + self.assertIn("server/geometry/grid_mesh.py", names) + self.assertIn("server/geometry/surface.py", names) + self.assertIn("server/jobs/background.py", names) + self.assertIn("server/models/ben2.py", names) + self.assertIn("server/models/da3.py", names) + self.assertIn("server/models/moge2.py", names) + self.assertIn("server/jobs/depth.py", names) + self.assertIn("server/jobs/surface.py", names) + self.assertIn("server/jobs/debug_da3.py", names) + self.assertIn("server/jobs/debug_moge2.py", names) + self.assertIn("server/jobs/debug_ben2.py", names) self.assertNotIn("server/processing.py", names) self.assertNotIn("server/progress.py", names) self.assertIn("model_catalog.py", names) self.assertIn("server/model_download.py", names) - self.assertIn("server/onnx_ben2.py", names) - self.assertIn("server/onnx_da3.py", names) - self.assertIn("server/onnx_moge2.py", names) - self.assertNotIn("server/jobs.py", names) + self.assertIn("server/models/onnx_ben2.py", names) + self.assertIn("server/models/onnx_da3.py", names) + self.assertIn("server/models/onnx_moge2.py", names) + self.assertNotIn("server/da3.py", names) + self.assertNotIn("server/moge2.py", names) self.assertIn("server/__init__.py", names) self.assertIn( "wheels/blendjob-0.1.13-py3-none-any.whl", diff --git a/tests/test_server.py b/tests/test_server.py index 15881fa..6eedabc 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -113,13 +113,13 @@ class ServerTest(unittest.TestCase): def test_ben2_model_is_reused_for_matching_runtime_key(self): loads = [] - ben2 = ModuleType("server.onnx_ben2") + ben2 = ModuleType("server.models.onnx_ben2") 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, {"server.onnx_ben2": ben2}): + with patch.dict(sys.modules, {"server.models.onnx_ben2": ben2}): first, _load_ms = cache.get_ben2("model", device) second, cached_load_ms = cache.get_ben2("model", device) @@ -131,7 +131,7 @@ class ServerTest(unittest.TestCase): def test_da3_model_is_reused_for_matching_runtime_key(self): loads = [] - da3 = ModuleType("server.onnx_da3") + da3 = ModuleType("server.models.onnx_da3") da3.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) @@ -140,7 +140,7 @@ class ServerTest(unittest.TestCase): with patch.dict( sys.modules, { - "server.onnx_da3": da3, + "server.models.onnx_da3": da3, }, ): first = cache.get_da3("model", device) @@ -172,9 +172,9 @@ class ServerTest(unittest.TestCase): moge = object() - fake_moge = ModuleType("server.onnx_moge2") + fake_moge = ModuleType("server.models.onnx_moge2") fake_moge.create_session = lambda _model_dir, _device: moge - with patch.dict(sys.modules, {"server.onnx_moge2": fake_moge}): + with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}): self.assertIs(cache.get_moge2("moge", "cuda"), moge) self.assertIsNone(cache.da3_model) @@ -217,10 +217,8 @@ class ServerTest(unittest.TestCase): "download-required-models", "remove-background", "extract-frames", - "generate-da3-depth", - "generate-moge2-depth", - "generate-da3-surface", - "generate-moge2-surface", + "generate-depth", + "generate-surface", "debug-da3", "debug-moge2", "debug-ben2", @@ -265,7 +263,17 @@ class ServerTest(unittest.TestCase): self.server_module, "ModelManager", return_value=fake_models ): server = self.create_server() - context = server.submit("remove-background", {}, job_id="job") + started = threading.Event() + + @server.job("cancellable") + def cancellable(context, _parameters): + started.set() + while not context._is_cancelled(): + time.sleep(0.001) + context.check_cancelled() + + context = server.submit("cancellable", {}, job_id="job") + self.assertTrue(started.wait(1.0)) self.assertIsNotNone(server.cancel("job")) self.assertTrue(context._is_cancelled()) status = self.wait_for_job(context) diff --git a/tests/test_worker.py b/tests/test_worker.py index 86b8a0c..e11ebea 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -74,7 +74,7 @@ class FakeModels: return self.model_directory def get_ben2(self, _model_directory, _device): - from server.onnx_ben2 import create_session + from server.models.onnx_ben2 import create_session return create_session(self.model_directory, "cpu"), 0.0 @@ -108,31 +108,36 @@ class FakeJobContext: class WorkerTest(unittest.TestCase): @classmethod def setUpClass(cls): - cls.media = load_server_module("media") + cls.media = load_server_module("media.input") cls.geometry = load_server_module("geometry") - cls.background = load_server_module("background") - cls.moge2 = load_server_module("moge2") - cls.da3 = load_server_module("da3") + cls.surface = load_server_module("geometry.surface") + cls.background = load_server_module("jobs.background") + cls.ben2 = load_server_module("models.ben2") + cls.da3_model = load_server_module("models.da3") + cls.moge2_model = load_server_module("models.moge2") + cls.surface_outputs = load_server_module("outputs.surface") + cls.debug_da3_outputs = load_server_module("outputs.debug_da3") + cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2") + cls.depth_job = load_server_module("jobs.depth") + cls.surface_job = load_server_module("jobs.surface") + cls.debug_da3_job = load_server_module("jobs.debug_da3") + cls.debug_moge2_job = load_server_module("jobs.debug_moge2") + cls.debug_ben2_job = load_server_module("jobs.debug_ben2") + cls.frames_job = load_server_module("jobs.frames") + cls.foreground_job = load_server_module("jobs.foreground") cls.download = load_server_module("model_download") - def test_foreground_preserves_source_rgb_and_ben2_alpha(self): - alpha = SimpleNamespace(size=(2, 1)) - source_rgba = SimpleNamespace(size=(2, 1), applied_alpha=None) - source_rgba.putalpha = lambda value: setattr( - source_rgba, - "applied_alpha", - value, - ) - source = SimpleNamespace(convert=lambda _mode: source_rgba) - foreground_rgba = SimpleNamespace( - getchannel=lambda _name: alpha, - ) - foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba) + def test_surface_has_one_model_independent_builder(self): + self.assertTrue(callable(self.surface.build_surface)) + self.assertTrue(callable(self.surface_outputs.write_surface_archive)) + self.assertFalse(hasattr(self.surface, "export_da3_surface")) + self.assertFalse(hasattr(self.surface, "export_moge2_surface")) - result = self.background.preserve_source_rgb(source, foreground) - - self.assertIs(result, source_rgba) - self.assertIs(result.applied_alpha, alpha) + def test_model_adapters_only_normalize_surface_input(self): + self.assertTrue(callable(self.da3_model.surface_frame)) + self.assertTrue(callable(self.moge2_model.surface_frame)) + self.assertFalse(hasattr(self.da3_model, "generate_surface")) + self.assertFalse(hasattr(self.moge2_model, "generate_surface")) def test_collect_images_is_sorted_and_filters_suffixes(self): with tempfile.TemporaryDirectory() as directory: @@ -273,70 +278,6 @@ class WorkerTest(unittest.TestCase): [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) - def test_npz_is_complete_requires_readable_archive(self): - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - valid_path = root / "valid.npz" - np.savez_compressed( - valid_path, - depth=np.zeros((4, 4)), - ) - self.assertTrue(self.media.npz_is_complete(valid_path)) - - empty_path = root / "empty.npz" - empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18) - self.assertFalse(self.media.npz_is_complete(empty_path)) - - def test_wait_for_exported_paths_waits_for_complete_npz(self): - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - npz_path = root / "exports" / "mini_npz" / "results.npz" - npz_path.parent.mkdir(parents=True) - npz_path.touch() - with patch.object( - self.media, - "npz_is_complete", - side_effect=[False, True], - ): - result = self.media.wait_for_exported_paths( - root, - "mini_npz", - timeout=2.0, - ) - self.assertEqual(result, [npz_path]) - - def test_move_exported_file_retries_windows_file_lock(self): - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - source = root / "source.npz" - destination = root / "destination.npz" - source.touch() - path_type = type(source) - original_replace = path_type.replace - calls = 0 - - def locked_once(path, target): - nonlocal calls - calls += 1 - if calls == 1: - error = PermissionError("file is in use") - error.winerror = 32 - raise error - return original_replace(path, target) - - with patch.object(path_type, "replace", new=locked_once), patch.object( - self.media.time, - "sleep", - ) as sleep: - result = self.media.move_exported_file( - source, - destination, - ) - - self.assertEqual(result, destination) - self.assertTrue(destination.is_file()) - sleep.assert_called_once_with(0.1) - def test_rejects_unsupported_input_file(self): with tempfile.TemporaryDirectory() as directory: input_path = Path(directory) / "input.txt" @@ -531,7 +472,7 @@ class WorkerTest(unittest.TestCase): ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext - fake_ben2 = ModuleType("server.onnx_ben2") + fake_ben2 = ModuleType("server.models.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() fake_ben2.infer = lambda _model, image: ( inference_options.append((image, {})) or FakeImage() @@ -543,7 +484,7 @@ class WorkerTest(unittest.TestCase): sys.modules, { "torch": fake_torch, - "server.onnx_ben2": fake_ben2, + "server.models.onnx_ben2": fake_ben2, "PIL": fake_pil, }, ): @@ -551,7 +492,7 @@ class WorkerTest(unittest.TestCase): input_path = root / "input.png" input_path.touch() context = FakeJobContext(root / "result", "remove-background") - self.background.remove_background( + self.background.run( context, {"input": str(input_path), "model": "BEN2", "device": "cpu"}, ) @@ -615,7 +556,7 @@ class WorkerTest(unittest.TestCase): ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext - fake_ben2 = ModuleType("server.onnx_ben2") + fake_ben2 = ModuleType("server.models.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() fake_ben2.infer = lambda _model, image: ( inference_options.append((image, {})) or FakeImage() @@ -627,7 +568,7 @@ class WorkerTest(unittest.TestCase): sys.modules, { "torch": fake_torch, - "server.onnx_ben2": fake_ben2, + "server.models.onnx_ben2": fake_ben2, "PIL": fake_pil, }, ): @@ -637,7 +578,7 @@ class WorkerTest(unittest.TestCase): (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() context = FakeJobContext(root / "result", "remove-background") - self.background.remove_background( + self.background.run( context, { "input": str(input_directory), @@ -673,7 +614,7 @@ class WorkerTest(unittest.TestCase): for index in range(1, 6): (input_directory / f"frame-{index:04d}.jpg").touch() context = FakeJobContext(root / "result", "extract-frames") - self.media.extract_frames( + self.frames_job.run( context, { "input": str(input_directory), @@ -694,62 +635,166 @@ class WorkerTest(unittest.TestCase): self.assertEqual(context.status["progress"], 1.0) self.assertIn("3 frame", context.status["message"]) - def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self): - inference_options = [] - mesh_calls = [] - + def test_da3_debug_saves_only_direct_api_prediction(self): 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), + is_metric=0, + sky=None, + scale_factor=None, ) - fake_da3 = ModuleType("server.onnx_da3") - fake_da3.create_session = lambda _model_dir, _device: object() + fake_da3 = ModuleType("server.models.onnx_da3") + calls = [] def fake_infer(_session, images, **options): - inference_options.append((images, options)) + calls.append((images, options)) return prediction fake_da3.infer = fake_infer - fake_runtime = ModuleType("server.onnx_runtime") + fake_runtime = ModuleType("server.models.onnx_runtime") + fake_runtime.provider_label = lambda _session: "CPU" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "input.png" + input_path.touch() + with patch.dict( + sys.modules, + { + "server.models.onnx_da3": fake_da3, + "server.models.onnx_runtime": fake_runtime, + }, + ): + context = FakeJobContext(root / "result", "debug-da3") + self.debug_da3_job.run( + context, + { + "input": str(input_path), + "model": "DA3", + "video_frames": 1, + "device": "cpu", + "process_res": 518, + "process_res_method": "lower_bound_resize", + "ref_view_strategy": "middle", + }, + ) + with np.load(context.directory / "prediction.npz") as result: + self.assertEqual( + set(result.files), + { + "depth", + "conf", + "extrinsics", + "intrinsics", + "processed_images", + "is_metric", + }, + ) + np.testing.assert_array_equal(result["conf"], prediction.conf) + self.assertFalse(any(context.directory.glob("*.glb"))) + self.assertFalse(any(context.directory.glob("*.png"))) + self.assertEqual(calls[0][0], [str(input_path)]) + self.assertEqual(calls[0][1]["process_res"], 518) + + def test_moge2_debug_saves_each_direct_api_dictionary(self): + prediction = { + "points": np.ones((2, 2, 3), dtype=np.float32), + "depth": np.ones((2, 2), dtype=np.float32), + "normal": np.ones((2, 2, 3), dtype=np.float32), + "mask": np.ones((2, 2), dtype=bool), + "intrinsics": np.eye(3, dtype=np.float32), + } + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + frames = root / "frames" + frames.mkdir() + (frames / "a.png").touch() + (frames / "b.png").touch() + context = FakeJobContext(root / "result", "debug-moge2") + with patch.object( + self.debug_moge2_job, + "infer_one", + return_value=prediction, + ) as infer_one: + self.debug_moge2_job.run( + context, + { + "input": str(frames), + "model": "MOGE2_VITS_NORMAL", + "video_frames": 2, + "device": "cpu", + "resolution_level": 5, + }, + ) + self.assertEqual(infer_one.call_count, 2) + archives = sorted(context.directory.glob("prediction.*.npz")) + self.assertEqual(len(archives), 2) + with np.load(archives[0]) as result: + self.assertEqual(set(result.files), set(prediction)) + self.assertFalse(any(context.directory.glob("*.glb"))) + + def test_ben2_debug_saves_direct_api_foreground(self): + class Foreground: + def save(self, path): + Path(path).write_bytes(b"foreground") + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "input.png" + input_path.touch() + context = FakeJobContext(root / "result", "debug-ben2") + with patch.object( + self.debug_ben2_job, + "infer_foreground", + return_value=(Foreground(), 1.0, 2.0), + ) as infer_foreground: + self.debug_ben2_job.run( + context, + { + "input": str(input_path), + "model": "BEN2_BASE", + "video_frames": 1, + "device": "cpu", + }, + ) + infer_foreground.assert_called_once() + self.assertEqual( + (context.directory / "foreground.png").read_bytes(), + b"foreground", + ) + def test_da3_removes_background_after_depth_using_original_input(self): + order = [] + inference_inputs = [] + fake_da3 = ModuleType("server.models.onnx_da3") + + def fake_infer(_session, images, **_options): + order.append("depth") + inference_inputs.append(images) + return 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.infer = fake_infer + fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" - def fake_mesh_export( - prediction, - images, - output_directory, - confidence_percentile, - filter_edges, - max_edge_length, - ): - mesh_calls.append( - ( - prediction, - images, - confidence_percentile, - filter_edges, - max_edge_length, - ) - ) - output_path = Path(output_directory) / "mesh.glb" - output_path.touch() - return output_path + class FakeForeground: + def convert(self, _mode): + return self - def fake_point_cloud(_prediction, output_directory, *_args): - (Path(output_directory) / "scene.glb").touch() + def save(self, path, **_options): + Path(path).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)), - ) + def fake_foreground(image, *_args, **_kwargs): + order.append("background") + self.assertEqual(image, input_path) + return FakeForeground(), 1.0, 2.0 with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -758,77 +803,40 @@ class WorkerTest(unittest.TestCase): with patch.dict( sys.modules, { - "server.onnx_da3": fake_da3, - "server.onnx_runtime": fake_runtime, + "server.models.onnx_da3": fake_da3, + "server.models.onnx_runtime": fake_runtime, }, ), patch.object( - self.da3, - "export_textured_mesh_glb", - side_effect=fake_mesh_export, - ), patch.object( - self.da3, - "export_da3_point_cloud", - side_effect=fake_point_cloud, - ), patch.object( - self.da3, - "export_da3_depth_preview", - side_effect=fake_depth_preview, - ), patch.object( - self.da3, - "export_da3_mini_npz", - side_effect=fake_mini_npz, + self.foreground_job, + "infer_foreground", + side_effect=fake_foreground, ): - debug_formats = ( - "glb", - "depth_vis", - "mini_npz", - "mesh", + context = FakeJobContext( + root / "result", + "generate-depth", + ) + self.depth_job.run( + context, + { + "input": str(input_path), + "model": "DA3", + "remove_background_model": "BEN2", + "video_frames": 1, + "export_format": "mini_npz", + "device": "cpu", + "process_res": 518, + "process_res_method": "upper_bound_resize", + "ref_view_strategy": "first", + "inference_mode": "per_frame", + }, ) - for export_format in debug_formats: - context = FakeJobContext( - root / export_format, - "debug-da3", - ) - self.da3.generate_da3( - context, - { - "input": str(input_path), - "model": "DA3", - "video_frames": 16, - "export_format": export_format, - "device": "cpu", - "process_res": 518, - "process_res_method": "lower_bound_resize", - "ref_view_strategy": "middle", - "conf_thresh_percentile": 35.0, - "num_max_points": 12345, - "mesh_filter_edges": True, - "mesh_max_edge_length": 0.2, - }, - ) - self.assertEqual(context.status["progress"], 1.0) - self.assertIn(export_format, context.status["message"]) - result_directory = context.directory - 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), 4) - self.assertEqual(len(mesh_calls), 1) - for images, options in inference_options: - self.assertEqual(images, [str(input_path)]) - self.assertEqual(options["process_res"], 518) - self.assertEqual( - options["process_res_method"], - "lower_bound_resize", - ) - self.assertEqual(options["ref_view_strategy"], "middle") - self.assertEqual(mesh_calls[0][1], [input_path]) - self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2)) + self.assertEqual(order, ["depth", "background"]) + self.assertEqual(inference_inputs, [[str(input_path)]]) def test_reconstruct_copies_frames_for_multi_frame_input(self): inference_inputs = [] - fake_da3 = ModuleType("server.onnx_da3") + fake_da3 = ModuleType("server.models.onnx_da3") fake_da3.create_session = lambda _model_dir, _device: object() def fake_infer(_session, images, **_options): @@ -842,7 +850,7 @@ class WorkerTest(unittest.TestCase): ) fake_da3.infer = fake_infer - fake_runtime = ModuleType("server.onnx_runtime") + fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" with tempfile.TemporaryDirectory() as directory: @@ -854,15 +862,15 @@ class WorkerTest(unittest.TestCase): with patch.dict( sys.modules, { - "server.onnx_da3": fake_da3, - "server.onnx_runtime": fake_runtime, + "server.models.onnx_da3": fake_da3, + "server.models.onnx_runtime": fake_runtime, }, ): context = FakeJobContext( root / "result", - "generate-da3-depth", + "generate-depth", ) - self.da3.generate_da3( + self.depth_job.run( context, { "input": str(input_directory), @@ -1099,68 +1107,51 @@ class WorkerTest(unittest.TestCase): ) np.testing.assert_array_equal(compact_faces, [[0, 1, 2]]) - def test_surface_background_faces_are_removed_before_export(self): + def test_surface_background_keeps_faces_with_any_foreground_vertex(self): faces = np.asarray( ((0, 1, 4), (0, 4, 3), (1, 2, 5), (1, 5, 4)), dtype=np.int32, ) foreground = np.asarray( - ((False, True, True), (False, True, True)), + ((False, False, False), (False, True, False)), dtype=bool, ) filtered = self.geometry.remove_background_faces(faces, foreground) - np.testing.assert_array_equal(filtered, faces[2:]) + np.testing.assert_array_equal(filtered, faces[[0, 1, 3]]) + + def test_surface_background_mask_keeps_any_nonzero_alpha(self): + from PIL import Image + + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + for alpha, expected in ((0, False), (1, True), (127, True)): + path = root / f"alpha-{alpha}.png" + Image.new("RGBA", (2, 2), (10, 20, 30, alpha)).save(path) + mask = self.geometry.foreground_mesh_mask(path, (2, 2)) + self.assertEqual(mask.dtype, np.bool_) + self.assertTrue(np.all(mask == expected)) def test_image_mesh_data_keeps_camera_geometry_and_uv(self): - class FakeImage: - def __enter__(self): - return self - - def __exit__(self, *_args): - return None - - def convert(self, mode): - self.mode = mode - return self - - def save(self, path, **_options): - Path(path).write_bytes(b"texture") - - fake_pil = ModuleType("PIL") - fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) - depth = np.full((1, 3, 3), 2.0, dtype=np.float32) - confidence = np.ones((1, 3, 3), dtype=np.float32) + depth = np.full((3, 3), 2.0, dtype=np.float32) + confidence = np.ones((3, 3), dtype=np.float32) intrinsics = np.asarray( - [[[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]]], + [[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32, ) - prediction = SimpleNamespace( + frame = self.surface.GeometryFrame( depth=depth, - conf=confidence, + points=self.geometry.unproject_camera_depth(depth, intrinsics), intrinsics=intrinsics, + valid_mask=np.ones(depth.shape, dtype=bool), + confidence=confidence, ) - progress = [] - - with tempfile.TemporaryDirectory() as directory, patch.dict( - sys.modules, - {"PIL": fake_pil}, - ), patch.object( - self.geometry, - "compact_mesh", - side_effect=lambda vertices, faces: ( - vertices, - faces, - ), - ): - result_path, texture_path = self.geometry.export_image_mesh_data( - prediction, - Path(directory) / "input.png", - Path(directory) / "result", - progress_callback=lambda step, total, label: progress.append( - (step, total, label) - ), + mesh = self.surface.build_surface(frame) + with tempfile.TemporaryDirectory() as directory: + result_path = self.surface_outputs.write_surface_archive( + mesh, + Path(directory) / "mesh.npz", ) with np.load(result_path) as result: self.assertEqual(result["vertices"].shape, (9, 3)) @@ -1170,7 +1161,7 @@ class WorkerTest(unittest.TestCase): self.assertNotIn("postprocess_method", result.files) self.assertEqual(result["point_depth"].shape, (9,)) self.assertEqual( - result["face_confidence_percentile"].shape, + result["face_confidence"].shape, (8,), ) self.assertEqual( @@ -1183,17 +1174,67 @@ class WorkerTest(unittest.TestCase): "image_size", "reference_depth", "point_depth", - "face_confidence_percentile", + "face_confidence", "face_depth_discontinuity", }, ) np.testing.assert_allclose(result["uv"][0], [0.0, 1.0]) np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0]) - self.assertTrue(texture_path.is_file()) - self.assertEqual([item[0] for item in progress], list(range(1, 7))) - self.assertTrue(all(item[1] == 6 for item in progress)) + self.assertNotIn("normal_map", result.files) - def test_prediction_attributes_keep_debug_structure(self): + def test_moge_surface_uses_the_same_archive_without_normal_data(self): + depth = np.full((3, 3), 2.0, dtype=np.float32) + intrinsics = np.asarray( + [[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + frame = self.surface.GeometryFrame( + depth=depth, + points=self.geometry.unproject_camera_depth(depth, intrinsics), + intrinsics=intrinsics, + valid_mask=np.ones(depth.shape, dtype=bool), + model_mask=np.asarray( + ( + (False, False, False), + (True, True, True), + (True, True, True), + ), + dtype=bool, + ), + ) + mesh = self.surface.build_surface(frame) + with tempfile.TemporaryDirectory() as directory: + path = self.surface_outputs.write_surface_archive( + mesh, + Path(directory) / "mesh.npz", + ) + with np.load(path) as result: + self.assertIn("face_mask", result.files) + self.assertNotIn("face_confidence", result.files) + self.assertNotIn("normal", result.files) + self.assertNotIn("normal_map", result.files) + + def test_moge_surface_normal_is_exported_as_a_texture(self): + from PIL import Image + + normal = np.asarray( + [[[0.0, 0.0, -1.0]]], + dtype=np.float32, + ) + with tempfile.TemporaryDirectory() as directory: + path = self.surface_outputs.write_normal_texture( + normal, + Path(directory) / "normal.png", + ) + with Image.open(path) as image: + self.assertEqual(image.mode, "RGB") + self.assertEqual(image.size, (1, 1)) + red, green, blue = image.getpixel((0, 0)) + self.assertIn(red, {127, 128}) + self.assertIn(green, {127, 128}) + self.assertEqual(blue, 255) + + def test_da3_debug_writer_keeps_api_values_without_derived_attributes(self): prediction = SimpleNamespace( depth=np.ones((1, 2, 3), dtype=np.float32), conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3), @@ -1202,18 +1243,18 @@ class WorkerTest(unittest.TestCase): ) with tempfile.TemporaryDirectory() as directory: - output_path = self.geometry.export_prediction_attributes( + output_path = self.debug_da3_outputs.write_prediction( prediction, - directory, + Path(directory) / "prediction.npz", ) with np.load(output_path) as attributes: self.assertEqual(attributes["depth"].shape, (1, 2, 3)) self.assertEqual(attributes["conf"].shape, (1, 2, 3)) - self.assertEqual( - attributes["confidence_percentile"].shape, - (1, 2, 3), - ) + self.assertNotIn("confidence_percentile", attributes.files) if __name__ == "__main__": unittest.main() + + + diff --git a/tools/build_geometry_nodes_asset.py b/tools/build_geometry_nodes_asset.py index cb5d1cc..47fb19f 100644 --- a/tools/build_geometry_nodes_asset.py +++ b/tools/build_geometry_nodes_asset.py @@ -381,6 +381,15 @@ def build_group( ) geometry_input_socket.structure_type = "SINGLE" + remesh_iterations = new_socket( + interface, + "Iterations", + "INPUT", + "NodeSocketInt", + ) + remesh_iterations.default_value = 2 + remesh_iterations.min_value = 0 + remesh_iterations.max_value = 8 resolution = new_socket( interface, "Resolution", @@ -390,15 +399,6 @@ def build_group( resolution.default_value = 0.1 resolution.min_value = 0.000001 resolution.subtype = "DISTANCE" - remesh_iterations = new_socket( - interface, - "Relax Iterations", - "INPUT", - "NodeSocketInt", - ) - remesh_iterations.default_value = 2 - remesh_iterations.min_value = 1 - remesh_iterations.max_value = 8 smooth_amount = new_socket( interface, "Smooth", @@ -485,7 +485,7 @@ def build_group( remesh_input = new_group_input( nodes, (-100.0, -500.0), - {"Resolution", "Relax Iterations", "Smooth"}, + {"Resolution", "Iterations", "Smooth"}, ) topology_input = new_group_input(nodes, (350.0, -700.0), {"Fix Topology"}) island_input = new_group_input(nodes, (3400.0, 100.0), {"Resolution"}) @@ -497,7 +497,7 @@ def build_group( if include_confidence: confidence_attribute = nodes.new("GeometryNodeInputNamedAttribute") confidence_attribute.data_type = "FLOAT" - confidence_attribute.inputs["Name"].default_value = "o_face_confidence" + confidence_attribute.inputs["Name"].default_value = "o_confidence" confidence_attribute.location = (-1800.0, 600.0) confidence_scale = nodes.new("ShaderNodeMath") confidence_scale.operation = "MULTIPLY" @@ -595,7 +595,7 @@ def build_group( group, culled_geometry, remesh_input.outputs["Resolution"], - remesh_input.outputs["Relax Iterations"], + remesh_input.outputs["Iterations"], topology_input.outputs["Fix Topology"], remesh_input.outputs["Smooth"], (-300.0, 100.0), diff --git a/tools/validate_depth_plane_asset.py b/tools/validate_depth_plane_asset.py index 669e3f6..d100eaa 100644 --- a/tools/validate_depth_plane_asset.py +++ b/tools/validate_depth_plane_asset.py @@ -4,7 +4,12 @@ import sys import bpy -sys.path.insert(0, str(Path(__file__).parents[1] / "src")) +PROJECT_ROOT = Path(__file__).parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "src")) +sys.path.insert( + 0, + str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), +) from snapform.blender.depth_plane import ( DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME, diff --git a/tools/validate_geometry_nodes_asset.py b/tools/validate_geometry_nodes_asset.py index 6ff650d..9846838 100644 --- a/tools/validate_geometry_nodes_asset.py +++ b/tools/validate_geometry_nodes_asset.py @@ -5,7 +5,12 @@ import bpy import numpy as np -sys.path.insert(0, str(Path(__file__).parents[1] / "src")) +PROJECT_ROOT = Path(__file__).parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "src")) +sys.path.insert( + 0, + str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), +) from snapform.blender.depth_plane import ( get_modifier_input, @@ -26,7 +31,7 @@ def set_input(modifier, group, name, value): def add_face_attributes(mesh, count): - confidence = mesh.attributes.new("o_face_confidence", "FLOAT", "FACE") + confidence = mesh.attributes.new("o_confidence", "FLOAT", "FACE") confidence.data.foreach_set("value", [1.0] * count) @@ -48,6 +53,15 @@ def main(): if item.item_type == "SOCKET" and item.in_out == "INPUT" ] assert all(item.structure_type == "SINGLE" for item in input_sockets) + assert [item.name for item in input_sockets[:4]] == [ + "Geometry", + "Iterations", + "Resolution", + "Smooth", + ] + iterations = next(item for item in input_sockets if item.name == "Iterations") + assert iterations.default_value == 2 + assert iterations.min_value == 0 resolution = next(item for item in input_sockets if item.name == "Resolution") assert abs(resolution.default_value - 0.1) < 1e-6 assert group.asset_data is None @@ -123,10 +137,11 @@ def main(): uv = mesh.uv_layers.new(name="UVMap") uv.data.foreach_set("uv", [0, 0, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1]) add_face_attributes(mesh, 2) + assert "o_confidence" in mesh.attributes assert not { "o_background", "o_background_face", - "o_confidence", + "o_face_confidence", "o_confidence_percentile", }.intersection(mesh.attributes.keys()) @@ -137,7 +152,7 @@ def main(): modifier.show_group_selector = False assert not modifier.show_group_selector set_input(modifier, group, "Resolution", 0.000001) - set_input(modifier, group, "Relax Iterations", 1) + set_input(modifier, group, "Iterations", 0) set_input(modifier, group, "Fix Topology", 0) set_input(modifier, group, "Smooth", 0.0) evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) diff --git a/tools/validate_moge_depth_plane.py b/tools/validate_moge_depth_plane.py index 7ab0427..da0d512 100644 --- a/tools/validate_moge_depth_plane.py +++ b/tools/validate_moge_depth_plane.py @@ -8,6 +8,10 @@ import numpy as np PROJECT_ROOT = Path(__file__).parents[1] sys.path.insert(0, str(PROJECT_ROOT / "src")) +sys.path.insert( + 0, + str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), +) from snapform.blender.depth_plane import ( MOGE_DEPTH_PLANE_NODE_GROUP_NAME, diff --git a/tools/validate_moge_surface.py b/tools/validate_moge_surface.py index b25fba1..fef8ae4 100644 --- a/tools/validate_moge_surface.py +++ b/tools/validate_moge_surface.py @@ -8,13 +8,37 @@ import numpy as np PROJECT_ROOT = Path(__file__).parents[1] sys.path.insert(0, str(PROJECT_ROOT / "src")) +sys.path.insert( + 0, + str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), +) from snapform.blender.depth_plane import MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME from snapform.blender.surface import ( MOGE_NODE_GROUP_NAME, build_image_mesh, - prepare_moge_surface_archive, ) +from snapform.server.geometry import GeometryFrame, build_surface +from snapform.server.outputs.surface import write_surface_archive + + +def write_test_normal_texture(normal, output_path): + converted = normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + rgb = np.clip((converted + 1.0) * 0.5, 0.0, 1.0) + alpha = np.ones((*rgb.shape[:2], 1), dtype=np.float32) + pixels = np.concatenate((rgb, alpha), axis=-1) + image = bpy.data.images.new( + "MoGe Test Normal", + width=rgb.shape[1], + height=rgb.shape[0], + alpha=True, + float_buffer=False, + ) + image.pixels.foreach_set(pixels.reshape(-1)) + image.file_format = "PNG" + image.filepath_raw = str(output_path) + image.save() + return output_path def main(): @@ -37,26 +61,38 @@ def main(): dtype=np.float32, ) image = bpy.data.images.new("MoGe Surface Test", width=width, height=height) + foreground = bpy.data.images.new( + "MoGe Surface Test.foreground", + width=width, + height=height, + ) source = bpy.data.objects.new("MoGe Surface Test", None) source.empty_display_type = "IMAGE" source.data = image bpy.context.collection.objects.link(source) with tempfile.TemporaryDirectory() as temporary: directory = Path(temporary) - results_path = directory / "results.npz" mesh_path = directory / "mesh.npz" - np.savez( - results_path, - depth=depth, - mask=mask, - points=points, - intrinsics=intrinsics, - normal=np.broadcast_to( - np.array([0.0, 0.0, -1.0], dtype=np.float32), - (1, height, width, 3), - ), + normal = np.broadcast_to( + np.array([0.0, 0.0, -1.0], dtype=np.float32), + (height, width, 3), + ) + pixel_intrinsics = intrinsics[0].copy() + pixel_intrinsics[0] *= width + pixel_intrinsics[1] *= height + pixel_intrinsics[2] = (0.0, 0.0, 1.0) + frame = GeometryFrame( + depth=depth[0], + points=points[0], + intrinsics=pixel_intrinsics, + valid_mask=np.isfinite(depth[0]) & (depth[0] > 0.0), + model_mask=mask[0], + ) + write_surface_archive(build_surface(frame), mesh_path) + normal_path = write_test_normal_texture( + normal, + directory / "normal.png", ) - prepare_moge_surface_archive(results_path, mesh_path) surface = build_image_mesh( bpy.context, source, @@ -65,8 +101,11 @@ def main(): True, False, True, + color_image=foreground, + normal_path=normal_path, ) assert len(surface.data.vertices) > 0 + assert all(len(polygon.vertices) == 3 for polygon in surface.data.polygons) assert surface.modifiers[0].node_group.name == MOGE_NODE_GROUP_NAME inputs = { item.name @@ -88,7 +127,7 @@ def main(): assert not { "o_confidence", "o_confidence_percentile", - "o_face_confidence", + "o_confidence", "o_background", "o_background_face", }.intersection(attribute_names) @@ -134,7 +173,7 @@ def main(): assert smooth_socket.default_value == 0.5 assert not any( node.bl_idname == "GeometryNodeInputNamedAttribute" - and node.inputs["Name"].default_value == "o_face_confidence" + and node.inputs["Name"].default_value == "o_confidence" for node in surface.modifiers[0].node_group.nodes ) assert MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME in bpy.data.node_groups @@ -154,6 +193,11 @@ def main(): assert normal_strength.min_value == 0.0 assert normal_map.inputs["Strength"].is_linked assert not surface.data.has_custom_normals + color_texture = next( + node for node in surface.data.materials[0].node_tree.nodes + if node.bl_idname == "ShaderNodeTexImage" and node.image == foreground + ) + assert color_texture.image.name == "MoGe Surface Test.foreground" normal_texture = next( node for node in surface.data.materials[0].node_tree.nodes if node.bl_idname == "ShaderNodeTexImage" @@ -172,7 +216,10 @@ def validate_real_result(results_path, image_path): bpy.context.collection.objects.link(source) mesh_path = results_path.with_name("mesh.npz") if not mesh_path.is_file(): - prepare_moge_surface_archive(results_path, mesh_path) + raise ValueError(f"Surface archive is missing: {mesh_path}") + normal_path = results_path.with_name("normal.png") + if not normal_path.is_file(): + raise ValueError(f"Surface normal texture is missing: {normal_path}") surface = build_image_mesh( bpy.context, source, @@ -181,6 +228,7 @@ def validate_real_result(results_path, image_path): True, False, True, + normal_path=normal_path, ) assert len(surface.data.vertices) > 100 assert len(surface.data.polygons) > 100 @@ -202,7 +250,7 @@ def validate_real_result(results_path, image_path): assert not { "o_confidence", "o_confidence_percentile", - "o_face_confidence", + "o_confidence", "o_background", "o_background_face", }.intersection(attribute_names)