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.
This commit is contained in:
18
README.md
18
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 逐帧保存。
|
||||
|
||||
## 任务结果保存在哪里
|
||||
|
||||
|
||||
@@ -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。
|
||||
|
||||
@@ -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` 资产；修改节点组、材质组或其接口时，应运行相应资产验证脚本，并再次运行完整测试。
|
||||
|
||||
## 上游与许可证
|
||||
|
||||
@@ -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`。
|
||||
|
||||
@@ -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)：模块边界、调用流程、项目结构、常用命令与许可证
|
||||
|
||||
@@ -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 保留输入图像颜色。
|
||||
|
||||
@@ -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 解码掩盖）。
|
||||
|
||||
|
||||
@@ -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
|
||||
<opaque-job-id>/
|
||||
├── scene.glb、*.ply、*.jpg 或 results.npz
|
||||
└── 其它对应 Debug 格式的产物
|
||||
├── results.npz、mesh.npz、texture.png 或 normal.png
|
||||
└── prediction.npz Debug 原始 API 返回
|
||||
|
||||
<opaque-job-id>/
|
||||
├── foreground.png 单图结果
|
||||
|
||||
Binary file not shown.
@@ -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"],
|
||||
|
||||
@@ -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=(
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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"}
|
||||
@@ -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)
|
||||
@@ -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
|
||||
35
src/snapform/server/geometry/__init__.py
Normal file
35
src/snapform/server/geometry/__init__.py
Normal file
@@ -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",
|
||||
)
|
||||
202
src/snapform/server/geometry/grid_mesh.py
Normal file
202
src/snapform/server/geometry/grid_mesh.py
Normal file
@@ -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)
|
||||
|
||||
31
src/snapform/server/geometry/projection.py
Normal file
31
src/snapform/server/geometry/projection.py
Normal file
@@ -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)
|
||||
133
src/snapform/server/geometry/surface.py
Normal file
133
src/snapform/server/geometry/surface.py
Normal file
@@ -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,
|
||||
)
|
||||
|
||||
5
src/snapform/server/jobs/__init__.py
Normal file
5
src/snapform/server/jobs/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from .depth import run as generate_depth
|
||||
from .surface import run as generate_surface
|
||||
|
||||
|
||||
__all__ = ("generate_depth", "generate_surface")
|
||||
78
src/snapform/server/jobs/background.py
Normal file
78
src/snapform/server/jobs/background.py
Normal file
@@ -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"}
|
||||
20
src/snapform/server/jobs/common.py
Normal file
20
src/snapform/server/jobs/common.py
Normal file
@@ -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)
|
||||
26
src/snapform/server/jobs/debug_ben2.py
Normal file
26
src/snapform/server/jobs/debug_ben2.py
Normal file
@@ -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": "."}
|
||||
21
src/snapform/server/jobs/debug_da3.py
Normal file
21
src/snapform/server/jobs/debug_da3.py
Normal file
@@ -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": "."}
|
||||
26
src/snapform/server/jobs/debug_moge2.py
Normal file
26
src/snapform/server/jobs/debug_moge2.py
Normal file
@@ -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": "."}
|
||||
49
src/snapform/server/jobs/depth.py
Normal file
49
src/snapform/server/jobs/depth.py
Normal file
@@ -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"}
|
||||
51
src/snapform/server/jobs/foreground.py
Normal file
51
src/snapform/server/jobs/foreground.py
Normal file
@@ -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
|
||||
21
src/snapform/server/jobs/frames.py
Normal file
21
src/snapform/server/jobs/frames.py
Normal file
@@ -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)}
|
||||
75
src/snapform/server/jobs/surface.py
Normal file
75
src/snapform/server/jobs/surface.py
Normal file
@@ -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"}
|
||||
@@ -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
|
||||
26
src/snapform/server/media/__init__.py
Normal file
26
src/snapform/server/media/__init__.py
Normal file
@@ -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",
|
||||
)
|
||||
158
src/snapform/server/media/input.py
Normal file
158
src/snapform/server/media/input.py
Normal file
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
8
src/snapform/server/models/__init__.py
Normal file
8
src/snapform/server/models/__init__.py
Normal file
@@ -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")
|
||||
78
src/snapform/server/models/ben2.py
Normal file
78
src/snapform/server/models/ben2.py
Normal file
@@ -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
|
||||
94
src/snapform/server/models/da3.py
Normal file
94
src/snapform/server/models/da3.py
Normal file
@@ -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],
|
||||
)
|
||||
119
src/snapform/server/models/moge2.py
Normal file
119
src/snapform/server/models/moge2.py
Normal file
@@ -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],
|
||||
)
|
||||
@@ -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)
|
||||
11
src/snapform/server/outputs/__init__.py
Normal file
11
src/snapform/server/outputs/__init__.py
Normal file
@@ -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",
|
||||
)
|
||||
24
src/snapform/server/outputs/debug_da3.py
Normal file
24
src/snapform/server/outputs/debug_da3.py
Normal file
@@ -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
|
||||
13
src/snapform/server/outputs/debug_moge2.py
Normal file
13
src/snapform/server/outputs/debug_moge2.py
Normal file
@@ -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
|
||||
64
src/snapform/server/outputs/depth.py
Normal file
64
src/snapform/server/outputs/depth.py
Normal file
@@ -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"
|
||||
)
|
||||
44
src/snapform/server/outputs/surface.py
Normal file
44
src/snapform/server/outputs/surface.py
Normal file
@@ -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
|
||||
@@ -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),
|
||||
):
|
||||
|
||||
@@ -12,6 +12,7 @@ SOURCE_DIR = (
|
||||
/ "src"
|
||||
/ "snapform"
|
||||
/ "server"
|
||||
/ "models"
|
||||
)
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ SOURCE_DIR = (
|
||||
/ "src"
|
||||
/ "snapform"
|
||||
/ "server"
|
||||
/ "models"
|
||||
)
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
for export_format in debug_formats:
|
||||
context = FakeJobContext(
|
||||
root / export_format,
|
||||
"debug-da3",
|
||||
root / "result",
|
||||
"generate-depth",
|
||||
)
|
||||
self.da3.generate_da3(
|
||||
self.depth_job.run(
|
||||
context,
|
||||
{
|
||||
"input": str(input_path),
|
||||
"model": "DA3",
|
||||
"video_frames": 16,
|
||||
"export_format": export_format,
|
||||
"remove_background_model": "BEN2",
|
||||
"video_frames": 1,
|
||||
"export_format": "mini_npz",
|
||||
"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,
|
||||
"process_res_method": "upper_bound_resize",
|
||||
"ref_view_strategy": "first",
|
||||
"inference_mode": "per_frame",
|
||||
},
|
||||
)
|
||||
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()
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
),
|
||||
(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)
|
||||
|
||||
Reference in New Issue
Block a user