From 13ef7052260dc5d469da2e9f59a2d8055c2e4669 Mon Sep 17 00:00:00 2001 From: nan Date: Thu, 20 Aug 2026 01:56:04 +0800 Subject: [PATCH] Add operators for model downloading and surface generation - Implemented `DownloadModel` and `DownloadRequiredModels` operators for downloading models. - Created `ClearModels` operator to clear model resources. - Added `ConvertToSurface` operator for converting image empties into textured surfaces. - Introduced `GenerateDa3Surface` and `GenerateMoge2Surface` operators for generating surfaces using DA3 and MoGe-2 models respectively. - Developed utility functions for handling image mesh data, assigning attributes, and managing geometry nodes. - Established runtime configuration and post-installation model downloading in `runtime.py`. --- README.md | 15 +- docs/depth-plane.md | 2 +- docs/development.md | 65 +- docs/features.md | 15 + docs/server.md | 2 +- docs/surface.md | 4 +- src/snapform/__init__.py | 25 +- src/snapform/blender/__init__.py | 1 - src/snapform/blender/depth_plane.py | 2271 ----------------- src/snapform/menu.py | 76 + src/snapform/operators.py | 735 ------ src/snapform/operators/__init__.py | 34 + src/snapform/operators/background.py | 112 + src/snapform/operators/balloon/__init__.py | 1 + src/snapform/operators/balloon/geometry.py | 346 +++ src/snapform/operators/balloon/operator.py | 456 ++++ src/snapform/operators/balloon/poisson.py | 123 + src/snapform/operators/balloon/teddy.py | 218 ++ src/snapform/operators/common.py | 111 + src/snapform/operators/debug.py | 91 + .../operators/depth_plane/__init__.py | 7 + src/snapform/operators/depth_plane/build.py | 546 ++++ src/snapform/operators/depth_plane/data.py | 262 ++ .../operators/depth_plane/operator.py | 481 ++++ src/snapform/operators/image.py | 451 ++++ src/snapform/operators/material.py | 693 +++++ src/snapform/operators/models.py | 48 + .../{blender => operators}/surface.py | 169 +- src/snapform/panel.py | 2 +- src/snapform/preferences.py | 2 +- src/snapform/properties.py | 2 +- src/snapform/{job_runtime.py => runtime.py} | 2 +- tests/test_addon.py | 390 ++- tests/test_image_depth.py | 71 +- tests/test_pack.py | 6 +- tools/validate_depth_plane_asset.py | 8 +- tools/validate_geometry_nodes_asset.py | 7 +- tools/validate_moge_depth_plane.py | 8 +- tools/validate_moge_surface.py | 9 +- 39 files changed, 4729 insertions(+), 3138 deletions(-) delete mode 100644 src/snapform/blender/__init__.py delete mode 100644 src/snapform/blender/depth_plane.py create mode 100644 src/snapform/menu.py delete mode 100644 src/snapform/operators.py create mode 100644 src/snapform/operators/__init__.py create mode 100644 src/snapform/operators/background.py create mode 100644 src/snapform/operators/balloon/__init__.py create mode 100644 src/snapform/operators/balloon/geometry.py create mode 100644 src/snapform/operators/balloon/operator.py create mode 100644 src/snapform/operators/balloon/poisson.py create mode 100644 src/snapform/operators/balloon/teddy.py create mode 100644 src/snapform/operators/common.py create mode 100644 src/snapform/operators/debug.py create mode 100644 src/snapform/operators/depth_plane/__init__.py create mode 100644 src/snapform/operators/depth_plane/build.py create mode 100644 src/snapform/operators/depth_plane/data.py create mode 100644 src/snapform/operators/depth_plane/operator.py create mode 100644 src/snapform/operators/image.py create mode 100644 src/snapform/operators/material.py create mode 100644 src/snapform/operators/models.py rename src/snapform/{blender => operators}/surface.py (79%) rename src/snapform/{job_runtime.py => runtime.py} (94%) diff --git a/README.md b/README.md index df6cf87..92fa071 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 ## 它能做什么 +- 从 Image Empty 套索提取不规则区域，并通过 Mode 选择 Poisson 或 Teddy / Chordal Axis 生成封闭 Balloon 网格；源 Empty 保持不变 - 一键把 Image Empty 转换为立体面片：让 AI 估算深度，或直接用原图作为深度 - 把单张 Image Empty 转换为透视纹理网格，并自动匹配原对象的 Transform 与 Offset - 一键去除图片背景，生成透明前景图 @@ -32,6 +33,7 @@ Environment Bootstrap 使用连续的 0–100%：uv 为 1–20%、Python 为 20 时显示 `Download Required Models`，一次补齐缺失项，不会重复下载已经就绪的模型。 Environment 未安装时，同一位置显示两个安装按钮。Environment 或必需模型未就绪时， 菜单仍显示 `Convert to Depth Plane`、`Convert to Surface` 和 `Remove Background`，但三个入口都会保持禁用。 +`Balloon Mesh (Lasso)` 是 Blender 本地几何功能，不依赖 Environment 或模型，在上述状态下仍可使用。 这些 UI 状态直接检查 Storage Root 的 `models/` 文件，不会仅因打开菜单或面板而启动 Job Server。 ## 快速上手 @@ -48,6 +50,17 @@ Environment 未安装时，同一位置显示两个安装按钮。Environment - `Keep Original` 默认关闭，位于分割线下方；开启后保留源 Image Empty 4. 点击确认，稍等片刻就得到带位移的立体面片。`O DA3 Depth Plane` 与 `O MoGe2 Depth Plane` 都使用 `.color` 的 `Alpha > 0.5` 划分前景范围，范围外由 `Background Depth` 控制。 +### 套索生成 Balloon + +选中单张图片的 Image Empty，右键打开 `SnapForm` 子菜单并选择 `Balloon Mesh (Lasso)`。按住左键沿目标轮廓绘制套索，松开左键即由独立的 `Extract Balloon` Operator 生成带原图 UV 的封闭网格，同时继续停留在 Lasso 状态等待下一次圈选。每次生成后源 Image Empty 都保持选中和 Active；中键、滚轮、触控板、NDOF 与 Alt+左键事件会交还 3D View，可在持续圈选期间旋转、缩放和平移视角。`Esc` 结束会话，右键只取消当前轨迹。 + +生成后可在左下角 `Adjust Last Operation` 调整： + +- `Mode`：在全局 Poisson 场和 Teddy / Chordal Axis 局部椭圆截面之间切换 +- `Thickness`：控制鼓起厚度，软上限为 `5.0`，不设硬上限 +- `Resolution`：相邻采样点的目标世界空间距离，类型为 Distance，默认 `0.1 Blender unit`；数值越小网格越密 + +每次释放产生的 `Extract Balloon` 都会独立完成并成为 Last Operation，因此持续 Lasso 尚未退出时即可在左下角调整当前 Balloon；每个结果也拥有独立的 Undo 步骤。调整参数会基于该次套索和保存的源对象 Matrix World 重建，不会跳回原点。两种模式共用世界距离采样、平滑边界、`.color` 图片副本和 `O Texture Layer` 材质。该流程不运行 BEN2 或任何其他 AI，也不需要安装 Environment。当前只支持静态图片，不处理 Movie 或 Image Sequence。 ### 去除背景 选中 Image Empty 后，右键打开 `SnapForm` 子菜单并选择 `Remove Background`，完成后 Empty 会自动 @@ -106,6 +119,6 @@ Session 逐帧推理，显存峰值不随视频帧数平方增长。Debug 只保 ## 技术文档 -Blender 侧只负责界面、任务提交与结果导入，AI 与文件输出由独立 Job Server 执行。详细功能见 +Blender 侧只负责界面、任务提交与结果导入，AI 与文件输出由独立 Job Server 执行。源码将插件级 `runtime.py`、`menu.py`、`panel.py` 保持在顶层，具体功能与共享图片、材质能力统一放在 `operators/`；每种任务的结果回写与对应 Operator 放在一起。详细功能见 [文档索引](docs/index.md)，模块边界、调用流程、项目结构与开发命令见 [程序设计、开发与许可证](docs/development.md)。CI 在 `main` push 与 Pull Request 上运行测试，仅由 `v*` tag 构建并发布扩展包；文档在 `main` 的文档相关文件变化时部署。 diff --git a/docs/depth-plane.md b/docs/depth-plane.md index 80fb0ed..39eeb2f 100644 --- a/docs/depth-plane.md +++ b/docs/depth-plane.md @@ -32,7 +32,7 @@ Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty 一致；Original 模式的视频或序列帧使用 `color.0001.*` 颜色层同时作为颜色与深度输入；开启 Remove Background 时同名颜色层包含 BEN2 修改后的 Alpha。 - DA3 的深度仍以 EXR 序列帧（`depth.0001.exr` 起）保存在 job 目录根下；MoGe-2 分别使用 `moge_depth`、`moge_mask` 与 `moge_normal` EXR 序列。 + DA3 动画深度以 EXR 序列帧（`depth.0001.exr` 起）保存在 Job 目录根下；MoGe-2 动画只为置换数据创建 `moge_depth.0001.exr` 起的 EXR 序列。MoGe-2 的 `results.npz` 保留完整 depth、normal 与 mask 数组，Server 另外输出对应 PNG 预览。 DA3 与 MoGe-2 的单帧推理进度只显示模型与阶段名称，不追加 `1/1`；多帧时继续显示当前帧与总数。`Shadeless`、`Remove Background` 与 `Keep Original` 均默认关闭，因此默认删除原对象。`Shadeless` 会以 `O Shadeless` 节点组替代普通 Principled BSDF 路径，并继续通过 Alpha 与 Transparent BSDF 混合。 3. DA3 或 MoGe-2 完成后会自动创建面片。DA3 面片几何先由 `mini_npz` 的相机内参与深度确定，再通过对象 Transform 匹配原 Image Empty 的可见尺寸、位置、旋转和 Offset。Original 模式直接使用源 Empty diff --git a/docs/development.md b/docs/development.md index 107e1db..0aba0d2 100644 --- a/docs/development.md +++ b/docs/development.md @@ -19,10 +19,13 @@ flowchart LR 核心边界如下： - `src/snapform/__init__.py` 是注册入口，集中注册 Blender 类型、`Scene.snapform_settings`、右键菜单与 Runtime -- `panel.py`、`preferences.py`、`properties.py` 只描述 UI 和 Blender 属性，不直接导入 Server 模型实现 -- `operators.py` 将 Blender 状态转换为普通字典参数；Job 完成后在主线程导入图片、网格、材质和相机 -- `blender/depth_plane.py` 与 `blender/surface.py` 实现 Image Empty 输入检查、几何构建和资产挂载 -- `job_runtime.py` 以入口字符串声明独立 Server，并配置 Environment、Storage Root 与首次安装后的推荐模型下载 +- `menu.py` 与 `panel.py` 是同级 Blender UI 入口；`preferences.py`、`properties.py` 分别保存插件配置和 RNA 状态，它们都不直接导入 Server 模型实现 +- `runtime.py` 创建唯一 `JobRuntime`，配置 Environment、Storage Root、Server 入口与首次安装后的推荐模型下载 +- `operators/` 按功能组织 Blender 操作。Depth Plane、Surface、Background 与 Balloon 各自包含入口 Operator、Job 请求和结果回写，不设置跨功能的 `results.py` +- `operators/common.py` 只保留跨 Job Operator 共用的参数校验、模型就绪检查、临时输入清理和请求参数构造；`operators/image.py` 与 `operators/material.py` 提供各功能 Operator 共用的 Image Empty、序列和材质操作，三者同层但不互相混合职责 +- `operators/models.py` 与 `operators/debug.py` 分别承载模型管理和 Debug 操作 +- `operators/balloon/` 以持续 `BalloonLasso` 控制器分发单次、可独立 Redo/Undo 的 `ExtractBalloon`；`geometry.py` 负责共用拓扑，`poisson.py` 与 `teddy.py` 分别计算两种厚度场 +- `operators/depth_plane/` 将转换入口、Job 请求、响应回写、深度数据和网格构建收拢在同一功能包；`operators/surface.py` 与 `operators/background.py` 采用相同的功能内聚原则 - `server/__init__.py` 创建 `JobServer`、注册公开 Job，并把 `ModelManager` 注册为 Resource - `server/jobs/` 按任务组织完整流程；DA3、MoGe-2、BEN2 Debug 各自拥有独立模块和 Job，不用模型分支总控函数 - `server/models/` 包含三个模型适配器及其 ONNX 实现，只负责模型载入、推理和 Prediction；不接收 Job Context、不写文件 @@ -48,52 +51,40 @@ flowchart LR │ └── snapform/ │ ├── __init__.py │ ├── blender_manifest.toml -│ ├── job_runtime.py -│ ├── operators.py +│ ├── runtime.py +│ ├── menu.py │ ├── panel.py │ ├── preferences.py │ ├── properties.py +│ ├── model_catalog.py │ ├── assets/ │ │ └── O_SnapForm.blend -│ ├── blender/ -│ │ ├── depth_plane.py -│ │ └── surface.py -│ ├── model_catalog.py +│ ├── operators/ +│ │ ├── common.py +│ │ ├── image.py +│ │ ├── material.py +│ │ ├── models.py +│ │ ├── debug.py +│ │ ├── background.py +│ │ ├── surface.py +│ │ ├── depth_plane/ +│ │ │ ├── operator.py +│ │ │ ├── data.py +│ │ │ └── build.py +│ │ └── balloon/ +│ │ ├── operator.py +│ │ ├── geometry.py +│ │ ├── poisson.py +│ │ └── teddy.py │ ├── server/ │ │ ├── __init__.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 │ │ ├── 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 +│ └── blendjob-*.whl ├── tests/ └── tools/ ``` diff --git a/docs/features.md b/docs/features.md index 3e742cb..a9de734 100644 --- a/docs/features.md +++ b/docs/features.md @@ -31,6 +31,7 @@ - Debug 与 Convert to Depth Plane 的模型选项直接显示下载状态；未下载模型带 `(Missing)`，已下载模型保持原名 - Debug 不运行点云、Surface、深度预览、置信度 percentile 或坐标变换等 SnapForm 后处理 +- Image Empty 右键菜单提供单一 Balloon 套索提取，并可在 Last Operation 中切换 Poisson 与 Teddy / Chordal Axis Mode；无需 Environment 或模型，源 Empty 保留 - Image Empty 右键菜单可直接运行 DA3，根据 `mini_npz` 深度生成面片 - Image Empty 右键菜单可直接运行 BEN2，保留原图并换用新建的透明背景图片 - 缺少必需的 `MoGe-2 ViT-S Normal` 或 `BEN2 Base` 时，Image Empty 的 `SnapForm` @@ -82,6 +83,20 @@ 完成后会打开 job 文件夹，其中的 `foreground.png` 带透明 Alpha；状态栏显示 BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。 +## Extract Balloon + +`Balloon Mesh (Lasso)` 是纯 Blender 侧入口，不提交 Job，也不调用 BEN2、深度模型或其他 AI，仅支持静态 Image Empty。它启动不进入 Undo/Redo 历史的持续 `BalloonLasso` 控制器；按住左键绘制套索，松开时控制器调用一次独立的 `ExtractBalloon` Operator 并立即完成，随后继续等待下一次圈选。`Esc` 结束控制器，右键只清除当前未完成轨迹。屏幕轨迹通过视图射线与 Empty 图片平面求交，再结合图片宽高比、Display Size、Image Offset 与 Object Transform 映射为原图像素坐标。 + +Operator 的 `Mode` 保留两种厚度算法： + +- `Poisson` 在共用二维网格上组装 cotangent Laplacian 与面积载荷，求解零边界 Poisson 场；高度取场值平方根并按 `Thickness` 缩放 +- `Teddy / Chordal Axis` 依据 [Teddy: A Sketching Interface for 3D Freeform Design](https://www.cs.princeton.edu/courses/archive/fall00/cs597b/papers/teddy-siggraph99.pdf) 的思路，仅对平滑轮廓做约束 Delaunay，形成袖段、分叉和端点轴，并用轴半径生成局部圆/椭圆式截面。当前实现将厚度场映射到共用均匀网格，并非逐面复刻论文的专用缝合拓扑 + +每次释放都会把该次套索、源对象名称和启动时的完整 Matrix World 传给一个带 `REGISTER`、`UNDO` 的 `ExtractBalloon`。它完成后立即成为 Last Operation，即使持续 Lasso 仍在运行，也可在左下角修改当前结果的 `Mode`、`Thickness` 和 `Resolution`；每次提取分别拥有 Undo 步骤。重建后会重新把源 Image Empty 设为选中和 Active，便于继续圈选，结果也不会跳回原点。Modal 对中键、滚轮、触控板、NDOF、Alt+左键、HUD 和其他未处理事件返回 `PASS_THROUGH`，所以持续状态不阻断 3D View 导航或参数面板交互。`Thickness` 默认 `1.0`、软上限 `5.0` 且无硬上限。`Resolution` 是 subtype 为 `DISTANCE` 的长度属性，直接表示相邻采样点的目标世界空间距离，默认 `0.1 Blender unit`；数值越小网格越密，不再使用密度倍率换算。 + +两种模式共用轮廓平滑、按弧长等距重采样和二维约束 Delaunay。边界与内部点以 `Resolution` 的世界距离采样，不随图片分辨率、Display Size 或对象 Scale 改变目标间距。前后表面在同一条零厚度边缘共享顶点，不生成直立侧壁，结果为封闭网格。UV 保持原图坐标；材质使用独立的 `O Texture Layer`，并复制源 Image datablock、按 Convert 系列相同规则命名为 `.color`。创建 Object 与 Mesh 时直接使用源 Image Empty 名称，不主动追加 Balloon 后缀；同一 Blender datablock 集合已有重名时仍由 Blender 自动追加数字编号。源 Image Empty 保持不变，Undo 可移除提取结果。 + +当前不提供添加/减去选区、边缘羽化和视频跟踪。 ## Image Empty 去背景 1. 确认 `Models` panel 中的 `BEN2 Base` 已就绪；首次安装 Environment 时会默认下载。 diff --git a/docs/server.md b/docs/server.md index 87c7f78..f6aeb28 100644 --- a/docs/server.md +++ b/docs/server.md @@ -107,7 +107,7 @@ Environment Bootstrap 的单次 0–100% 只包含 uv、Python 与 packages， AI、帧提取和 Models panel 下载的 Job 进度保存在 Server 内存中。Handler 只使用 `context.progress(value, message)` 发布进度，终态由 JobServer 内部管理；SnapForm 不再维护独立的 `server/progress.py` 或 Job 状态 JSON。Blender 每 0.25 秒轮询 `GET /jobs/{job_id}`，任务目录不写 `status.json`。Operator 直接提交参数字典；Job 被接受时由 `JobServer` 自动生成 ID 并创建 `jobs//`，提交响应同时返回 `job_id` 与 `directory`。路径通过 `JobContext.directory` 传给 Handler，SnapForm Server 强制将输出写入该目录。 -`job_runtime.py` 创建唯一的 `runtime = JobRuntime(...)`，以自身目录为 `entrypoint_root`，传入 `server:server` 相对入口字符串与显式 `snapform` namespace，并导出 `runtime.JobOperatorBase`。Runtime 独占 Storage Root、Environment 路径、安装进程、Server 控制和 Status Bar。Operator 未覆盖 `request(context)` 时会提交全部 Property；覆盖后返回值纯替换默认参数。可选的 `response(context, result)` 在 Blender 主线程直接完成结果导入。SnapForm 不再提供 Client 请求层、Job Mixin、Status Bar 或 Build 中转 Operator。 +`runtime.py` 创建唯一的 `runtime = JobRuntime(...)`，以 `src/snapform/` 为 `entrypoint_root`，传入 `server:server` 相对入口字符串与显式 `snapform` namespace，并导出 `runtime.JobOperatorBase`。Runtime 独占 Storage Root、Environment 路径、安装进程、Server 控制和 Status Bar。Operator 未覆盖 `request(context)` 时会提交全部 Property；覆盖后返回值纯替换默认参数。可选的 `response(context, result)` 在 Blender 主线程直接完成结果导入。SnapForm 不再提供 Client 请求层、Job Mixin、Status Bar 或 Build 中转 Operator。 Environment 安装是 Server 启动前唯一的 Bootstrap 流程：`environment.json` 记录期望安装的环境声明，`install-status.json` 在安装期间回传进度并在任务收尾后删除，`install.log` 持久记录安装输出。Bootstrap 不创建 Job 目录。Environment 成功后执行普通的 `post_install(runtime)`；默认模型下载此时已经是正式 Server Job，会创建自己的 Job 目录，并使用与其它任务相同的进度和取消协议。 diff --git a/docs/surface.md b/docs/surface.md index 69e8559..b9bfb9b 100644 --- a/docs/surface.md +++ b/docs/surface.md @@ -60,8 +60,8 @@ 以及面域 `o_mask`、`o_depth_discontinuity`，不伪造 confidence、stretch 或 background 属性， 也不写入 custom normal。MoGe-2 推理返回的 normal 直接编码为独立的 `normal.png`，Blender 将它 作为 Non-Color 贴图导入现有材质；normal 不进入 `mesh.npz`。 - 所有 MoGe-2 Image datablock 都以源图片名而非对象名为基名：depth、normal 与 mask 分别使用 - `图片名.depth`、`图片名.normal`、`图片名.mask`，不附加 `.moge2`。 + MoGe-2 Surface 独立导入的 normal Image datablock 以源图片名而非对象名为基名，使用 + `图片名.normal`，不附加 `.moge2`。 7. 网格以 `reference_depth` 截面匹配 Image Empty 的显示矩形：该截面的投影四角与 Empty 的可见 边界重合，主点同时匹配 Empty Offset。网格对象直接沿用源 Empty 的 `matrix_world`，因此位置、旋转、 缩放及父级关系产生的世界变换保持一致。更近的顶点沿 Empty 局部 `+Z` 方向凸出，三角面正面朝向 diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index 2b66022..688fdaa 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -1,36 +1,25 @@ import bpy -from .blender.depth_plane import ( - ConvertToDepthPlane, - RemoveImageBackground, - SnapFormImageMenu, - draw_image_context_menu, -) -from .blender.surface import ConvertToSurface -from .operators import ( - JOB_OPERATOR_CLASSES, -) -from .job_runtime import runtime +from .runtime import runtime +from .operators import CLASSES as OPERATOR_CLASSES +from .preferences import SnapFormPreferences +from .properties import SnapFormSettings +from .menu import SnapFormImageMenu, draw_image_context_menu from .panel import ( Ben2RemovalDebugPanel, - ServerPanel, Da3OutputDebugPanel, DebugPanel, ModelsPanel, Moge2ApiDebugPanel, + ServerPanel, ) -from .preferences import SnapFormPreferences -from .properties import SnapFormSettings CLASSES = ( SnapFormPreferences, SnapFormSettings, - *JOB_OPERATOR_CLASSES, + *OPERATOR_CLASSES, SnapFormImageMenu, - RemoveImageBackground, - ConvertToDepthPlane, - ConvertToSurface, ServerPanel, ModelsPanel, DebugPanel, diff --git a/src/snapform/blender/__init__.py b/src/snapform/blender/__init__.py deleted file mode 100644 index c741b19..0000000 --- a/src/snapform/blender/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Blender-process geometry construction and import logic.""" diff --git a/src/snapform/blender/depth_plane.py b/src/snapform/blender/depth_plane.py deleted file mode 100644 index 2817ad7..0000000 --- a/src/snapform/blender/depth_plane.py +++ /dev/null @@ -1,2271 +0,0 @@ -import re -import shutil -import tempfile -from pathlib import Path - -import bpy - - -DEPTH_KEY_CANDIDATES = ("depth", "depths") -INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts") -DISPLACEMENT_STRENGTH = 1.0 -DISPLACEMENT_MIDLEVEL = 0.0 -ORIGINAL_DISPLACEMENT_STRENGTH = 1.0 -ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0 -BUMP_BASE_DISTANCE = 0.001 -DEPTH_PLANE_NODE_GROUP_NAME = "O DA3 Depth Plane" -MOGE_DEPTH_PLANE_NODE_GROUP_NAME = "O MoGe2 Depth Plane" -NODE_ASSET_NAME = "O_SnapForm.blend" -MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement" -BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth" -DEPTH_SCALE_ATTRIBUTE = "o_depth_scale" -SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer" -MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Surface Layer" -DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer" -MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer" -SHADELESS_NODE_GROUP_NAME = "O Shadeless" -VIEW_TRANSFORM_COLOR_SPACES = ( - ("ACES 1.3", "ACES 1.3 sRGB"), - ("ACES 2.0", "ACES 2.0 sRGB"), - ("AgX", "AgX Base sRGB"), - ("Filmic", "Filmic sRGB"), - ("Khronos PBR Neutral", "Khronos PBR Neutral sRGB"), -) -DA3_PATCH_SIZE = 14 -DA3_MIN_PROCESS_RES = 252 -DA3_MAX_PROCESS_RES = 1008 -IMAGE_NAME_SUFFIXES = ( - ".jpeg", - ".jpg", - ".png", - ".webp", - ".tiff", - ".tif", - ".exr", - ".hdr", - ".bmp", - ".tga", -) - - -VIDEO_NAME_SUFFIXES = ( - ".mp4", - ".mov", - ".avi", - ".mkv", - ".webm", - ".m4v", - ".mts", - ".mpg", - ".mpeg", -) - - -def image_depth_model_items(owner, context): - from ..properties import depth_plane_model_enum_items - - return depth_plane_model_enum_items(owner, context) - - -def is_image_empty(candidate): - return bool( - candidate - and getattr(candidate, "type", None) == "EMPTY" - and getattr(candidate, "empty_display_type", None) == "IMAGE" - and getattr(candidate, "data", None) is not None - ) - - -def material_color_space_name(scene): - view_settings = getattr(scene, "view_settings", None) - view_transform = getattr(view_settings, "view_transform", "") - normalized = str(view_transform).strip().casefold() - for view_name, color_space in VIEW_TRANSFORM_COLOR_SPACES: - if normalized.startswith(view_name.casefold()): - return color_space - return "sRGB" - - -def set_material_color_space(image, scene=None): - if scene is None: - scene = getattr(bpy.context, "scene", None) - settings = image.colorspace_settings - preferred = material_color_space_name(scene) - original = settings.name - for color_space in dict.fromkeys((preferred, "sRGB", original)): - if not color_space: - continue - try: - settings.name = color_space - except (TypeError, ValueError): - continue - if settings.name == color_space: - return color_space - return settings.name - - -def is_animated_image(image): - return getattr(image, "source", None) in {"MOVIE", "SEQUENCE"} - - -def video_total_frames(image): - """Return the total animated frame count, or 1 for a still image.""" - source = getattr(image, "source", None) - if source == "MOVIE": - return int(getattr(image, "frame_duration", 1) or 1) - if source == "SEQUENCE": - return len(collect_sequence_frames(image)) or 1 - return 1 - - -def image_user_settings(source_object): - """Return the Empty's image user settings as a four-item tuple.""" - image_user = getattr(source_object, "image_user", None) - if image_user is None: - return 1, 0, 0, False - return ( - int(getattr(image_user, "frame_start", 1) or 1), - int(getattr(image_user, "frame_offset", 0) or 0), - int(getattr(image_user, "frame_duration", 0) or 0), - bool(getattr(image_user, "use_cyclic", False)), - ) - - -def extract_frame_range(source_object): - """Return the clamped (start_frame, frame_count) range extracted from the source.""" - total_frames = video_total_frames(getattr(source_object, "data", None)) - if total_frames <= 1: - return 1, total_frames - _frame_start, frame_offset, image_frames, _use_cyclic = ( - image_user_settings(source_object) - ) - frame_count = ( - min(image_frames, total_frames) - if image_frames >= 1 - else total_frames - ) - start_frame = min(1 + frame_offset, total_frames - frame_count + 1) - start_frame = max(start_frame, 1) - return start_frame, frame_count - - -def default_extract_frames(source_object): - """Return the Empty's Image Frames value, or the full frame count when unset.""" - return extract_frame_range(source_object)[1] - - -def image_path(image): - filepath_from_user = getattr(image, "filepath_from_user", None) - filepath = ( - filepath_from_user() - if callable(filepath_from_user) - else getattr(image, "filepath", "") - ) - if not filepath: - return Path() - return Path(bpy.path.abspath(filepath)).expanduser() - - -def image_is_packed(image): - if getattr(image, "packed_file", None) is not None: - return True - return bool(getattr(image, "packed_files", ())) - - -def remove_temporary_input(path): - """Remove a temporary input file or frame directory.""" - path = Path(path) - try: - if path.is_dir(): - shutil.rmtree(path, ignore_errors=True) - else: - path.unlink(missing_ok=True) - except OSError: - pass - - -def image_base_name(image): - name = image.name - lower_name = name.lower() - for suffix in IMAGE_NAME_SUFFIXES + VIDEO_NAME_SUFFIXES: - suffix_index = lower_name.rfind(suffix) - if suffix_index <= 0: - continue - trailing = name[suffix_index + len(suffix) :] - if trailing and not ( - trailing.startswith(".") and trailing[1:].isdigit() - ): - continue - return f"{name[:suffix_index]}{trailing}" - return name - - -def depth_data_name(source_object): - return f"{image_base_name(source_object.data)}.depth" - - -def normal_data_name(source_object): - return f"{image_base_name(source_object.data)}.normal" - - -def mask_data_name(source_object): - return f"{image_base_name(source_object.data)}.mask" - - -def copy_depth_image(source_object): - """Return a renamed copy of the source image for use as the depth map.""" - depth_image = source_object.data.copy() - depth_image.name = depth_data_name(source_object) - return depth_image - - -def color_data_name(source_object): - return f"{image_base_name(source_object.data)}.color" - - -def scaled_process_resolution(image, scale_percent): - """Return the scaled long edge clamped and aligned for DA3 inference.""" - width, height = image.size - if width <= 0 or height <= 0: - raise ValueError("The source image has no readable dimensions") - if scale_percent < 1.0 or scale_percent > 100.0: - raise ValueError("Resolution Scale must be between 1% and 100%") - process_res = max(width, height) * scale_percent / 100.0 - process_res = max(DA3_MIN_PROCESS_RES, process_res) - process_res = min(DA3_MAX_PROCESS_RES, process_res) - return int(process_res / DA3_PATCH_SIZE + 0.5) * DA3_PATCH_SIZE - - -def process_resolution_scale_property(): - return bpy.props.FloatProperty( - name="Resolution Scale", - description=( - "Scale the source image long edge for DA3 processing; the result " - "is limited to 252-1008 px and aligned to a multiple of 14" - ), - default=100.0, - min=1.0, - max=100.0, - step=100, - precision=0, - subtype="PERCENTAGE", - ) - - -def draw_process_resolution(layout, owner, image): - process_res = scaled_process_resolution( - image, - owner.process_res_factor, - ) - layout.prop( - owner, - "process_res_factor", - text=f"Resolution Scale ({process_res} px)", - ) - - -def save_packed_image(image, scene): - output_path = Path(tempfile.mkdtemp(prefix="snapform-input-")) / "input.png" - image_settings = scene.render.image_settings - original_file_format = image_settings.file_format - try: - image_settings.file_format = "PNG" - image.save_render(str(output_path), scene=scene) - except (OSError, RuntimeError) as error: - output_path.unlink(missing_ok=True) - raise RuntimeError("Unable to export the packed image") from error - finally: - image_settings.file_format = original_file_format - if not output_path.is_file(): - raise RuntimeError("Blender did not export the packed image") - return output_path - - -def collect_sequence_frames(image): - """Return the numbered frame files of a SEQUENCE image, or [] when absent.""" - source_path = image_path(image) - if not is_animated_image(image) or not source_path.is_file(): - return [] - match = re.match( - r"^(?P.*?)(?P\d+)(?P\.[^.]+)$", - source_path.name, - ) - if match is None: - return [] - prefix = match.group("prefix") - suffix = match.group("suffix").lower() - frames = [] - for path in source_path.parent.iterdir(): - frame_match = re.match( - r"^(?P.*?)(?P\d+)(?P\.[^.]+)$", - path.name, - ) - if ( - frame_match is not None - and frame_match.group("prefix") == prefix - and frame_match.group("suffix").lower() == suffix - ): - frames.append((int(frame_match.group("digits")), path)) - return [path for _digits, path in sorted(frames)] - - -def prepare_image_input(image, scene): - source_path = image_path(image) - if image_is_packed(image): - return save_packed_image(image, scene), True - if getattr(image, "source", None) == "MOVIE": - if not source_path.is_file(): - raise RuntimeError( - "The image empty must reference an existing video file" - ) - return source_path, False - sequence_frames = collect_sequence_frames(image) - if sequence_frames: - temporary_directory = Path(tempfile.mkdtemp(prefix="snapform-input-")) - try: - for index, frame_path in enumerate(sequence_frames): - suffix = frame_path.suffix.lower() or ".jpg" - shutil.copy2( - frame_path, - temporary_directory / f"frame.{index + 1:04d}{suffix}", - ) - except Exception: - shutil.rmtree(temporary_directory, ignore_errors=True) - raise - return temporary_directory, True - if not source_path.is_file(): - raise RuntimeError( - "The image empty must reference an existing image file" - ) - return source_path, False - - -def configure_image_user( - image_user, - frame_count=None, - frame_start=1, - frame_offset=0, - use_cyclic=None, -): - """Map the image to scene frames and enable auto refresh.""" - if image_user is None: - return - if hasattr(image_user, "frame_start"): - image_user.frame_start = frame_start - if hasattr(image_user, "frame_offset"): - image_user.frame_offset = frame_offset - if hasattr(image_user, "use_auto_refresh"): - image_user.use_auto_refresh = True - if use_cyclic is not None and hasattr(image_user, "use_cyclic"): - image_user.use_cyclic = use_cyclic - if frame_count is not None and hasattr(image_user, "frame_duration"): - image_user.frame_duration = frame_count - - -def pack_sequence_if_possible(image, expected_frames): - """Pack a frame sequence when Blender supports it; otherwise keep external files.""" - if expected_frames <= 1: - return False - try: - image.pack() - except (OSError, RuntimeError): - return False - if len(getattr(image, "packed_files", ())) >= expected_frames: - return True - try: - image.unpack(method="REMOVE") - except (OSError, RuntimeError): - pass - return False - - -def sorted_frame_paths(directory, prefix=None): - """Return image files in a directory sorted by their trailing frame number.""" - def frame_key(path): - match = re.search(r"(\d+)(?=\.[^.]+$)", path.name) - return (int(match.group(1)) if match is not None else 0, path.name) - - return sorted( - ( - path - for path in Path(directory).iterdir() - if path.is_file() - and path.suffix.lower() in IMAGE_NAME_SUFFIXES - and (prefix is None or path.stem.startswith(prefix)) - ), - key=frame_key, - ) - - -def force_sequence_source(image): - """Ensure a loaded image is a SEQUENCE datablock for its numbered frames.""" - if getattr(image, "source", None) == "SEQUENCE": - return image - try: - image.source = "SEQUENCE" - except (OSError, RuntimeError, TypeError) as error: - raise RuntimeError( - "Blender cannot load the frame sequence" - ) from error - if getattr(image, "source", None) != "SEQUENCE": - raise RuntimeError("Blender cannot load the frame sequence") - return image - - -def create_sequence_image(name, frame_directory, frame_prefix="frame."): - """Load a numbered frame directory as a SEQUENCE image datablock.""" - frame_paths = sorted_frame_paths(frame_directory, frame_prefix) - if not frame_paths: - raise RuntimeError(f"No image frames found in: {frame_directory}") - try: - image = bpy.data.images.load( - str(frame_paths[0]), - check_existing=False, - ) - force_sequence_source(image) - except (OSError, RuntimeError) as error: - if "image" in locals(): - bpy.data.images.remove(image, do_unlink=True) - raise RuntimeError("Unable to load the frame sequence") from error - image.name = name - pack_sequence_if_possible(image, len(frame_paths)) - return image - - -def create_depth_sequence_image(name, depth_frames, scene, frame_directory): - """Write one EXR per depth frame and return a SEQUENCE image datablock.""" - import numpy as np - - frame_directory = Path(frame_directory) - frame_directory.mkdir(parents=True, exist_ok=True) - height, width = depth_frames[0].shape - image = bpy.data.images.new( - name, - width=width, - height=height, - alpha=False, - float_buffer=True, - ) - image.colorspace_settings.name = "Non-Color" - image_settings = scene.render.image_settings - original_file_format = image_settings.file_format - frame_paths = [] - try: - image_settings.file_format = "OPEN_EXR" - for frame_index, depth_values in enumerate(depth_frames): - # NPZ arrays use a top-left origin, while Blender image pixels start below. - depth_pixels = np.flipud(depth_values) - rgba = np.empty((height, width, 4), dtype=np.float32) - rgba[..., 0] = depth_pixels - rgba[..., 1] = depth_pixels - rgba[..., 2] = depth_pixels - rgba[..., 3] = 1.0 - image.pixels.foreach_set(rgba.ravel()) - frame_path = ( - frame_directory / f"depth.{frame_index + 1:04d}.exr" - ) - image.save_render(str(frame_path), scene=scene) - if not frame_path.is_file(): - raise RuntimeError(f"Blender did not write: {frame_path}") - frame_paths.append(frame_path) - finally: - image_settings.file_format = original_file_format - bpy.data.images.remove(image, do_unlink=True) - if not frame_paths: - raise RuntimeError("Unable to write the depth frame sequence") - try: - sequence_image = bpy.data.images.load( - str(frame_paths[0]), - check_existing=False, - ) - force_sequence_source(sequence_image) - except (OSError, RuntimeError) as error: - if "sequence_image" in locals(): - bpy.data.images.remove(sequence_image, do_unlink=True) - raise RuntimeError( - "Unable to load the depth frame sequence" - ) from error - sequence_image.name = name - sequence_image.colorspace_settings.name = "Non-Color" - pack_sequence_if_possible(sequence_image, len(frame_paths)) - return sequence_image - - -def create_color_image( - source_object, - output_directory, - assign_to_source=True, -): - output_directory = Path(output_directory) - color_paths = sorted_frame_paths(output_directory, "color.") - if not color_paths: - raise RuntimeError( - f"Color result has no images: {output_directory}" - ) - source_image = source_object.data - animated = is_animated_image(source_image) - try: - result_image = bpy.data.images.load( - str(color_paths[0]), - check_existing=False, - ) - result_image.alpha_mode = "CHANNEL_PACKED" - if animated: - force_sequence_source(result_image) - except (OSError, RuntimeError) as error: - if "result_image" in locals(): - bpy.data.images.remove(result_image, do_unlink=True) - raise RuntimeError("Unable to load the color result") from error - - original_fake_user = source_image.use_fake_user - try: - result_image.name = color_data_name(source_object) - if animated: - pack_sequence_if_possible(result_image, len(color_paths)) - else: - result_image.pack() - if assign_to_source: - image_user = getattr(source_object, "image_user", None) - ( - original_frame_start, - _original_frame_offset, - _original_frame_duration, - original_use_cyclic, - ) = image_user_settings(source_object) - source_image.use_fake_user = True - source_object.data = result_image - if animated: - configure_image_user( - image_user, - len(color_paths), - frame_start=original_frame_start, - use_cyclic=original_use_cyclic, - ) - except (OSError, RuntimeError) as error: - source_image.use_fake_user = original_fake_user - bpy.data.images.remove(result_image, do_unlink=True) - raise RuntimeError("Unable to prepare the color result") from error - return result_image - - -def plane_bounds(source_object): - image_width, image_height = source_object.data.size - if image_width <= 0 or image_height <= 0: - raise ValueError("The source image has no readable dimensions") - - aspect = image_width / image_height - display_size = float(source_object.empty_display_size) - if aspect >= 1.0: - width = display_size - height = display_size / aspect - else: - width = display_size * aspect - height = display_size - - offset = source_object.empty_image_offset - x_min = float(offset[0]) * width - y_min = float(offset[1]) * height - return ( - x_min, - x_min + width, - y_min, - y_min + height, - ) - - -def plane_fit_transform(source_object, camera_bounds): - target_bounds = plane_bounds(source_object) - camera_width = camera_bounds[1] - camera_bounds[0] - camera_height = camera_bounds[3] - camera_bounds[2] - if camera_width <= 0.0 or camera_height <= 0.0: - raise ValueError("The camera plane has invalid dimensions") - - target_width = target_bounds[1] - target_bounds[0] - target_height = target_bounds[3] - target_bounds[2] - scale_x = target_width / camera_width - scale_y = target_height / camera_height - scale_z = (scale_x * scale_y) ** 0.5 - camera_center_x = (camera_bounds[0] + camera_bounds[1]) * 0.5 - camera_center_y = (camera_bounds[2] + camera_bounds[3]) * 0.5 - target_center_x = (target_bounds[0] + target_bounds[1]) * 0.5 - target_center_y = (target_bounds[2] + target_bounds[3]) * 0.5 - return ( - scale_x, - scale_y, - scale_z, - target_center_x - camera_center_x * scale_x, - target_center_y - camera_center_y * scale_y, - ) - - -def apply_plane_transform(plane_object, source_object, camera_bounds): - from mathutils import Matrix - - scale_x, scale_y, scale_z, offset_x, offset_y = plane_fit_transform( - source_object, - camera_bounds, - ) - fit_matrix = Matrix( - ( - (scale_x, 0.0, 0.0, offset_x), - (0.0, scale_y, 0.0, offset_y), - (0.0, 0.0, 1.0, 0.0), - (0.0, 0.0, 0.0, 1.0), - ) - ) - plane_object.matrix_world = source_object.matrix_world @ fit_matrix - return abs(float(plane_object.scale.z)) * abs(scale_z) - - -def apply_object_scale_to_mesh(plane_object): - """Bake object scale into the plane while keeping location and rotation.""" - from mathutils import Matrix - - scale = plane_object.scale.copy() - plane_object.data.transform( - Matrix.Diagonal((scale.x, scale.y, scale.z, 1.0)) - ) - plane_object.data.update() - plane_object.scale = (1.0, 1.0, 1.0) - - -def scale_depth_values(depth, depth_scale): - return prepare_depth_values(depth) * depth_scale - - -def extract_depth_arrays(npz_path): - try: - import numpy as np - except ImportError as error: - raise RuntimeError("Blender Python does not provide NumPy") from error - - try: - with np.load(npz_path, allow_pickle=False) as archive: - depth_key = next( - ( - key - for key in DEPTH_KEY_CANDIDATES - if key in archive.files - ), - None, - ) - if depth_key is None: - keys = ", ".join(archive.files) or "none" - raise ValueError( - f"Mini NPZ has no depth array; available keys: {keys}" - ) - depth = np.asarray(archive[depth_key], dtype=np.float32) - except OSError as error: - raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error - - if depth.ndim >= 3 and depth.shape[-1] == 1: - depth = depth[..., 0] - if depth.ndim == 2: - depth = depth[None, ...] - if depth.ndim != 3: - raise ValueError( - f"Unsupported depth array shape: {tuple(depth.shape)}" - ) - return [depth[index] for index in range(depth.shape[0])] - - -def extract_depth_array(npz_path): - return extract_depth_arrays(npz_path)[0] - - -def extract_intrinsics_array(npz_path): - try: - import numpy as np - except ImportError as error: - raise RuntimeError("Blender Python does not provide NumPy") from error - - try: - with np.load(npz_path, allow_pickle=False) as archive: - intrinsics_key = next( - ( - key - for key in INTRINSICS_KEY_CANDIDATES - if key in archive.files - ), - None, - ) - if intrinsics_key is None: - keys = ", ".join(archive.files) or "none" - raise ValueError( - "Mini NPZ has no camera intrinsics; " - f"available keys: {keys}" - ) - intrinsics = np.asarray( - archive[intrinsics_key], - dtype=np.float32, - ) - except OSError as error: - raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error - - if intrinsics.ndim == 3: - intrinsics = intrinsics[0] - if intrinsics.shape != (3, 3): - raise ValueError( - f"Unsupported intrinsics shape: {tuple(intrinsics.shape)}" - ) - return intrinsics - - -def prepare_depth_values(depth): - import numpy as np - - valid = np.isfinite(depth) - if not valid.any(): - raise ValueError("Depth array contains no finite values") - values = depth.astype(np.float32, copy=True) - values[~valid] = 0.0 - return values - - -def depth_displacement_values(depth): - """Convert raw DA3 camera distance to positive outward displacement.""" - import numpy as np - - return 1.0 - np.asarray(depth, dtype=np.float32) - - -def camera_plane_bounds( - source_object, - depth, - intrinsics, - displacement_strength, - reference_depth=None, -): - import numpy as np - - scale = abs(float(displacement_strength)) - if scale <= 1e-8: - raise ValueError("Displacement Strength must not be zero") - - if reference_depth is None: - valid = np.isfinite(depth) & (depth > 0.0) - if not valid.any(): - raise ValueError("Depth array contains no positive finite values") - reference_depth = float(np.median(depth[valid])) - reference_distance = reference_depth * scale - - height, width = depth.shape - fx = float(intrinsics[0, 0]) - fy = float(intrinsics[1, 1]) - cx = float(intrinsics[0, 2]) - cy = float(intrinsics[1, 2]) - if fx <= 0.0 or fy <= 0.0: - raise ValueError("Camera intrinsics contain an invalid focal length") - - source_width, source_height = source_object.data.size - if source_width <= 0 or source_height <= 0: - raise ValueError("The source image has no readable dimensions") - source_aspect = source_width / source_height - - # Predicted fx/fy may differ even for square pixels. Preserve their projected - # area while enforcing the source image aspect ratio to avoid image stretch. - raw_width = width * reference_distance / fx - raw_height = height * reference_distance / fy - plane_area = raw_width * raw_height - plane_width = float(np.sqrt(plane_area * source_aspect)) - plane_height = float(np.sqrt(plane_area / source_aspect)) - - x_min = -(cx / width) * plane_width - x_max = (1.0 - cx / width) * plane_width - y_min = -(1.0 - cy / height) * plane_height - y_max = (cy / height) * plane_height - - # Preserve an intentional Image Empty origin offset around the camera frame. - offset = source_object.empty_image_offset - x_offset = (float(offset[0]) + 0.5) * plane_width - y_offset = (float(offset[1]) + 0.5) * plane_height - return ( - ( - x_min + x_offset, - x_max + x_offset, - y_min + y_offset, - y_max + y_offset, - ), - reference_depth, - ) - - -def create_depth_image(name, depth_values): - import numpy as np - - height, width = depth_values.shape - image = bpy.data.images.new( - name, - width=width, - height=height, - alpha=False, - float_buffer=True, - ) - image.colorspace_settings.name = "Non-Color" - image.file_format = "OPEN_EXR" - - # NPZ arrays use a top-left origin, while Blender image pixels start below. - depth_pixels = np.flipud(depth_values) - rgba = np.empty((height, width, 4), dtype=np.float32) - rgba[..., 0] = depth_pixels - rgba[..., 1] = depth_pixels - rgba[..., 2] = depth_pixels - rgba[..., 3] = 1.0 - image.pixels.foreach_set(rgba.ravel()) - # Pack as EXR so values outside 0..1 survive saving and reopening the blend. - image.pack() - return image - - -def create_moge_mask_image(name, mask_values): - import numpy as np - - height, width = mask_values.shape - image = bpy.data.images.new( - name, - width=width, - height=height, - alpha=True, - float_buffer=True, - ) - image.colorspace_settings.name = "Non-Color" - values = np.flipud(mask_values.astype(np.float32)) - rgba = np.ones((height, width, 4), dtype=np.float32) - rgba[..., :3] = values[..., None] - rgba[..., 3] = values - image.pixels.foreach_set(rgba.ravel()) - image.pack() - return image - - -def create_moge_normal_image(name, normal_values, convert_coordinates=True): - import numpy as np - - converted = normal_values.astype(np.float32, copy=True) - if convert_coordinates: - converted *= np.array([1.0, -1.0, -1.0], dtype=np.float32) - converted = np.clip(converted * 0.5 + 0.5, 0.0, 1.0) - height, width = converted.shape[:2] - image = bpy.data.images.new( - name, - width=width, - height=height, - alpha=False, - float_buffer=True, - ) - image.colorspace_settings.name = "Non-Color" - rgba = np.ones((height, width, 4), dtype=np.float32) - rgba[..., :3] = np.flipud(converted) - image.pixels.foreach_set(rgba.ravel()) - image.pack() - return image - - -def create_moge_sequence_image(name, frames, scene, directory, prefix, kind): - import numpy as np - - directory = Path(directory) - directory.mkdir(parents=True, exist_ok=True) - height, width = frames[0].shape[:2] - image = bpy.data.images.new( - name, - width=width, - height=height, - alpha=True, - float_buffer=True, - ) - image.colorspace_settings.name = "Non-Color" - image_settings = scene.render.image_settings - original_format = image_settings.file_format - paths = [] - try: - image_settings.file_format = "OPEN_EXR" - for index, values in enumerate(frames): - rgba = np.ones((height, width, 4), dtype=np.float32) - if kind == "depth": - flipped = np.flipud(values.astype(np.float32)) - rgba[..., :3] = flipped[..., None] - elif kind == "mask": - flipped = np.flipud(values.astype(np.float32)) - rgba[..., :3] = flipped[..., None] - rgba[..., 3] = flipped - else: - converted = values.astype(np.float32, copy=True) - converted *= np.array([1.0, -1.0, -1.0], dtype=np.float32) - rgba[..., :3] = np.flipud( - np.clip(converted * 0.5 + 0.5, 0.0, 1.0) - ) - image.pixels.foreach_set(rgba.ravel()) - path = directory / f"{prefix}.{index + 1:04d}.exr" - image.save_render(str(path), scene=scene) - paths.append(path) - finally: - image_settings.file_format = original_format - bpy.data.images.remove(image, do_unlink=True) - sequence = bpy.data.images.load(str(paths[0]), check_existing=False) - force_sequence_source(sequence) - sequence.name = name - sequence.colorspace_settings.name = "Non-Color" - pack_sequence_if_possible(sequence, len(paths)) - return sequence - - -def moge_material_node_group(): - existing = bpy.data.node_groups.get(MOGE_MATERIAL_NODE_GROUP_NAME) - if existing is not None: - return existing - group = bpy.data.node_groups.new( - MOGE_MATERIAL_NODE_GROUP_NAME, - "ShaderNodeTree", - ) - for name, socket_type in ( - ("Color", "NodeSocketColor"), - ("Alpha", "NodeSocketFloat"), - ("Depth", "NodeSocketFloat"), - ("Bump Scale", "NodeSocketFloat"), - ): - socket = group.interface.new_socket( - name=name, - in_out="INPUT", - socket_type=socket_type, - ) - if name == "Alpha": - socket.default_value = 1.0 - if name == "Bump Scale": - socket.default_value = 1.0 - for name, socket_type in ( - ("Color", "NodeSocketColor"), - ("Roughness", "NodeSocketFloat"), - ("Alpha", "NodeSocketFloat"), - ("Normal", "NodeSocketVector"), - ("Displacement", "NodeSocketVector"), - ): - group.interface.new_socket( - name=name, - in_out="OUTPUT", - socket_type=socket_type, - ) - nodes = group.nodes - links = group.links - group_input = nodes.new("NodeGroupInput") - group_input.location = (-900.0, 100.0) - bump = nodes.new("ShaderNodeBump") - bump.location = (250.0, -100.0) - bump_distance = nodes.new("ShaderNodeMath") - bump_distance.operation = "MULTIPLY" - bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE - bump_distance.location = (0.0, -100.0) - background = nodes.new("ShaderNodeAttribute") - background.attribute_type = "GEOMETRY" - background.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE - background.location = (-650.0, -400.0) - material_mode = nodes.new("ShaderNodeAttribute") - material_mode.attribute_type = "GEOMETRY" - material_mode.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE - material_mode.location = (-650.0, -600.0) - strength = nodes.new("ShaderNodeAttribute") - strength.attribute_type = "GEOMETRY" - strength.attribute_name = DEPTH_SCALE_ATTRIBUTE - strength.location = (-400.0, -600.0) - clamp_depth = nodes.new("ShaderNodeMath") - clamp_depth.operation = "MAXIMUM" - clamp_depth.location = (-350.0, -350.0) - enable_depth = nodes.new("ShaderNodeMath") - enable_depth.operation = "MULTIPLY" - enable_depth.location = (-100.0, -350.0) - scale_depth = nodes.new("ShaderNodeMath") - scale_depth.operation = "MULTIPLY" - scale_depth.location = (150.0, -350.0) - displacement = nodes.new("ShaderNodeDisplacement") - displacement.inputs["Midlevel"].default_value = 0.0 - displacement.inputs["Scale"].default_value = 1.0 - displacement.location = (400.0, -350.0) - group_output = nodes.new("NodeGroupOutput") - group_output.location = (300.0, 100.0) - links.new(group_input.outputs["Color"], group_output.inputs["Color"]) - links.new(group_input.outputs["Alpha"], group_output.inputs["Alpha"]) - links.new(group_input.outputs["Color"], bump.inputs["Height"]) - links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0]) - links.new(bump_distance.outputs[0], bump.inputs["Distance"]) - links.new(bump.outputs["Normal"], group_output.inputs["Normal"]) - links.new(group_input.outputs["Depth"], clamp_depth.inputs[0]) - links.new(background.outputs["Fac"], clamp_depth.inputs[1]) - links.new(clamp_depth.outputs[0], enable_depth.inputs[0]) - links.new(material_mode.outputs["Fac"], enable_depth.inputs[1]) - links.new(enable_depth.outputs[0], scale_depth.inputs[0]) - links.new(strength.outputs["Fac"], scale_depth.inputs[1]) - links.new(scale_depth.outputs[0], displacement.inputs["Height"]) - links.new(displacement.outputs["Displacement"], group_output.inputs["Displacement"]) - group_output.inputs["Roughness"].default_value = 0.5 - group.color_tag = "SHADER" - group.description = "Apply MoGe-2 color and derive bump detail from it." - group.use_fake_user = True - return group - - -def create_moge_material( - name, - color_image, - depth_image, - depth_scale, - shadeless=False, - frame_count=1, - frame_start=1, - use_cyclic=False, - scene=None, -): - material = bpy.data.materials.new(name) - material.use_nodes = True - if hasattr(material, "surface_render_method"): - material.surface_render_method = "DITHERED" - nodes = material.node_tree.nodes - nodes.clear() - links = material.node_tree.links - uv = nodes.new("ShaderNodeUVMap") - uv.uv_map = "UVMap" - uv.location = (-900.0, 200.0) - color = nodes.new("ShaderNodeTexImage") - set_material_color_space(color_image, scene) - color.image = color_image - color.interpolation = "Linear" - color.extension = "EXTEND" - color.location = (-650.0, 300.0) - depth = nodes.new("ShaderNodeTexImage") - depth.image = depth_image - depth.interpolation = "Linear" - depth.extension = "EXTEND" - depth.location = (-650.0, -100.0) - if frame_count > 1: - configure_image_user( - color.image_user, - frame_count, - frame_start=frame_start, - use_cyclic=use_cyclic, - ) - configure_image_user( - depth.image_user, - frame_count, - frame_start=frame_start, - use_cyclic=use_cyclic, - ) - layer = nodes.new("ShaderNodeGroup") - layer.node_tree = moge_material_node_group() - layer.name = MOGE_MATERIAL_NODE_GROUP_NAME - layer.location = (-300.0, 200.0) - output = nodes.new("ShaderNodeOutputMaterial") - output.location = (500.0, 200.0) - links.new(uv.outputs["UV"], color.inputs["Vector"]) - links.new(uv.outputs["UV"], depth.inputs["Vector"]) - links.new(color.outputs["Color"], layer.inputs["Color"]) - links.new(color.outputs["Alpha"], layer.inputs["Alpha"]) - links.new(depth.outputs["Color"], layer.inputs["Depth"]) - layer.inputs["Bump Scale"].default_value = depth_scale - if shadeless: - shader = nodes.new("ShaderNodeGroup") - shader.node_tree = load_shadeless_node_group() - shader.name = SHADELESS_NODE_GROUP_NAME - transparent = nodes.new("ShaderNodeBsdfTransparent") - mix = nodes.new("ShaderNodeMixShader") - links.new(layer.outputs["Color"], shader.inputs["Color"]) - links.new(layer.outputs["Alpha"], mix.inputs["Fac"]) - links.new(transparent.outputs["BSDF"], mix.inputs[1]) - links.new(shader.outputs["Shader"], mix.inputs[2]) - links.new(mix.outputs["Shader"], output.inputs["Surface"]) - 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) - return material - - -def grid_segments(bounds, subdivision_level): - if subdivision_level < 0: - raise ValueError("Subdivision must not be negative") - - width = bounds[1] - bounds[0] - height = bounds[3] - bounds[2] - if width <= 0.0 or height <= 0.0: - raise ValueError("The plane has invalid dimensions") - - short_segments = 2**subdivision_level - if width >= height: - return ( - max(short_segments, int(short_segments * width / height + 0.5)), - short_segments, - ) - return ( - short_segments, - max(short_segments, int(short_segments * height / width + 0.5)), - ) - - -def create_plane_mesh(source_object, bounds=None, subdivision_level=5): - x_min, x_max, y_min, y_max = ( - bounds if bounds is not None else plane_bounds(source_object) - ) - if subdivision_level is None: - x_segments, y_segments = 1, 1 - else: - x_segments, y_segments = grid_segments( - (x_min, x_max, y_min, y_max), - subdivision_level, - ) - vertices = [] - uv_coordinates = [] - for y_index in range(y_segments + 1): - v = y_index / y_segments - y = y_min + (y_max - y_min) * v - for x_index in range(x_segments + 1): - u = x_index / x_segments - x = x_min + (x_max - x_min) * u - vertices.append((x, y, 0.0)) - uv_coordinates.append((u, v)) - - faces = [] - row_width = x_segments + 1 - for y_index in range(y_segments): - for x_index in range(x_segments): - bottom_left = y_index * row_width + x_index - faces.append( - ( - bottom_left, - bottom_left + 1, - bottom_left + row_width + 1, - bottom_left + row_width, - ) - ) - - mesh = bpy.data.meshes.new(source_object.name) - mesh.from_pydata(vertices, (), faces) - mesh.update() - for polygon in mesh.polygons: - polygon.use_smooth = True - uv_layer = mesh.uv_layers.new(name="UVMap") - for loop in mesh.loops: - uv_layer.data[loop.index].uv = uv_coordinates[ - loop.vertex_index - ] - return mesh - - -def material_node_group(depth_plane=False, surface_group_name=None): - group_name = ( - DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME - if depth_plane - else surface_group_name or SURFACE_MATERIAL_NODE_GROUP_NAME - ) - existing = bpy.data.node_groups.get(group_name) - if existing is not None: - return existing - group = bpy.data.node_groups.new(group_name, "ShaderNodeTree") - interface = group.interface - input_sockets = [ - ("Color", "NodeSocketColor"), - ("Alpha", "NodeSocketFloat"), - ("Bump Scale", "NodeSocketFloat"), - ] - if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME: - input_sockets.insert(2, ("Normal", "NodeSocketColor")) - input_sockets.insert(3, ("Normal Strength", "NodeSocketFloat")) - if depth_plane: - input_sockets.insert(2, ("Depth", "NodeSocketFloat")) - for name, socket_type in input_sockets: - socket = interface.new_socket( - name=name, - in_out="INPUT", - socket_type=socket_type, - ) - if name == "Alpha": - socket.default_value = 1.0 - socket.min_value = 0.0 - socket.max_value = 1.0 - socket.subtype = "FACTOR" - elif name == "Bump Scale": - socket.default_value = 1.0 - socket.min_value = -10000.0 - socket.max_value = 10000.0 - elif name == "Normal Strength": - socket.default_value = 1.0 - socket.min_value = 0.0 - socket.max_value = 10.0 - output_sockets = [ - ("Color", "NodeSocketColor"), - ("Roughness", "NodeSocketFloat"), - ("Alpha", "NodeSocketFloat"), - ("Normal", "NodeSocketVector"), - ] - if depth_plane: - output_sockets.append(("Displacement", "NodeSocketVector")) - for name, socket_type in output_sockets: - interface.new_socket( - name=name, - in_out="OUTPUT", - socket_type=socket_type, - ) - - nodes = group.nodes - links = group.links - group_input = nodes.new("NodeGroupInput") - group_input.location = (-500.0, 100.0) - group_output = nodes.new("NodeGroupOutput") - group_output.location = (300.0, 100.0) - roughness_curve = nodes.new("ShaderNodeFloatCurve") - roughness_curve.mapping.initialize() - curve_points = roughness_curve.mapping.curves[0].points - curve_points[0].location = (0.0, 0.5) - curve_points[-1].location = (1.0, 0.1) - roughness_curve.mapping.update() - roughness_curve.location = (-250.0, 350.0) - bump = nodes.new("ShaderNodeBump") - bump.location = (-250.0, 0.0) - if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME: - normal_map = nodes.new("ShaderNodeNormalMap") - normal_map.space = "OBJECT" - normal_map.location = (-500.0, -150.0) - links.new(group_input.outputs["Normal"], normal_map.inputs["Color"]) - links.new( - group_input.outputs["Normal Strength"], - normal_map.inputs["Strength"], - ) - links.new(normal_map.outputs["Normal"], bump.inputs["Normal"]) - displacement = None - - processed_alpha = group_input.outputs["Alpha"] - bump_depth = group_input.outputs["Color"] - if depth_plane: - displacement = nodes.new("ShaderNodeDisplacement") - displacement.inputs["Midlevel"].default_value = 0.0 - displacement.inputs["Scale"].default_value = 1.0 - group_input.location = (-1200.0, 100.0) - group_output.location = (1050.0, 100.0) - roughness_curve.location = (-300.0, 450.0) - bump.location = (750.0, 50.0) - displacement.location = (750.0, -300.0) - background_attribute = nodes.new("ShaderNodeAttribute") - background_attribute.attribute_type = "GEOMETRY" - background_attribute.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE - background_attribute.location = (-950.0, -500.0) - material_attribute = nodes.new("ShaderNodeAttribute") - material_attribute.attribute_type = "GEOMETRY" - material_attribute.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE - material_attribute.location = (-50.0, -650.0) - strength_attribute = nodes.new("ShaderNodeAttribute") - strength_attribute.attribute_type = "GEOMETRY" - strength_attribute.attribute_name = DEPTH_SCALE_ATTRIBUTE - strength_attribute.location = (200.0, -650.0) - foreground_mask = nodes.new("ShaderNodeMath") - foreground_mask.operation = "GREATER_THAN" - foreground_mask.inputs[1].default_value = 0.5 - foreground_mask.location = (-950.0, -250.0) - clamp_foreground = nodes.new("ShaderNodeMath") - clamp_foreground.operation = "MAXIMUM" - clamp_foreground.location = (-650.0, -400.0) - masked_depth = nodes.new("ShaderNodeMix") - masked_depth.data_type = "FLOAT" - masked_depth.location = (-350.0, -350.0) - enabled_depth = nodes.new("ShaderNodeMath") - enabled_depth.operation = "MULTIPLY" - enabled_depth.location = (250.0, -350.0) - scaled_depth = nodes.new("ShaderNodeMath") - scaled_depth.operation = "MULTIPLY" - scaled_depth.location = (450.0, -350.0) - links.new(group_input.outputs["Alpha"], foreground_mask.inputs[0]) - links.new(group_input.outputs["Depth"], clamp_foreground.inputs[0]) - links.new(background_attribute.outputs["Fac"], clamp_foreground.inputs[1]) - links.new(foreground_mask.outputs[0], masked_depth.inputs[0]) - links.new(background_attribute.outputs["Fac"], masked_depth.inputs[2]) - links.new(clamp_foreground.outputs[0], masked_depth.inputs[3]) - links.new(masked_depth.outputs["Result"], enabled_depth.inputs[0]) - links.new(material_attribute.outputs["Fac"], enabled_depth.inputs[1]) - links.new(enabled_depth.outputs[0], scaled_depth.inputs[0]) - links.new(strength_attribute.outputs["Fac"], scaled_depth.inputs[1]) - processed_alpha = foreground_mask.outputs[0] - processed_depth = scaled_depth.outputs[0] - - links.new(group_input.outputs["Color"], group_output.inputs["Color"]) - links.new(group_input.outputs["Color"], roughness_curve.inputs["Value"]) - links.new( - roughness_curve.outputs["Value"], - group_output.inputs["Roughness"], - ) - links.new(processed_alpha, group_output.inputs["Alpha"]) - links.new(bump_depth, bump.inputs["Height"]) - bump_distance = nodes.new("ShaderNodeMath") - bump_distance.operation = "MULTIPLY" - bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE - bump_distance.location = (500.0, 100.0) if depth_plane else (-250.0, 150.0) - links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0]) - links.new(bump_distance.outputs[0], bump.inputs["Distance"]) - links.new(bump.outputs["Normal"], group_output.inputs["Normal"]) - if displacement is not None: - links.new(processed_depth, displacement.inputs["Height"]) - links.new( - displacement.outputs["Displacement"], - group_output.inputs["Displacement"], - ) - group.color_tag = "SHADER" - group.description = ( - "Prepare DA3 depth-plane material outputs." - if depth_plane - else f"Prepare {group_name} material outputs." - ) - group.use_fake_user = True - return group - - -def create_material( - name, - color_image, - depth_scale, - displacement_image=None, - displacement_midlevel=0.0, - displacement_scale=1.0, - shadeless=False, - frame_count=1, - frame_start=1, - color_frame_offset=0, - color_frame_duration=None, - depth_frame_offset=0, - depth_frame_duration=None, - use_cyclic=None, - surface_group_name=None, - normal_image=None, - scene=None, -): - material = bpy.data.materials.new(name) - material.use_nodes = True - if hasattr(material, "surface_render_method"): - material.surface_render_method = "DITHERED" - nodes = material.node_tree.nodes - nodes.clear() - - color_texture = nodes.new("ShaderNodeTexImage") - uv_map = nodes.new("ShaderNodeUVMap") - output = nodes.new("ShaderNodeOutputMaterial") - - uv_map.uv_map = "UVMap" - uv_map.location = (-1250.0, 350.0) - material_image = color_image - set_material_color_space(material_image, scene) - color_texture.image = material_image - color_texture.interpolation = "Linear" - color_texture.extension = "EXTEND" - color_texture.location = (-1000.0, 350.0) - output.location = (850.0, 350.0) - extracted_frames = frame_count if frame_count > 1 else None - color_frames = ( - color_frame_duration - if color_frame_duration is not None - else extracted_frames - ) - depth_frames = ( - depth_frame_duration - if depth_frame_duration is not None - else extracted_frames - ) - material_frames = color_frames - material_frame_offset = color_frame_offset - if frame_count > 1 or is_animated_image(material_image): - configure_image_user( - color_texture.image_user, - material_frames, - frame_start=frame_start, - frame_offset=material_frame_offset, - use_cyclic=use_cyclic, - ) - - links = material.node_tree.links - links.new(uv_map.outputs["UV"], color_texture.inputs["Vector"]) - alpha = color_texture.outputs["Alpha"] - material_inputs = nodes.new("ShaderNodeGroup") - is_depth_plane = displacement_image is not None - material_inputs.node_tree = material_node_group( - is_depth_plane, - surface_group_name=surface_group_name, - ) - material_inputs.name = material_inputs.node_tree.name - material_inputs.label = material_inputs.name - material_inputs.location = (-300.0, 350.0) - material_inputs.inputs["Bump Scale"].default_value = depth_scale - links.new(color_texture.outputs["Color"], material_inputs.inputs["Color"]) - links.new(alpha, material_inputs.inputs["Alpha"]) - if normal_image is not None and "Normal" in material_inputs.inputs: - normal_texture = nodes.new("ShaderNodeTexImage") - normal_texture.image = normal_image - normal_texture.interpolation = "Linear" - normal_texture.extension = "EXTEND" - normal_texture.location = (-1000.0, -50.0) - links.new(uv_map.outputs["UV"], normal_texture.inputs["Vector"]) - links.new(normal_texture.outputs["Color"], material_inputs.inputs["Normal"]) - if shadeless: - shader = nodes.new("ShaderNodeGroup") - shader.node_tree = load_shadeless_node_group() - shader.name = SHADELESS_NODE_GROUP_NAME - transparent = nodes.new("ShaderNodeBsdfTransparent") - mix_shader = nodes.new("ShaderNodeMixShader") - shader.location = (0.0, 550.0) - transparent.location = (0.0, 150.0) - mix_shader.location = (300.0, 350.0) - links.new( - material_inputs.outputs["Color"], - shader.inputs["Color"], - ) - links.new( - material_inputs.outputs["Alpha"], - mix_shader.inputs["Fac"], - ) - links.new( - transparent.outputs["BSDF"], - mix_shader.inputs[1], - ) - links.new( - shader.outputs["Shader"], - mix_shader.inputs[2], - ) - links.new( - mix_shader.outputs["Shader"], - output.inputs["Surface"], - ) - else: - shader = nodes.new("ShaderNodeBsdfPrincipled") - shader.location = (250.0, 400.0) - - links.new( - 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"], - ) - links.new(shader.outputs["BSDF"], output.inputs["Surface"]) - - if displacement_image is not None: - depth_texture = nodes.new("ShaderNodeTexImage") - depth_texture.image = displacement_image - depth_texture.interpolation = "Linear" - depth_texture.extension = "EXTEND" - if frame_count > 1 or is_animated_image(displacement_image): - configure_image_user( - depth_texture.image_user, - depth_frames, - frame_start=frame_start, - frame_offset=depth_frame_offset, - use_cyclic=use_cyclic, - ) - depth_texture.location = (-1000.0, -400.0) - links.new(uv_map.outputs["UV"], depth_texture.inputs["Vector"]) - displacement_height = depth_texture.outputs["Color"] - if displacement_midlevel != 0.0 or displacement_scale != 1.0: - depth_adjust = nodes.new("ShaderNodeMath") - depth_adjust.operation = "SUBTRACT" - depth_adjust.inputs[1].default_value = displacement_midlevel - depth_adjust.location = (500.0, -550.0) - depth_scale_node = nodes.new("ShaderNodeMath") - depth_scale_node.operation = "MULTIPLY" - depth_scale_node.inputs[1].default_value = displacement_scale - depth_scale_node.location = (700.0, -550.0) - links.new(displacement_height, depth_adjust.inputs[0]) - links.new(depth_adjust.outputs[0], depth_scale_node.inputs[0]) - links.new(depth_scale_node.outputs[0], material_inputs.inputs["Depth"]) - else: - links.new(displacement_height, material_inputs.inputs["Depth"]) - links.new( - material_inputs.outputs["Displacement"], - output.inputs["Displacement"], - ) - set_displacement_only(material) - return material - - -def set_displacement_only(material): - for owner, property_name in ( - (material, "surface_displacement_method"), - (material, "displacement_method"), - (getattr(material, "cycles", None), "displacement_method"), - ): - if owner is not None and hasattr(owner, property_name): - setattr(owner, property_name, "DISPLACEMENT") - return - - -def node_asset_path(): - return Path(__file__).parent.parent / "assets" / NODE_ASSET_NAME - - -def load_asset_node_group(group_name): - existing = bpy.data.node_groups.get(group_name) - if existing is not None: - return existing - asset_path = node_asset_path() - if not asset_path.is_file(): - raise RuntimeError(f"Missing node asset: {asset_path}") - use_pack = bpy.app.version >= (5, 0, 0) - if not use_pack: - directory = str(asset_path) + "/NodeTree/" - result = bpy.ops.wm.append( - filepath=directory + group_name, - directory=directory, - filename=group_name, - link=False, - do_reuse_local_id=True, - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to append the node asset") - node_group = bpy.data.node_groups.get(group_name) - if node_group is None: - raise RuntimeError("Appended node asset is unavailable") - node_group.use_fake_user = True - return node_group - with bpy.data.libraries.load(str(asset_path), link=True) as (source, target): - if group_name not in source.node_groups: - raise RuntimeError(f"Node asset is missing {group_name}") - target.node_groups = [group_name] - linked_group = target.node_groups[0] - node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group) - if linked_group != node_group and linked_group.users == 0: - bpy.data.node_groups.remove(linked_group) - node_group.use_fake_user = True - return node_group - - -def load_shadeless_node_group(): - return load_asset_node_group(SHADELESS_NODE_GROUP_NAME) - - -def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME): - return load_asset_node_group(group_name) - - -def modifier_input_identifier(node_group, name): - for item in node_group.interface.items_tree: - if item.item_type != "SOCKET" or item.in_out != "INPUT": - continue - if item.name == name: - return item.identifier - raise RuntimeError(f"Geometry Nodes input is missing: {name}") - - -def modifier_input_slots(modifier): - """Return the runtime input collection, or None for the legacy API.""" - properties = getattr(modifier, "properties", None) - if properties is None: - return None - return getattr(properties, "inputs", None) - - -def set_modifier_input(modifier, identifier, value): - """Assign a value to a Geometry Nodes modifier input socket.""" - inputs = modifier_input_slots(modifier) - if inputs is not None: - # Blender 5.2+ exposes inputs as runtime RNA keyed by socket identifier. - getattr(inputs, identifier).value = value - else: - # Blender 4.x exposes inputs as modifier ID properties. - modifier[identifier] = value - - -def get_modifier_input(modifier, identifier): - """Read a Geometry Nodes modifier input socket value.""" - inputs = modifier_input_slots(modifier) - if inputs is not None: - return getattr(inputs, identifier).value - return modifier[identifier] - - -def modifier_input_driver_path(modifier, identifier): - """Return the driver data path for a modifier input socket.""" - if modifier_input_slots(modifier) is not None: - return ( - f'modifiers["{modifier.name}"]' - f".properties.inputs.{identifier}.value" - ) - return f'modifiers["{modifier.name}"]["{identifier}"]' - - -def drive_modifier_input(owner_object, modifier, identifier, expression="frame"): - """Add a scripted driver to a modifier input so it follows the scene frame.""" - driver = owner_object.driver_add( - modifier_input_driver_path(modifier, identifier) - ) - driver.driver.type = "SCRIPTED" - driver.driver.expression = expression - return driver - - -def add_depth_plane_processing( - plane_object, - depth_image, - color_image, - material, - frame_start=1, - frame_duration=1, - use_cyclic=False, - animate_frame=False, -): - node_group = load_depth_plane_node_group() - modifier = plane_object.modifiers.new( - name=DEPTH_PLANE_NODE_GROUP_NAME, - type="NODES", - ) - modifier.node_group = node_group - modifier.show_group_selector = False - for name, value in ( - ("Depth Image", depth_image), - ("Color Image", color_image), - ("Material", material), - ("Frame", 1), - ("Frame Start", frame_start), - ("Frames", max(int(frame_duration or 1), 1)), - ("Cyclic", use_cyclic), - ): - identifier = modifier_input_identifier(node_group, name) - set_modifier_input(modifier, identifier, value) - if animate_frame: - frame_identifier = modifier_input_identifier(node_group, "Frame") - drive_modifier_input(plane_object, modifier, frame_identifier) - plane_object.update_tag(refresh={"DATA"}) - return modifier - - -def add_moge_depth_plane_processing( - plane_object, - depth_image, - color_image, - material, - frame_start=1, - frame_duration=1, - use_cyclic=False, - animate_frame=False, -): - node_group = load_depth_plane_node_group( - MOGE_DEPTH_PLANE_NODE_GROUP_NAME - ) - modifier = plane_object.modifiers.new( - name=MOGE_DEPTH_PLANE_NODE_GROUP_NAME, - type="NODES", - ) - modifier.node_group = node_group - modifier.show_group_selector = False - for name, value in ( - ("Depth Image", depth_image), - ("Color Image", color_image), - ("Material", material), - ("Frame", 1), - ("Frame Start", frame_start), - ("Frames", max(int(frame_duration or 1), 1)), - ("Cyclic", use_cyclic), - ): - identifier = modifier_input_identifier(node_group, name) - set_modifier_input(modifier, identifier, value) - if animate_frame: - frame_identifier = modifier_input_identifier(node_group, "Frame") - drive_modifier_input(plane_object, modifier, frame_identifier) - plane_object.update_tag(refresh={"DATA"}) - return modifier - - -def finish_depth_plane( - context, - source_object, - plane_object, - delete_source, -): - source_name = source_object.name - for selected in context.selected_objects: - selected.select_set(False) - plane_object.select_set(True) - context.view_layer.objects.active = plane_object - if delete_source: - bpy.data.objects.remove(source_object, do_unlink=True) - plane_object.name = source_name - return plane_object - - -def build_image_depth_plane( - context, - source_object, - delete_source, - shadeless=False, - color_image=None, - depth_frames_directory=None, - depth_image=None, -): - source_name = source_object.name - frame_count = 1 - frames_image = None - if depth_frames_directory is not None: - frames_image = create_sequence_image( - depth_data_name(source_object), - depth_frames_directory, - frame_prefix="color.", - ) - frame_count = len( - sorted_frame_paths(depth_frames_directory, "color.") - ) - ( - source_frame_start, - source_frame_offset, - source_frame_duration, - source_use_cyclic, - ) = image_user_settings(source_object) - uses_extracted = frames_image is not None - texture_frame_count = ( - frame_count if uses_extracted else source_frame_duration - ) - texture_frame_offset = 0 if uses_extracted else source_frame_offset - depth_source_image = depth_image or source_object.data - if depth_image is None: - if frames_image is not None: - depth_source_image = frames_image - elif not is_animated_image(source_object.data): - depth_source_image = copy_depth_image(source_object) - material_image = color_image - if material_image is None: - raise RuntimeError("Depth Plane material image is required") - mesh = create_plane_mesh(source_object, subdivision_level=None) - plane_object = bpy.data.objects.new(source_name, mesh) - context.collection.objects.link(plane_object) - plane_object.matrix_world = source_object.matrix_world.copy() - - depth_scale = abs(float(plane_object.scale.z)) - material = create_material( - source_name, - material_image, - depth_scale, - displacement_image=depth_source_image, - displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL, - displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH, - shadeless=shadeless, - frame_count=frame_count, - frame_start=source_frame_start, - color_frame_offset=texture_frame_offset, - color_frame_duration=texture_frame_count, - depth_frame_offset=texture_frame_offset, - depth_frame_duration=texture_frame_count, - use_cyclic=source_use_cyclic, - scene=context.scene, - ) - mesh.materials.append(material) - add_depth_plane_processing( - plane_object, - depth_source_image, - material_image, - material, - frame_start=source_frame_start, - frame_duration=frame_count, - use_cyclic=source_use_cyclic, - animate_frame=is_animated_image(source_object.data), - ) - plane_object["snapform_frame_count"] = frame_count - return finish_depth_plane( - context, - source_object, - plane_object, - delete_source, - ) - - -def build_depth_plane( - context, - source_object, - npz_path, - delete_source, - shadeless=False, - depth_name="", - color_image=None, - depth_frames_directory=None, -): - source_name = source_object.name - raw_depth_arrays = extract_depth_arrays(npz_path) - raw_depth = raw_depth_arrays[0] - intrinsics = extract_intrinsics_array(npz_path) - raw_depth_values = prepare_depth_values(raw_depth) - bounds, reference_depth = camera_plane_bounds( - source_object, - raw_depth_values, - intrinsics, - DISPLACEMENT_STRENGTH, - ) - depth_arrays = [ - depth_displacement_values(frame) for frame in raw_depth_arrays - ] - depth_name = depth_name or depth_data_name(source_object) - frame_count = len(depth_arrays) - material_image = color_image - if material_image is None: - raise RuntimeError("Depth Plane material image is required") - ( - source_frame_start, - source_frame_offset, - source_frame_duration, - source_use_cyclic, - ) = image_user_settings(source_object) - mesh = create_plane_mesh(source_object, bounds, None) - plane_object = bpy.data.objects.new( - source_name, - mesh, - ) - context.collection.objects.link(plane_object) - depth_scale = apply_plane_transform( - plane_object, - source_object, - bounds, - ) - apply_object_scale_to_mesh(plane_object) - scaled_depth_arrays = [ - scale_depth_values(frame, depth_scale) for frame in depth_arrays - ] - depth_image = create_depth_image( - depth_name, - scaled_depth_arrays[0], - ) - if ( - is_animated_image(source_object.data) - and depth_frames_directory is not None - ): - bpy.data.images.remove(depth_image, do_unlink=True) - depth_image = create_depth_sequence_image( - depth_name, - scaled_depth_arrays, - context.scene, - depth_frames_directory, - ) - - material = create_material( - source_name, - material_image, - depth_scale, - displacement_image=depth_image, - displacement_midlevel=DISPLACEMENT_MIDLEVEL, - displacement_scale=DISPLACEMENT_STRENGTH, - shadeless=shadeless, - frame_count=frame_count, - frame_start=source_frame_start, - color_frame_offset=0 if color_image is not None else source_frame_offset, - color_frame_duration=( - frame_count if color_image is not None else source_frame_duration - ), - depth_frame_offset=0, - depth_frame_duration=frame_count, - use_cyclic=source_use_cyclic, - scene=context.scene, - ) - mesh.materials.append(material) - add_depth_plane_processing( - plane_object, - depth_image, - material_image, - material, - frame_start=source_frame_start, - frame_duration=frame_count, - use_cyclic=source_use_cyclic, - animate_frame=is_animated_image(source_object.data), - ) - plane_object["snapform_reference_depth"] = reference_depth - plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0] - plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2] - plane_object["snapform_depth_scale"] = depth_scale - plane_object["snapform_frame_count"] = frame_count - - return finish_depth_plane( - context, - source_object, - plane_object, - delete_source, - ) - - -def build_moge2_depth_plane( - context, - source_object, - npz_path, - delete_source, - shadeless=False, - color_image=None, - data_frames_directory=None, -): - import numpy as np - - try: - with np.load(npz_path, allow_pickle=False) as archive: - depth = np.asarray(archive["depth"], dtype=np.float32) - mask = np.asarray(archive["mask"], dtype=bool) - intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32) - except (KeyError, OSError) as error: - raise ValueError(f"Invalid MoGe-2 NPZ: {npz_path}") from error - if depth.ndim == 2: - depth = depth[None, ...] - mask = mask[None, ...] - if intrinsics.ndim == 2: - intrinsics = intrinsics[None, ...] - if not (len(depth) == len(mask) == len(intrinsics)): - raise ValueError("MoGe-2 arrays have different frame counts") - - first_valid = mask[0] & np.isfinite(depth[0]) & (depth[0] > 0.0) - if not first_valid.any(): - raise ValueError("MoGe-2 mask contains no valid depth") - reference_depth = float(np.median(depth[0][first_valid])) - height, width = depth[0].shape - pixel_intrinsics = intrinsics[0].copy() - pixel_intrinsics[0] *= width - pixel_intrinsics[1] *= height - bounds, _ = camera_plane_bounds( - source_object, - depth[0], - pixel_intrinsics, - 1.0, - reference_depth=reference_depth, - ) - displacement = [reference_depth - frame for frame in depth] - source_frame_start, _, _, source_use_cyclic = image_user_settings( - source_object - ) - mesh = create_plane_mesh(source_object, bounds, None) - plane_object = bpy.data.objects.new(source_object.name, mesh) - context.collection.objects.link(plane_object) - depth_scale = apply_plane_transform(plane_object, source_object, bounds) - apply_object_scale_to_mesh(plane_object) - depth_image = create_depth_image( - depth_data_name(source_object), - displacement[0] * depth_scale, - ) - if ( - is_animated_image(source_object.data) - and data_frames_directory is not None - ): - bpy.data.images.remove(depth_image, do_unlink=True) - depth_image = create_moge_sequence_image( - depth_data_name(source_object), - [frame * depth_scale for frame in displacement], - context.scene, - data_frames_directory, - "moge_depth", - "depth", - ) - material_image = color_image - if material_image is None: - raise RuntimeError("Depth Plane material image is required") - material = create_moge_material( - source_object.name, - material_image, - depth_image, - depth_scale, - shadeless=shadeless, - frame_count=len(depth), - frame_start=source_frame_start, - use_cyclic=source_use_cyclic, - scene=context.scene, - ) - mesh.materials.append(material) - modifier = add_moge_depth_plane_processing( - plane_object, - depth_image, - material_image, - material, - frame_start=source_frame_start, - frame_duration=len(depth), - use_cyclic=source_use_cyclic, - animate_frame=is_animated_image(source_object.data), - ) - plane_object["snapform_reference_depth"] = reference_depth - plane_object["snapform_depth_scale"] = depth_scale - plane_object["snapform_frame_count"] = len(depth) - return finish_depth_plane( - context, - source_object, - plane_object, - delete_source, - ) - - -class RemoveImageBackground(bpy.types.Operator): - bl_idname = "snapform.remove_image_background" - bl_label = "Remove Background" - bl_description = "Remove the background from this Image Empty with BEN2" - bl_options = {"REGISTER"} - - @classmethod - def poll(cls, context): - from ..job_runtime import runtime - from ..properties import required_models_ready - - return ( - is_image_empty(context.object) - and runtime.environment_ready() - and not runtime.server_busy() - and required_models_ready() - ) - - def execute(self, context): - source_object = context.object - total_frames = video_total_frames(source_object.data) - if total_frames > 1: - context.scene.snapform_settings.video_frames = ( - default_extract_frames(source_object) - ) - try: - source_path, delete_input = prepare_image_input( - source_object.data, - context.scene, - ) - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - - result = bpy.ops.snapform.process_image_background( - input_path=str(source_path), - source_object_name=source_object.name, - delete_input=delete_input, - ) - if "CANCELLED" not in result: - return {"FINISHED"} - if delete_input: - remove_temporary_input(source_path) - return {"CANCELLED"} - - -class ConvertToDepthPlane(bpy.types.Operator): - bl_idname = "snapform.convert_to_depth_plane" - bl_label = "Convert to Depth Plane" - bl_description = "Convert an image to a displaced depth plane" - bl_options = {"REGISTER"} - - shadeless: bpy.props.BoolProperty( - name="Shadeless", - description="Display the source color without scene lighting", - default=False, - ) - remove_background: bpy.props.BoolProperty( - name="Remove Background", - description="Remove the background during conversion with BEN2", - default=False, - ) - model: bpy.props.EnumProperty( - name="Depth From", - description="Model or original image used to build depth", - items=image_depth_model_items, - default=0, - ) - process_res_factor: process_resolution_scale_property() - resolution_level: bpy.props.IntProperty( - name="Resolution Level", - description="MoGe-2 inference resolution level; higher is finer and slower", - default=5, - min=0, - max=9, - ) - keep_original: bpy.props.BoolProperty( - name="Keep Original", - description="Keep the source Image Empty after conversion", - default=False, - ) - - @classmethod - def poll(cls, context): - from ..job_runtime import runtime - from ..properties import required_models_ready - - return ( - is_image_empty(context.object) - and runtime.environment_ready() - and not runtime.server_busy() - and required_models_ready() - ) - - def invoke(self, context, _event): - return context.window_manager.invoke_props_dialog(self, width=340) - - def draw(self, context): - layout = self.layout - layout.prop(self, "model") - if self.model in {"SMALL", "BASE"}: - draw_process_resolution(layout, self, context.object.data) - elif self.model.startswith("MOGE2_"): - layout.prop(self, "resolution_level") - layout.prop(self, "remove_background") - layout.prop(self, "shadeless") - layout.separator(type="LINE") - layout.prop(self, "keep_original") - - def execute(self, context): - source_object = context.object - if video_total_frames(source_object.data) > 1: - context.scene.snapform_settings.video_frames = ( - default_extract_frames(source_object) - ) - if self.model == "ORIGINAL": - try: - source_path, delete_input = prepare_image_input( - source_object.data, - context.scene, - ) - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - run_original = ( - bpy.ops.snapform.generate_original_background - if self.remove_background - else bpy.ops.snapform.extract_original_frames - ) - result = run_original( - input_path=str(source_path), - source_object_name=source_object.name, - delete_source=not self.keep_original, - delete_input=delete_input, - shadeless=self.shadeless, - ) - if "CANCELLED" not in result: - return {"FINISHED"} - if delete_input: - remove_temporary_input(source_path) - return {"CANCELLED"} - - if self.model.startswith("MOGE2_"): - try: - source_path, delete_input = prepare_image_input( - source_object.data, - context.scene, - ) - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - result = bpy.ops.snapform.generate_moge2_depth( - input_path=str(source_path), - source_object_name=source_object.name, - resolution_level=self.resolution_level, - model=self.model, - delete_source=not self.keep_original, - delete_input=delete_input, - shadeless=self.shadeless, - remove_background=self.remove_background, - ) - if "CANCELLED" not in result: - return {"FINISHED"} - if delete_input: - remove_temporary_input(source_path) - return {"CANCELLED"} - - image = source_object.data - try: - process_res = scaled_process_resolution( - image, - self.process_res_factor, - ) - except ValueError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - try: - source_path, delete_input = prepare_image_input( - image, - context.scene, - ) - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - - result = bpy.ops.snapform.generate_da3_depth( - input_path=str(source_path), - source_object_name=source_object.name, - depth_data_name=depth_data_name(source_object), - model=self.model, - process_res=process_res, - delete_source=not self.keep_original, - delete_input=delete_input, - shadeless=self.shadeless, - remove_background=self.remove_background, - ) - if "CANCELLED" in result: - if delete_input: - remove_temporary_input(source_path) - return {"CANCELLED"} - return {"FINISHED"} - - -class SnapFormImageMenu(bpy.types.Menu): - bl_idname = "SNAPFORM_MT_image_context" - bl_label = "SnapForm" - - def draw(self, _context): - from ..job_runtime import runtime - from ..properties import required_models_ready - - layout = self.layout - layout.operator_context = "INVOKE_DEFAULT" - environment_ready = runtime.environment_ready() - if not environment_ready: - layout.operator( - "snapform.install_environment", - text="Install Environment", - icon="IMPORT", - ) - layout.operator( - "snapform.install_environment_mirror", - text="Install Environment (Mirror)", - icon="IMPORT", - ) - models_ready = required_models_ready() - if environment_ready and not models_ready: - layout.operator( - "snapform.download_required_models", - text="Download Required Models", - icon="IMPORT", - ) - features_ready = environment_ready and models_ready - if not features_ready: - layout.separator() - depth_plane_row = layout.row() - depth_plane_row.enabled = features_ready - depth_plane_row.operator( - "snapform.convert_to_depth_plane", - text="Convert to Depth Plane", - icon="MOD_DISPLACE", - ) - surface_row = layout.row() - surface_row.enabled = features_ready - surface_row.operator( - "snapform.convert_to_surface", - text="Convert to Surface", - icon="MESH_GRID", - ) - layout.separator() - background_row = layout.row() - background_row.enabled = features_ready - background_row.operator( - "snapform.remove_image_background", - text="Remove Background", - icon="IMAGE_ALPHA", - ) - - -def draw_image_context_menu(self, context): - if not is_image_empty(context.object): - return - self.layout.separator() - self.layout.menu( - SnapFormImageMenu.bl_idname, - text="SnapForm", - icon="PLUGIN", - ) diff --git a/src/snapform/menu.py b/src/snapform/menu.py new file mode 100644 index 0000000..8063f0e --- /dev/null +++ b/src/snapform/menu.py @@ -0,0 +1,76 @@ +import bpy + +from .operators.image import is_image_empty + + +class SnapFormImageMenu(bpy.types.Menu): + bl_idname = "SNAPFORM_MT_image_context" + bl_label = "SnapForm" + + def draw(self, _context): + from .runtime import runtime + from .properties import required_models_ready + + layout = self.layout + layout.operator_context = "INVOKE_DEFAULT" + layout.operator( + "snapform.balloon_lasso", + text="Balloon Mesh (Lasso)", + icon="MOD_SKIN", + ) + layout.separator() + environment_ready = runtime.environment_ready() + if not environment_ready: + layout.operator( + "snapform.install_environment", + text="Install Environment", + icon="IMPORT", + ) + layout.operator( + "snapform.install_environment_mirror", + text="Install Environment (Mirror)", + icon="IMPORT", + ) + models_ready = required_models_ready() + if environment_ready and not models_ready: + layout.operator( + "snapform.download_required_models", + text="Download Required Models", + icon="IMPORT", + ) + features_ready = environment_ready and models_ready + if not features_ready: + layout.separator() + depth_plane_row = layout.row() + depth_plane_row.enabled = features_ready + depth_plane_row.operator( + "snapform.convert_to_depth_plane", + text="Convert to Depth Plane", + icon="MOD_DISPLACE", + ) + surface_row = layout.row() + surface_row.enabled = features_ready + surface_row.operator( + "snapform.convert_to_surface", + text="Convert to Surface", + icon="MESH_GRID", + ) + layout.separator() + background_row = layout.row() + background_row.enabled = features_ready + background_row.operator( + "snapform.remove_image_background", + text="Remove Background", + icon="IMAGE_ALPHA", + ) + + +def draw_image_context_menu(self, context): + if not is_image_empty(context.object): + return + self.layout.separator() + self.layout.menu( + SnapFormImageMenu.bl_idname, + text="SnapForm", + icon="PLUGIN", + ) diff --git a/src/snapform/operators.py b/src/snapform/operators.py deleted file mode 100644 index 6e43c85..0000000 --- a/src/snapform/operators.py +++ /dev/null @@ -1,735 +0,0 @@ -import shutil -from pathlib import Path - -import bpy - -from .job_runtime import JobOperatorBase, runtime -from .properties import ( - BEN2_MODEL_KEY, - DEFAULT_MOGE2_MODEL_KEY, - REQUIRED_MODEL_KEYS, - model_ready, - model_record, -) - - -IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 -IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize" -BEN2_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} -BEN2_VIDEO_SUFFIXES = { - ".mp4", - ".mov", - ".avi", - ".mkv", - ".webm", - ".m4v", - ".mts", - ".mpg", - ".mpeg", -} -HIDDEN = {"HIDDEN", "SKIP_SAVE"} - - -def input_path(value): - path = Path(bpy.path.abspath(value)) - if path.is_file() or path.is_dir(): - return path - raise ValueError("Select a valid image, video, or frame directory") - - -def background_input_path(value): - path = input_path(value) - if path.is_dir() or path.suffix.lower() in ( - BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES - ): - return path - raise ValueError( - "BEN2 requires an image, video, or directory of image frames" - ) - - -def source_start_frame(source_object_name): - from .blender.depth_plane import extract_frame_range - - source_object = bpy.data.objects.get(source_object_name) - return extract_frame_range(source_object)[0] if source_object else 1 - - -def production_device(): - addon = bpy.context.preferences.addons.get(__package__) - return getattr(getattr(addon, "preferences", None), "device", "AUTO") - - -def require_environment(): - if not runtime.environment_ready(): - raise RuntimeError("Install the SnapForm Environment first") - - -def require_model(model_key): - model = model_record(model_key) - if model is None: - raise RuntimeError(f"Unknown model: {model_key}") - if not model_ready(model_key): - raise RuntimeError(f"Download {model['label']} first") - return model - - -def cleanup_temporary_input(path, delete_input): - if not delete_input or not path: - return - temporary_path = Path(path) - try: - if temporary_path.is_dir(): - shutil.rmtree(temporary_path, ignore_errors=True) - return - temporary_path.unlink(missing_ok=True) - temporary_path.parent.rmdir() - except OSError: - pass - - -def open_job_directory(directory): - directory = Path(directory) - if not directory.is_dir(): - raise RuntimeError(f"Job directory does not exist: {directory}") - result = bpy.ops.wm.path_open(filepath=str(directory)) - if "FINISHED" not in result: - raise RuntimeError(f"Unable to open job directory: {directory}") - - -def push_undo(message): - result = bpy.ops.ed.undo_push(message=message) - if "FINISHED" not in result: - raise RuntimeError(f"Unable to create the {message} Undo step") - - -def build_da3_depth_response( - context, - result, - *, - source_object_name, - depth_data_name, - remove_background, - delete_source, - shadeless, -): - from .blender.depth_plane import ( - build_depth_plane, - create_color_image, - is_animated_image, - is_image_empty, - ) - - directory = result.directory - source_object = bpy.data.objects.get(source_object_name) - if not is_image_empty(source_object): - raise RuntimeError("The source image empty is no longer available") - animated = is_animated_image(source_object.data) - color_image = create_color_image( - source_object, - directory, - assign_to_source=False, - ) - build_depth_plane( - context, - source_object, - result.file("depth"), - delete_source, - shadeless=shadeless, - depth_name=depth_data_name, - color_image=color_image, - depth_frames_directory=directory if animated else None, - ) - push_undo("Convert to Depth Plane") - - -def build_moge2_depth_response( - context, - result, - *, - source_object_name, - remove_background, - delete_source, - shadeless, -): - from .blender.depth_plane import ( - build_moge2_depth_plane, - create_color_image, - is_image_empty, - ) - - directory = result.directory - source_object = bpy.data.objects.get(source_object_name) - if not is_image_empty(source_object): - raise RuntimeError("The source image empty is no longer available") - color_image = create_color_image( - source_object, - directory, - assign_to_source=False, - ) - build_moge2_depth_plane( - context, - source_object, - result.file("depth"), - delete_source, - shadeless=shadeless, - color_image=color_image, - data_frames_directory=directory, - ) - push_undo("Convert to Depth Plane") - - -def build_original_depth_response( - context, - result, - *, - source_object_name, - background, - delete_source, - shadeless, -): - from .blender.depth_plane import ( - build_image_depth_plane, - create_color_image, - is_animated_image, - is_image_empty, - ) - - directory = result.directory - source_object = bpy.data.objects.get(source_object_name) - if not is_image_empty(source_object): - raise RuntimeError("The source image empty is no longer available") - animated = is_animated_image(source_object.data) - color_image = create_color_image( - source_object, - directory, - assign_to_source=False, - ) - build_image_depth_plane( - context, - source_object, - delete_source, - shadeless=shadeless, - color_image=color_image, - depth_frames_directory=directory if animated else None, - ) - push_undo("Convert to Depth Plane") - - -def build_surface_response( - context, - result, - *, - source_object_name, - delete_source, - create_camera, - shadeless, - moge2, -): - from .blender.depth_plane import create_color_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") - color_image = create_color_image( - source_object, - result.directory, - assign_to_source=False, - ) - build_image_mesh( - context, - source_object, - result.file("surface"), - delete_source, - create_camera, - shadeless, - moge2, - color_image=color_image, - normal_path=( - result.directory / "normal.png" - if moge2 - else None - ), - ) - push_undo("Convert to Surface") - - -def apply_background_response(result, source_object_name): - from .blender.depth_plane import create_color_image, is_image_empty - - source_object = bpy.data.objects.get(source_object_name) - if not is_image_empty(source_object): - raise RuntimeError("The source image empty is no longer available") - create_color_image(source_object, result.directory) - push_undo("Remove Background") - - -def da3_parameters(settings, path, model, process_res, start_frame): - return { - "input": str(path), - "model": model, - "process_res": process_res, - "process_res_method": "upper_bound_resize", - "video_frames": settings.video_frames, - "start_frame": start_frame, - "device": production_device().lower(), - "ref_view_strategy": settings.ref_view_strategy, - } - - -def moge2_parameters(settings, path, model, resolution_level, start_frame): - return { - "input": str(path), - "model": model, - "resolution_level": int(resolution_level), - "video_frames": settings.video_frames, - "start_frame": start_frame, - "device": production_device().lower(), - } - - -class DownloadModel(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.download_model" - bl_label = "Download Model" - job_type = "download-model" - - model: bpy.props.StringProperty(options=HIDDEN) - - def request(self, _context): - require_environment() - if not bpy.app.online_access: - raise RuntimeError("Online Access is disabled in Blender Preferences") - model = model_record(self.model) - label = model["label"] if model else self.model - self.starting_message = f"Starting {label} download" - return {"model": self.model} - - -class DownloadRequiredModels(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.download_required_models" - bl_label = "Download Required Models" - job_type = "download-required-models" - starting_message = "Starting required model downloads" - - def request(self, _context): - require_environment() - if not bpy.app.online_access: - raise RuntimeError("Online Access is disabled in Blender Preferences") - return {"models": list(REQUIRED_MODEL_KEYS)} - - -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.model - require_model(model) - return { - **da3_parameters( - settings, - input_path(settings.input_path), - model, - settings.process_res, - 1, - ), - "device": settings.device.lower(), - "inference_mode": "joint", - "remove_background_model": "", - } - - def response(self, _context, result): - open_job_directory(result.directory) - return f"Saved DA3 API prediction to {result.directory}" - - -class DebugMoge2(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.debug_moge2" - bl_label = "Debug MoGe-2" - job_type = "debug-moge2" - starting_message = "Starting MoGe-2 inference" - - def request(self, context): - require_environment() - settings = context.scene.snapform_settings - require_model(settings.moge_model) - return { - **moge2_parameters( - settings, - input_path(settings.input_path), - settings.moge_model, - settings.moge_resolution_level, - 1, - ), - "device": settings.device.lower(), - } - - def response(self, _context, result): - open_job_directory(result.directory) - return f"Exported MoGe-2 data to {result.directory}" - - -class DebugBen2(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.debug_ben2" - bl_label = "Debug BEN2" - job_type = "debug-ben2" - starting_message = "Starting BEN2 background removal" - - def request(self, context): - require_environment() - settings = context.scene.snapform_settings - require_model(BEN2_MODEL_KEY) - return { - "input": str(background_input_path(settings.background_input_path)), - "model": BEN2_MODEL_KEY, - "device": settings.device.lower(), - "video_frames": settings.video_frames, - "start_frame": 1, - } - - def response(self, _context, result): - open_job_directory(result.directory) - return f"Removed background to {result.directory}" - - -class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.generate_da3_depth" - bl_label = "Generate DA3 Depth" - job_type = "generate-depth" - starting_message = "Starting depth estimation" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - depth_data_name: bpy.props.StringProperty(options=HIDDEN) - model: bpy.props.StringProperty(default="BASE", options=HIDDEN) - process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) - delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) - remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - require_model(self.model) - if self.remove_background: - require_model(BEN2_MODEL_KEY) - settings = context.scene.snapform_settings - return { - **da3_parameters( - settings, - input_path(self.input_path), - self.model, - self.process_res, - source_start_frame(self.source_object_name), - ), - "inference_mode": "per_frame", - "remove_background_model": ( - BEN2_MODEL_KEY if self.remove_background else "" - ), - } - - def response(self, context, result): - build_da3_depth_response( - context, - result, - source_object_name=self.source_object_name, - depth_data_name=self.depth_data_name, - remove_background=self.remove_background, - delete_source=self.delete_source, - shadeless=self.shadeless, - ) - return "Depth plane is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.generate_moge2_depth" - bl_label = "Generate MoGe-2 Depth" - job_type = "generate-depth" - starting_message = "Starting MoGe-2 inference" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) - resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) - delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) - remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - require_model(self.model) - if self.remove_background: - require_model(BEN2_MODEL_KEY) - settings = context.scene.snapform_settings - return { - **moge2_parameters( - settings, - input_path(self.input_path), - self.model, - self.resolution_level, - source_start_frame(self.source_object_name), - ), - "remove_background_model": ( - BEN2_MODEL_KEY if self.remove_background else "" - ), - } - - def response(self, context, result): - build_moge2_depth_response( - context, - result, - source_object_name=self.source_object_name, - remove_background=self.remove_background, - delete_source=self.delete_source, - shadeless=self.shadeless, - ) - return "MoGe-2 depth plane is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.generate_da3_surface" - bl_label = "Generate DA3 Surface" - job_type = "generate-surface" - starting_message = "Starting perspective surface generation" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - model: bpy.props.StringProperty(default="BASE", options=HIDDEN) - process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) - delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) - shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) - remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - require_model(self.model) - if self.remove_background: - require_model(BEN2_MODEL_KEY) - settings = context.scene.snapform_settings - return { - **da3_parameters( - settings, - input_path(self.input_path), - self.model, - self.process_res, - source_start_frame(self.source_object_name), - ), - "remove_background_model": ( - BEN2_MODEL_KEY if self.remove_background else "" - ), - } - - def response(self, context, result): - build_surface_response( - context, - result, - source_object_name=self.source_object_name, - delete_source=self.delete_source, - create_camera=self.create_camera, - shadeless=self.shadeless, - moge2=False, - ) - return "Perspective surface is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.generate_moge2_surface" - bl_label = "Generate MoGe-2 Surface" - job_type = "generate-surface" - starting_message = "Starting MoGe-2 inference" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) - resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) - delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) - shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) - remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - require_model(self.model) - if self.remove_background: - require_model(BEN2_MODEL_KEY) - settings = context.scene.snapform_settings - return { - **moge2_parameters( - settings, - input_path(self.input_path), - self.model, - self.resolution_level, - source_start_frame(self.source_object_name), - ), - "remove_background_model": ( - BEN2_MODEL_KEY if self.remove_background else "" - ), - } - - def response(self, context, result): - build_surface_response( - context, - result, - source_object_name=self.source_object_name, - delete_source=self.delete_source, - create_camera=self.create_camera, - shadeless=self.shadeless, - moge2=True, - ) - return "MoGe-2 surface is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class ProcessImageBackground(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.process_image_background" - bl_label = "Process Image Background" - job_type = "remove-background" - starting_message = "Starting Image Empty background removal" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - require_model(BEN2_MODEL_KEY) - settings = context.scene.snapform_settings - return { - "input": str(background_input_path(self.input_path)), - "model": BEN2_MODEL_KEY, - "device": production_device().lower(), - "video_frames": settings.video_frames, - "start_frame": source_start_frame(self.source_object_name), - } - - def response(self, _context, result): - apply_background_response(result, self.source_object_name) - return "Image background is removed" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class GenerateOriginalBackground(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.generate_original_background" - bl_label = "Generate Original Background" - job_type = "remove-background" - starting_message = "Starting background removal" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - require_model(BEN2_MODEL_KEY) - settings = context.scene.snapform_settings - return { - "input": str(background_input_path(self.input_path)), - "model": BEN2_MODEL_KEY, - "device": production_device().lower(), - "video_frames": settings.video_frames, - "start_frame": source_start_frame(self.source_object_name), - } - - def response(self, context, result): - build_original_depth_response( - context, - result, - source_object_name=self.source_object_name, - background=True, - delete_source=self.delete_source, - shadeless=self.shadeless, - ) - return "Depth plane is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class ExtractOriginalFrames(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.extract_original_frames" - bl_label = "Extract Original Frames" - job_type = "extract-frames" - starting_message = "Extracting source frames" - - input_path: bpy.props.StringProperty(options=HIDDEN) - source_object_name: bpy.props.StringProperty(options=HIDDEN) - delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) - delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) - shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) - - def request(self, context): - require_environment() - settings = context.scene.snapform_settings - return { - "input": str(input_path(self.input_path)), - "video_frames": settings.video_frames, - "start_frame": source_start_frame(self.source_object_name), - } - - def response(self, context, result): - build_original_depth_response( - context, - result, - source_object_name=self.source_object_name, - background=False, - delete_source=self.delete_source, - shadeless=self.shadeless, - ) - return "Depth plane is ready" - - def cleanup(self): - cleanup_temporary_input(self.input_path, self.delete_input) - - -class ClearModels(bpy.types.Operator): - bl_idname = "snapform.clear_models" - bl_label = "Clear Models" - - def execute(self, _context): - try: - runtime.clear_resource("model_manager") - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - -JOB_OPERATOR_CLASSES = ( - DownloadModel, - DownloadRequiredModels, - DebugDa3, - DebugMoge2, - DebugBen2, - GenerateDa3Depth, - GenerateMoge2Depth, - GenerateDa3Surface, - GenerateMoge2Surface, - ProcessImageBackground, - GenerateOriginalBackground, - ExtractOriginalFrames, - ClearModels, -) diff --git a/src/snapform/operators/__init__.py b/src/snapform/operators/__init__.py new file mode 100644 index 0000000..7edf027 --- /dev/null +++ b/src/snapform/operators/__init__.py @@ -0,0 +1,34 @@ +from .background import ProcessImageBackground, RemoveImageBackground +from .balloon import BalloonLasso, ExtractBalloon +from .debug import DebugBen2, DebugDa3, DebugMoge2 +from .depth_plane import ( + ConvertToDepthPlane, + ExtractOriginalFrames, + GenerateDa3Depth, + GenerateMoge2Depth, + GenerateOriginalBackground, +) +from .models import ClearModels, DownloadModel, DownloadRequiredModels +from .surface import ConvertToSurface, GenerateDa3Surface, GenerateMoge2Surface + + +CLASSES = ( + DownloadModel, + DownloadRequiredModels, + DebugDa3, + DebugMoge2, + DebugBen2, + GenerateDa3Depth, + GenerateMoge2Depth, + GenerateDa3Surface, + GenerateMoge2Surface, + ProcessImageBackground, + GenerateOriginalBackground, + ExtractOriginalFrames, + ClearModels, + ExtractBalloon, + BalloonLasso, + RemoveImageBackground, + ConvertToDepthPlane, + ConvertToSurface, +) diff --git a/src/snapform/operators/background.py b/src/snapform/operators/background.py new file mode 100644 index 0000000..0f834ff --- /dev/null +++ b/src/snapform/operators/background.py @@ -0,0 +1,112 @@ +import bpy + +from ..properties import BEN2_MODEL_KEY +from ..runtime import JobOperatorBase +from .common import ( + HIDDEN, + background_input_path, + cleanup_temporary_input, + production_device, + push_undo, + require_environment, + require_model, + source_start_frame, +) + +from .image import ( + default_extract_frames, + is_image_empty, + prepare_image_input, + remove_temporary_input, + video_total_frames, +) + + +class RemoveImageBackground(bpy.types.Operator): + bl_idname = "snapform.remove_image_background" + bl_label = "Remove Background" + bl_description = "Remove the background from this Image Empty with BEN2" + bl_options = {"REGISTER"} + + @classmethod + def poll(cls, context): + from ..runtime import runtime + from ..properties import required_models_ready + + return ( + is_image_empty(context.object) + and runtime.environment_ready() + and not runtime.server_busy() + and required_models_ready() + ) + + def execute(self, context): + source_object = context.object + total_frames = video_total_frames(source_object.data) + if total_frames > 1: + context.scene.snapform_settings.video_frames = ( + default_extract_frames(source_object) + ) + try: + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + + result = bpy.ops.snapform.process_image_background( + input_path=str(source_path), + source_object_name=source_object.name, + delete_input=delete_input, + ) + if "CANCELLED" not in result: + return {"FINISHED"} + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + + + + + + +def apply_background_response(result, source_object_name): + from .image import create_color_image, is_image_empty + + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + create_color_image(source_object, result.directory) + push_undo("Remove Background") + + +class ProcessImageBackground(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.process_image_background" + bl_label = "Process Image Background" + job_type = "remove-background" + starting_message = "Starting Image Empty background removal" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + "input": str(background_input_path(self.input_path)), + "model": BEN2_MODEL_KEY, + "device": production_device().lower(), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + } + + def response(self, _context, result): + apply_background_response(result, self.source_object_name) + return "Image background is removed" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) diff --git a/src/snapform/operators/balloon/__init__.py b/src/snapform/operators/balloon/__init__.py new file mode 100644 index 0000000..63ce0d2 --- /dev/null +++ b/src/snapform/operators/balloon/__init__.py @@ -0,0 +1 @@ +from .operator import BalloonLasso, ExtractBalloon diff --git a/src/snapform/operators/balloon/geometry.py b/src/snapform/operators/balloon/geometry.py new file mode 100644 index 0000000..807627e --- /dev/null +++ b/src/snapform/operators/balloon/geometry.py @@ -0,0 +1,346 @@ +import math + +from .poisson import poisson_balloon_heights +from .teddy import teddy_balloon_heights + + +POINT_SPACING = 0.1 +INFLATION = 0.45 + + +def polygon_area(points): + import numpy as np + + points = np.asarray(points, dtype=np.float64) + following = np.roll(points, -1, axis=0) + return 0.5 * float( + np.sum( + points[:, 0] * following[:, 1] + - following[:, 0] * points[:, 1] + ) + ) + + +def resample_closed_polygon(polygon, spacing): + import numpy as np + + points = np.asarray(polygon, dtype=np.float64) + if len(points) < 3: + raise ValueError("The lasso does not define a valid image region") + segments = np.roll(points, -1, axis=0) - points + lengths = np.linalg.norm(segments, axis=1) + points = points[lengths > 1e-8] + if len(points) < 3: + raise ValueError("The lasso region is too small to extract a Balloon") + segments = np.roll(points, -1, axis=0) - points + lengths = np.linalg.norm(segments, axis=1) + perimeter = float(lengths.sum()) + if perimeter <= spacing * 3.0: + raise ValueError("The lasso region is too small to extract a Balloon") + count = max(3, int(round(perimeter / spacing))) + distances = np.arange(count, dtype=np.float64) * perimeter / count + cumulative = np.concatenate(([0.0], np.cumsum(lengths))) + indices = np.searchsorted(cumulative[1:], distances, side="right") + factors = (distances - cumulative[indices]) / lengths[indices] + return points[indices] + segments[indices] * factors[:, None] + + +def smooth_lasso_contour(polygon, spacing, iterations=4): + contour = resample_closed_polygon(polygon, spacing * 0.5) + original_area = abs(polygon_area(contour)) + for _index in range(iterations): + import numpy as np + + contour = ( + np.roll(contour, 1, axis=0) * 0.125 + + contour * 0.75 + + np.roll(contour, -1, axis=0) * 0.125 + ) + smooth_area = abs(polygon_area(contour)) + if original_area > 0.0 and smooth_area > 0.0: + center = contour.mean(axis=0) + contour = center + (contour - center) * math.sqrt( + original_area / smooth_area + ) + contour = resample_closed_polygon(contour, spacing) + if polygon_area(contour) < 0.0: + contour = contour[::-1].copy() + return contour + + +def points_inside_polygon(points, polygon): + import numpy as np + + points = np.asarray(points, dtype=np.float64) + polygon = np.asarray(polygon, dtype=np.float64) + inside = np.zeros(len(points), dtype=bool) + previous = np.roll(polygon, 1, axis=0) + for current, prior in zip(polygon, previous): + crosses = (current[1] > points[:, 1]) != ( + prior[1] > points[:, 1] + ) + denominator = prior[1] - current[1] + if abs(denominator) < 1e-12: + continue + crossing_x = ( + (prior[0] - current[0]) + * (points[:, 1] - current[1]) + / denominator + + current[0] + ) + inside ^= crosses & (points[:, 0] < crossing_x) + return inside + + +def distances_to_polygon(points, polygon): + import numpy as np + + points = np.asarray(points, dtype=np.float64) + polygon = np.asarray(polygon, dtype=np.float64) + squared = np.full(len(points), np.inf, dtype=np.float64) + for start, end in zip(polygon, np.roll(polygon, -1, axis=0)): + segment = end - start + length_squared = float(segment.dot(segment)) + if length_squared <= 1e-16: + continue + factor = np.clip( + (points - start) @ segment / length_squared, + 0.0, + 1.0, + ) + nearest = start + factor[:, None] * segment + squared = np.minimum( + squared, + np.sum((points - nearest) ** 2, axis=1), + ) + return np.sqrt(squared) + + +def uniform_interior_points(contour, spacing): + import numpy as np + + contour = np.asarray(contour, dtype=np.float64) + minimum = contour.min(axis=0) + maximum = contour.max(axis=0) + row_spacing = spacing * math.sqrt(3.0) * 0.5 + row_count = int( + math.ceil((maximum[1] - minimum[1]) / row_spacing) + ) + 1 + column_count = int( + math.ceil((maximum[0] - minimum[0]) / spacing) + ) + 1 + candidates = [] + for row in range(row_count): + y = minimum[1] + row * row_spacing + offset = spacing * 0.5 if row % 2 else 0.0 + for column in range(column_count): + candidates.append( + (minimum[0] + offset + column * spacing, y) + ) + points = np.asarray(candidates, dtype=np.float64) + points = points[points_inside_polygon(points, contour)] + if not len(points): + return points + clearance = distances_to_polygon(points, contour) + return points[clearance >= spacing * 0.4] + + +def metric_transform(world_metric=None): + import numpy as np + + metric = ( + np.eye(2, dtype=np.float64) + if world_metric is None + else np.asarray(world_metric, dtype=np.float64) + ) + if metric.shape != (2, 2): + raise ValueError("The Image Empty has an invalid world transform") + try: + return np.linalg.cholesky(metric) + except np.linalg.LinAlgError as error: + raise ValueError( + "The Image Empty has a degenerate world transform" + ) from error + + +def prepare_planar_points( + polygon, + image_size, + local_bounds, + spacing=POINT_SPACING, + world_metric=None, +): + import numpy as np + + image_width, image_height = image_size + x_min, x_max, y_min, y_max = local_bounds + local_width = x_max - x_min + local_height = y_max - y_min + pixels = np.asarray(polygon, dtype=np.float64) + local = np.empty_like(pixels) + local[:, 0] = x_min + local_width * pixels[:, 0] / image_width + local[:, 1] = y_max - local_height * pixels[:, 1] / image_height + if float(np.ptp(local, axis=0).max()) <= 0.0: + raise ValueError("The lasso does not overlap the source image") + if spacing <= 0.0: + raise ValueError("Balloon point spacing must be greater than zero") + + planar = local @ metric_transform(world_metric) + contour = smooth_lasso_contour(planar, spacing) + interior = uniform_interior_points(contour, spacing) + return contour, interior, spacing + +def boundary_vertices_from_faces(vertex_count, faces): + from collections import Counter + + import numpy as np + + edge_users = Counter() + for face in faces: + for start, end in zip(face, face[1:] + face[:1]): + edge_users[tuple(sorted((start, end)))] += 1 + boundary = np.zeros(vertex_count, dtype=bool) + for (start, end), users in edge_users.items(): + if users == 1: + boundary[start] = True + boundary[end] = True + return boundary + + +def balloon_heights( + algorithm, + planar_points, + planar_faces, + contour, + boundary_mask, + inflation, +): + if algorithm == "POISSON": + return poisson_balloon_heights( + planar_points, + planar_faces, + boundary_mask, + inflation, + ) + if algorithm == "TEDDY": + return teddy_balloon_heights( + planar_points, + planar_faces, + contour, + boundary_mask, + inflation, + ) + raise ValueError(f"Unknown Balloon algorithm: {algorithm}") + +def build_closed_geometry( + planar_points, + planar_faces, + boundary_mask, + heights, + local_bounds, + world_metric=None, + world_z_scale=1.0, +): + import numpy as np + + points = np.asarray(planar_points, dtype=np.float64) + boundary_mask = np.asarray(boundary_mask, dtype=bool) + heights = np.asarray(heights, dtype=np.float64) + local_points = points @ np.linalg.inv(metric_transform(world_metric)) + if world_z_scale <= 1e-10: + raise ValueError("The Image Empty has a degenerate world transform") + local_heights = heights / world_z_scale + x_min, x_max, y_min, y_max = local_bounds + width = x_max - x_min + height = y_max - y_min + vertices = [] + vertex_uv = [] + front_indices = [] + back_indices = [] + for index, point in enumerate(local_points): + u = (point[0] - x_min) / width + v = (point[1] - y_min) / height + front_indices.append(len(vertices)) + vertices.append((point[0], point[1], local_heights[index])) + vertex_uv.append((u, v)) + if boundary_mask[index]: + back_indices.append(front_indices[-1]) + else: + back_indices.append(len(vertices)) + vertices.append((point[0], point[1], -local_heights[index])) + vertex_uv.append((u, v)) + faces = [] + for triangle in planar_faces: + a, b, c = (points[index] for index in triangle) + signed_area = (b[0] - a[0]) * (c[1] - a[1]) - ( + b[1] - a[1] + ) * (c[0] - a[0]) + ordered = triangle if signed_area > 0.0 else triangle[::-1] + faces.append(tuple(front_indices[index] for index in ordered)) + faces.append( + tuple(back_indices[index] for index in reversed(ordered)) + ) + return vertices, faces, np.asarray(vertex_uv, dtype=np.float32) + +def triangulate_balloon( + polygon, + image_size, + local_bounds, + spacing=POINT_SPACING, + world_metric=None, + world_z_scale=1.0, + inflation=INFLATION, + algorithm="POISSON", +): + import numpy as np + from mathutils import Vector + from mathutils.geometry import delaunay_2d_cdt + + contour, interior, spacing = prepare_planar_points( + polygon, + image_size, + local_bounds, + spacing=spacing, + world_metric=world_metric, + ) + input_points = np.concatenate((contour, interior), axis=0) + boundary_count = len(contour) + boundary_edges = [ + (index, (index + 1) % boundary_count) + for index in range(boundary_count) + ] + result = delaunay_2d_cdt( + [Vector(point) for point in input_points], + boundary_edges, + [list(range(boundary_count))], + 1, + max(spacing * 1e-6, 1e-9), + ) + planar_points = np.asarray( + [tuple(point) for point in result[0]], + dtype=np.float64, + ) + planar_faces = [tuple(face) for face in result[2] if len(face) == 3] + if not planar_faces: + raise ValueError("Unable to triangulate the lasso region") + boundary_mask = boundary_vertices_from_faces( + len(planar_points), + planar_faces, + ) + heights = balloon_heights( + algorithm, + planar_points, + planar_faces, + contour, + boundary_mask, + inflation, + ) + return build_closed_geometry( + planar_points, + planar_faces, + boundary_mask, + heights, + local_bounds, + world_metric=world_metric, + world_z_scale=world_z_scale, + ) diff --git a/src/snapform/operators/balloon/operator.py b/src/snapform/operators/balloon/operator.py new file mode 100644 index 0000000..7a240b0 --- /dev/null +++ b/src/snapform/operators/balloon/operator.py @@ -0,0 +1,456 @@ +import json +import math + +import bpy + +from .geometry import INFLATION, POINT_SPACING, triangulate_balloon +from ..image import color_data_name, is_image_empty, plane_bounds +from ..material import TEXTURE_MATERIAL_NODE_GROUP_NAME, create_material +from ..surface import assign_vertex_uv + + +MIN_PATH_DISTANCE = 2.0 + + +def simplify_screen_path(path): + pixels = [] + for point in path: + point = (float(point[0]), float(point[1])) + if not pixels or math.dist(point, pixels[-1]) >= MIN_PATH_DISTANCE: + pixels.append(point) + return pixels + + +def screen_path_to_image_pixels( + path, + region, + region_data, + matrix_world, + local_bounds, + image_size, +): + from bpy_extras import view3d_utils + + inverse = matrix_world.inverted() + x_min, x_max, y_min, y_max = local_bounds + image_width, image_height = image_size + local_width = x_max - x_min + local_height = y_max - y_min + pixels = [] + for point in simplify_screen_path(path): + ray_origin = inverse @ view3d_utils.region_2d_to_origin_3d( + region, + region_data, + point, + ) + ray_direction = inverse.to_3x3() @ view3d_utils.region_2d_to_vector_3d( + region, + region_data, + point, + ) + if abs(ray_direction.z) <= 1e-10: + raise ValueError("The Image Empty is edge-on to the current view") + distance = -ray_origin.z / ray_direction.z + local = ray_origin + ray_direction * distance + pixels.append( + ( + (local.x - x_min) * image_width / local_width, + (y_max - local.y) * image_height / local_height, + ) + ) + return pixels + + +def world_plane_metric(matrix_world): + from mathutils import Vector + + linear = matrix_world.to_3x3() + basis_x = linear @ Vector((1.0, 0.0, 0.0)) + basis_y = linear @ Vector((0.0, 1.0, 0.0)) + basis_z = linear @ Vector((0.0, 0.0, 1.0)) + metric = ( + (basis_x.dot(basis_x), basis_x.dot(basis_y)), + (basis_x.dot(basis_y), basis_y.dot(basis_y)), + ) + return metric, basis_z.length + + +def serialize_matrix(matrix): + return json.dumps( + [float(matrix[row][column]) for row in range(4) for column in range(4)], + separators=(",", ":"), + ) + + +def deserialize_matrix(data): + from mathutils import Matrix + + values = json.loads(data) + if len(values) != 16: + raise ValueError("The stored Image Empty transform is invalid") + return Matrix(tuple(tuple(values[row * 4 : row * 4 + 4]) for row in range(4))) + + +def balloon_geometry_parameters(thickness, resolution): + if thickness <= 0.0: + raise ValueError("Balloon thickness must be greater than zero") + if resolution <= 0.0: + raise ValueError("Balloon resolution must be greater than zero") + return resolution, INFLATION * thickness + + +def extract_balloon_object( + context, + source_object, + polygon, + matrix_world, + algorithm, + thickness=1.0, + resolution=POINT_SPACING, + select_result=True, +): + image = source_object.data + image_size = tuple(image.size) + local_bounds = plane_bounds(source_object) + world_metric, world_z_scale = world_plane_metric(matrix_world) + point_spacing, inflation = balloon_geometry_parameters( + thickness, + resolution, + ) + vertices, faces, vertex_uv = triangulate_balloon( + polygon, + image_size, + local_bounds, + spacing=point_spacing, + world_metric=world_metric, + world_z_scale=world_z_scale, + inflation=inflation, + algorithm=algorithm, + ) + + source_name = source_object.name + balloon_name = source_name + color_image = image.copy() + color_image.name = color_data_name(source_object) + material = None + mesh = None + balloon = None + try: + material = create_material( + balloon_name, + color_image, + 0.0, + surface_group_name=TEXTURE_MATERIAL_NODE_GROUP_NAME, + scene=context.scene, + ) + mesh = bpy.data.meshes.new(balloon_name) + mesh.from_pydata(vertices, [], faces) + mesh.update() + assign_vertex_uv(mesh, vertex_uv) + for face in mesh.polygons: + face.use_smooth = True + + balloon = bpy.data.objects.new(balloon_name, mesh) + context.collection.objects.link(balloon) + balloon.matrix_world = matrix_world.copy() + balloon["snapform_source_type"] = "BALLOON" + balloon["snapform_balloon_algorithm"] = algorithm + balloon["snapform_source_object"] = source_name + balloon["snapform_inflation"] = inflation + balloon["snapform_resolution"] = point_spacing + balloon.data.materials.append(material) + except Exception: + if balloon is not None: + bpy.data.objects.remove(balloon, do_unlink=True) + if mesh is not None and mesh.users == 0: + bpy.data.meshes.remove(mesh) + if material is not None and material.users == 0: + bpy.data.materials.remove(material) + if color_image.users == 0: + bpy.data.images.remove(color_image) + raise + + bpy.ops.object.select_all(action="DESELECT") + selected_object = balloon if select_result else source_object + selected_object.select_set(True) + context.view_layer.objects.active = selected_object + return balloon + + +def thickness_property(): + return bpy.props.FloatProperty( + name="Thickness", + description="Scale the Balloon thickness while keeping the lasso region", + default=1.0, + min=0.05, + soft_max=5.0, + ) + + +def mode_property(): + return bpy.props.EnumProperty( + name="Mode", + description="Algorithm used to inflate the saved lasso region", + items=( + ( + "POISSON", + "Poisson", + "Use a global Poisson field for smooth inflation", + ), + ( + "TEDDY", + "Teddy / Chordal Axis", + "Use a Teddy-style chordal axis and local oval sections", + ), + ), + default="POISSON", + ) + + +def resolution_property(): + return bpy.props.FloatProperty( + name="Resolution", + description=( + "World-space distance between sampled mesh points; smaller " + "values create more geometry" + ), + default=POINT_SPACING, + min=0.005, + soft_min=0.02, + soft_max=0.5, + max=10.0, + subtype="DISTANCE", + unit="LENGTH", + ) + + +def hidden_string_property(): + return bpy.props.StringProperty(options={"HIDDEN", "SKIP_SAVE"}) + + +def hidden_bool_property(default=False): + return bpy.props.BoolProperty( + default=default, + options={"HIDDEN", "SKIP_SAVE"}, + ) + + +def event_over_interface_region(context, event): + area = getattr(context, "area", None) + if area is None: + return False + mouse_x = getattr(event, "mouse_x", -1) + mouse_y = getattr(event, "mouse_y", -1) + for region in area.regions: + if region.type == "WINDOW": + continue + if ( + region.x <= mouse_x < region.x + region.width + and region.y <= mouse_y < region.y + region.height + ): + return True + return False + + +VIEW_NAVIGATION_EVENTS = { + "MIDDLEMOUSE", + "WHEELUPMOUSE", + "WHEELDOWNMOUSE", + "WHEELINMOUSE", + "WHEELOUTMOUSE", + "WHEELLEFTMOUSE", + "WHEELRIGHTMOUSE", + "TRACKPADPAN", + "TRACKPADZOOM", + "MOUSEROTATE", + "MOUSESMARTZOOM", + "NDOF_MOTION", +} + + +class ExtractBalloon(bpy.types.Operator): + bl_idname = "snapform.extract_balloon" + bl_label = "Extract Balloon" + bl_description = "Build one closed Balloon mesh from a saved lasso region" + bl_options = {"REGISTER", "UNDO"} + + mode: mode_property() + thickness: thickness_property() + resolution: resolution_property() + source_object_name: hidden_string_property() + source_matrix_data: hidden_string_property() + selection_data: hidden_string_property() + keep_source_selected: hidden_bool_property() + + def draw(self, context): + layout = self.layout + layout.prop(self, "mode") + layout.prop(self, "thickness") + layout.prop(self, "resolution") + + def execute(self, context): + source_object = bpy.data.objects.get(self.source_object_name) + if not is_image_empty(source_object): + self.report({"ERROR"}, "The source Image Empty is no longer available") + return {"CANCELLED"} + try: + polygon = json.loads(self.selection_data) + if len(polygon) < 3: + raise ValueError("The stored lasso region is invalid") + matrix_world = deserialize_matrix(self.source_matrix_data) + extract_balloon_object( + context, + source_object, + polygon, + matrix_world, + self.mode, + thickness=self.thickness, + resolution=self.resolution, + select_result=not self.keep_source_selected, + ) + except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + +class BalloonLasso(bpy.types.Operator): + bl_idname = "snapform.balloon_lasso" + bl_label = "Balloon Mesh" + bl_description = "Continuously draw lasso regions and extract Balloon meshes" + bl_options = {"INTERNAL"} + + mode: mode_property() + thickness: thickness_property() + resolution: resolution_property() + source_object_name: hidden_string_property() + source_matrix_data: hidden_string_property() + + _path = None + _drawing = False + _handle = None + + @classmethod + def poll(cls, context): + return context.area is not None and context.area.type == "VIEW_3D" + + def _draw_lasso(self, context): + if not self._path: + return + import gpu + from gpu_extras.batch import batch_for_shader + + vertices = list(self._path) + if len(vertices) > 2: + vertices.append(vertices[0]) + shader = gpu.shader.from_builtin("UNIFORM_COLOR") + batch = batch_for_shader(shader, "LINE_STRIP", {"pos": vertices}) + gpu.state.blend_set("ALPHA") + gpu.state.line_width_set(2.0) + shader.bind() + shader.uniform_float("color", (0.15, 0.65, 1.0, 0.95)) + batch.draw(shader) + gpu.state.line_width_set(1.0) + gpu.state.blend_set("NONE") + + def _finish(self, context): + if self._handle is not None: + bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW") + self._handle = None + if context.area is not None: + context.area.header_text_set(None) + context.area.tag_redraw() + + def invoke(self, context, event): + source_object = context.active_object + if not is_image_empty(source_object): + self.report({"ERROR"}, "Select an Image Empty before extracting") + return {"CANCELLED"} + self.source_object_name = source_object.name + self.source_matrix_data = serialize_matrix(source_object.matrix_world) + self._path = [] + self._drawing = False + self._handle = bpy.types.SpaceView3D.draw_handler_add( + self._draw_lasso, + (context,), + "WINDOW", + "POST_PIXEL", + ) + context.window_manager.modal_handler_add(self) + context.area.header_text_set( + "Balloon Mesh: drag LMB repeatedly, navigate normally, Esc to finish" + ) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + def modal(self, context, event): + if event.type == "ESC": + self._finish(context) + return {"CANCELLED"} + + if ( + event.type in VIEW_NAVIGATION_EVENTS + or event_over_interface_region(context, event) + or (event.type == "LEFTMOUSE" and event.alt) + ): + return {"PASS_THROUGH"} + + if event.type == "RIGHTMOUSE" and event.value == "PRESS": + self._drawing = False + self._path = [] + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + if event.type == "LEFTMOUSE" and event.value == "PRESS": + self._drawing = True + self._path = [(event.mouse_region_x, event.mouse_region_y)] + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + if event.type == "MOUSEMOVE" and self._drawing: + self._path.append((event.mouse_region_x, event.mouse_region_y)) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + if event.type == "LEFTMOUSE" and event.value == "RELEASE" and self._drawing: + self._drawing = False + path = simplify_screen_path(self._path) + self._path = [] + if len(path) < 3: + self.report({"WARNING"}, "Draw a larger lasso") + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + source_object = bpy.data.objects.get(self.source_object_name) + if not is_image_empty(source_object): + self.report({"ERROR"}, "The source Image Empty is no longer available") + self._finish(context) + return {"CANCELLED"} + try: + matrix_world = deserialize_matrix(self.source_matrix_data) + polygon = screen_path_to_image_pixels( + path, + context.region, + context.region_data, + matrix_world, + plane_bounds(source_object), + tuple(source_object.data.size), + ) + result = bpy.ops.snapform.extract_balloon( + "EXEC_DEFAULT", + source_object_name=self.source_object_name, + source_matrix_data=self.source_matrix_data, + selection_data=json.dumps(polygon, separators=(",", ":")), + mode=self.mode, + thickness=self.thickness, + resolution=self.resolution, + keep_source_selected=True, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to extract the Balloon") + except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error: + self.report({"ERROR"}, str(error)) + context.area.tag_redraw() + return {"RUNNING_MODAL"} + + return {"PASS_THROUGH"} diff --git a/src/snapform/operators/balloon/poisson.py b/src/snapform/operators/balloon/poisson.py new file mode 100644 index 0000000..86381ca --- /dev/null +++ b/src/snapform/operators/balloon/poisson.py @@ -0,0 +1,123 @@ +import math + + +def cotangent_poisson_field( + planar_points, + planar_faces, + boundary_mask, + tolerance=1e-8, +): + import numpy as np + + points = np.asarray(planar_points, dtype=np.float64) + boundary_mask = np.asarray(boundary_mask, dtype=bool) + loads = np.zeros(len(points), dtype=np.float64) + edge_weights = {} + + def add_edge(start, end, weight): + edge = tuple(sorted((int(start), int(end)))) + edge_weights[edge] = edge_weights.get(edge, 0.0) + weight + + for face in planar_faces: + first, second, third = (int(index) for index in face) + a, b, c = points[[first, second, third]] + cross = abs( + (b[0] - a[0]) * (c[1] - a[1]) + - (b[1] - a[1]) * (c[0] - a[0]) + ) + if cross <= 1e-14: + continue + area = cross * 0.5 + loads[[first, second, third]] += area / 3.0 + cotangent_a = float((b - a).dot(c - a) / cross) + cotangent_b = float((a - b).dot(c - b) / cross) + cotangent_c = float((a - c).dot(b - c) / cross) + add_edge(second, third, cotangent_a * 0.5) + add_edge(first, third, cotangent_b * 0.5) + add_edge(first, second, cotangent_c * 0.5) + + interior = np.flatnonzero(~boundary_mask) + field = np.zeros(len(points), dtype=np.float64) + if not len(interior): + return field + compact = np.full(len(points), -1, dtype=np.int64) + compact[interior] = np.arange(len(interior), dtype=np.int64) + diagonal = np.zeros(len(interior), dtype=np.float64) + edge_starts = [] + edge_ends = [] + weights = [] + for (start, end), weight in edge_weights.items(): + start_compact = compact[start] + end_compact = compact[end] + if start_compact >= 0: + diagonal[start_compact] += weight + if end_compact >= 0: + diagonal[end_compact] += weight + if start_compact >= 0 and end_compact >= 0: + edge_starts.append(start_compact) + edge_ends.append(end_compact) + weights.append(weight) + if np.any(diagonal <= 1e-14): + raise ValueError("Unable to build a stable Poisson Balloon field") + + edge_starts = np.asarray(edge_starts, dtype=np.int64) + edge_ends = np.asarray(edge_ends, dtype=np.int64) + weights = np.asarray(weights, dtype=np.float64) + + def multiply(values): + result = diagonal * values + if len(weights): + np.add.at(result, edge_starts, -weights * values[edge_ends]) + np.add.at(result, edge_ends, -weights * values[edge_starts]) + return result + + right_hand_side = loads[interior] + solution = np.zeros(len(interior), dtype=np.float64) + residual = right_hand_side.copy() + preconditioned = residual / diagonal + direction = preconditioned.copy() + residual_product = float(residual.dot(preconditioned)) + target = max( + float(np.linalg.norm(right_hand_side)) * tolerance, + 1e-12, + ) + max_iterations = max(64, min(len(interior) * 4, 4000)) + for _iteration in range(max_iterations): + if float(np.linalg.norm(residual)) <= target: + break + multiplied = multiply(direction) + denominator = float(direction.dot(multiplied)) + if denominator <= 1e-20 or not math.isfinite(denominator): + raise ValueError("Unable to solve the Poisson Balloon field") + step = residual_product / denominator + solution += step * direction + residual -= step * multiplied + preconditioned = residual / diagonal + next_product = float(residual.dot(preconditioned)) + if not math.isfinite(next_product): + raise ValueError("Unable to solve the Poisson Balloon field") + if abs(residual_product) <= 1e-30: + break + direction = preconditioned + ( + next_product / residual_product + ) * direction + residual_product = next_product + + field[interior] = np.maximum(solution, 0.0) + return field + + +def poisson_balloon_heights( + planar_points, + planar_faces, + boundary_mask, + inflation, +): + import numpy as np + + field = cotangent_poisson_field( + planar_points, + planar_faces, + boundary_mask, + ) + return 2.0 * inflation * np.sqrt(field) diff --git a/src/snapform/operators/balloon/teddy.py b/src/snapform/operators/balloon/teddy.py new file mode 100644 index 0000000..a97ad82 --- /dev/null +++ b/src/snapform/operators/balloon/teddy.py @@ -0,0 +1,218 @@ + +def distances_to_contour(points, contour): + import numpy as np + + points = np.asarray(points, dtype=np.float64) + contour = np.asarray(contour, dtype=np.float64) + squared = np.full(len(points), np.inf, dtype=np.float64) + for start, end in zip(contour, np.roll(contour, -1, axis=0)): + segment = end - start + length_squared = float(segment.dot(segment)) + if length_squared <= 1e-16: + continue + factor = np.clip( + (points - start) @ segment / length_squared, + 0.0, + 1.0, + ) + nearest = start + factor[:, None] * segment + squared = np.minimum( + squared, + np.sum((points - nearest) ** 2, axis=1), + ) + return np.sqrt(squared) + + +def chordal_axis(contour, epsilon=1e-9): + import numpy as np + from mathutils import Vector + from mathutils.geometry import delaunay_2d_cdt + + contour = np.asarray(contour, dtype=np.float64) + boundary_count = len(contour) + boundary_edges = [ + (index, (index + 1) % boundary_count) + for index in range(boundary_count) + ] + result = delaunay_2d_cdt( + [Vector(point) for point in contour], + boundary_edges, + [list(range(boundary_count))], + 1, + epsilon, + ) + points = np.asarray( + [tuple(point) for point in result[0]], + dtype=np.float64, + ) + triangles = [tuple(face) for face in result[2] if len(face) == 3] + boundary = { + tuple(sorted(edge)) + for edge in boundary_edges + } + node_indices = {} + nodes = [] + segments = set() + + def edge_node(edge): + key = ("edge", *tuple(sorted(edge))) + if key not in node_indices: + node_indices[key] = len(nodes) + nodes.append((points[edge[0]] + points[edge[1]]) * 0.5) + return node_indices[key] + + for face_index, face in enumerate(triangles): + edges = [ + (face[0], face[1]), + (face[1], face[2]), + (face[2], face[0]), + ] + internal = [ + edge + for edge in edges + if tuple(sorted(edge)) not in boundary + ] + internal_nodes = [edge_node(edge) for edge in internal] + if len(internal_nodes) == 2: + segments.add(tuple(sorted(internal_nodes))) + elif len(internal_nodes) == 3: + key = ("face", face_index) + center_index = len(nodes) + node_indices[key] = center_index + nodes.append(points[list(face)].mean(axis=0)) + for midpoint_index in internal_nodes: + segments.add(tuple(sorted((center_index, midpoint_index)))) + + if not nodes: + center = contour.mean(axis=0) + return np.asarray((center,), dtype=np.float64), [], np.asarray( + (distances_to_contour((center,), contour)[0],), + dtype=np.float64, + ) + + nodes = np.asarray(nodes, dtype=np.float64) + radii = distances_to_contour(nodes, contour) + active = set(segments) + while len(active) > 1: + degree = [0] * len(nodes) + for start, end in active: + degree[start] += 1 + degree[end] += 1 + removable = [] + for edge in active: + start, end = edge + leaf = start if degree[start] == 1 else end if degree[end] == 1 else None + if leaf is None: + continue + length = float(np.linalg.norm(nodes[start] - nodes[end])) + if length < max(radii[leaf] * 0.5, epsilon): + removable.append(edge) + if not removable: + break + active.difference_update(removable) + + if active: + used = sorted({index for edge in active for index in edge}) + remap = {old: new for new, old in enumerate(used)} + nodes = nodes[used] + radii = radii[used] + active = { + tuple(sorted((remap[start], remap[end]))) + for start, end in active + } + else: + widest = int(np.argmax(radii)) + nodes = nodes[[widest]] + radii = radii[[widest]] + + return nodes, sorted(active), radii + + +def smooth_interior_heights( + heights, + planar_faces, + boundary_mask, + iterations=2, + factor=0.3, +): + import numpy as np + + values = np.asarray(heights, dtype=np.float64).copy() + boundary_mask = np.asarray(boundary_mask, dtype=bool) + neighbors = [set() for _index in range(len(values))] + for face in planar_faces: + for start, end in zip(face, face[1:] + face[:1]): + neighbors[start].add(end) + neighbors[end].add(start) + for _iteration in range(iterations): + updated = values.copy() + for index in np.flatnonzero(~boundary_mask): + adjacent = [ + neighbor + for neighbor in neighbors[index] + if not boundary_mask[neighbor] + ] + if adjacent: + average = float(values[adjacent].mean()) + updated[index] = values[index] * (1.0 - factor) + average * factor + updated[boundary_mask] = 0.0 + values = updated + return values + + +def teddy_balloon_heights( + planar_points, + planar_faces, + contour, + boundary_mask, + inflation, +): + import numpy as np + + points = np.asarray(planar_points, dtype=np.float64) + boundary_mask = np.asarray(boundary_mask, dtype=bool) + nodes, segments, radii = chordal_axis(contour) + squared_heights = np.zeros(len(points), dtype=np.float64) + + for start, end in segments: + axis_start = nodes[start] + axis_end = nodes[end] + axis = axis_end - axis_start + length_squared = float(axis.dot(axis)) + if length_squared <= 1e-16: + continue + factors = np.clip( + (points - axis_start) @ axis / length_squared, + 0.0, + 1.0, + ) + centers = axis_start + factors[:, None] * axis + local_radii = radii[start] + factors * (radii[end] - radii[start]) + distance_squared = np.sum((points - centers) ** 2, axis=1) + squared_heights = np.maximum( + squared_heights, + local_radii * local_radii - distance_squared, + ) + + degree = np.zeros(len(nodes), dtype=np.int64) + for start, end in segments: + degree[start] += 1 + degree[end] += 1 + for index in np.flatnonzero(degree != 2): + distance_squared = np.sum((points - nodes[index]) ** 2, axis=1) + squared_heights = np.maximum( + squared_heights, + radii[index] * radii[index] - distance_squared, + ) + + missing = (~boundary_mask) & (squared_heights <= 0.0) + if np.any(missing): + fallback = distances_to_contour(points[missing], contour) + squared_heights[missing] = fallback * fallback + heights = inflation * np.sqrt(np.maximum(squared_heights, 0.0)) + heights[boundary_mask] = 0.0 + return smooth_interior_heights( + heights, + planar_faces, + boundary_mask, + ) diff --git a/src/snapform/operators/common.py b/src/snapform/operators/common.py new file mode 100644 index 0000000..89619f0 --- /dev/null +++ b/src/snapform/operators/common.py @@ -0,0 +1,111 @@ +import shutil +from pathlib import Path + +import bpy + +from .image import extract_frame_range +from ..runtime import runtime +from ..properties import model_ready, model_record + + +BEN2_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} +BEN2_VIDEO_SUFFIXES = { + ".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".mts", ".mpg", ".mpeg" +} +HIDDEN = {"HIDDEN", "SKIP_SAVE"} +ADDON_PACKAGE = __package__.partition(".")[0] + + +def input_path(value): + path = Path(bpy.path.abspath(value)) + if path.is_file() or path.is_dir(): + return path + raise ValueError("Select a valid image, video, or frame directory") + + +def background_input_path(value): + path = input_path(value) + if path.is_dir() or path.suffix.lower() in ( + BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES + ): + return path + raise ValueError( + "BEN2 requires an image, video, or directory of image frames" + ) + + +def source_start_frame(source_object_name): + source_object = bpy.data.objects.get(source_object_name) + return extract_frame_range(source_object)[0] if source_object else 1 + + +def production_device(): + addon = bpy.context.preferences.addons.get(ADDON_PACKAGE) + return getattr(getattr(addon, "preferences", None), "device", "AUTO") + + +def require_environment(): + if not runtime.environment_ready(): + raise RuntimeError("Install the SnapForm Environment first") + + +def require_model(model_key): + model = model_record(model_key) + if model is None: + raise RuntimeError(f"Unknown model: {model_key}") + if not model_ready(model_key): + raise RuntimeError(f"Download {model['label']} first") + return model + + +def cleanup_temporary_input(path, delete_input): + if not delete_input or not path: + return + temporary_path = Path(path) + try: + if temporary_path.is_dir(): + shutil.rmtree(temporary_path, ignore_errors=True) + return + temporary_path.unlink(missing_ok=True) + temporary_path.parent.rmdir() + except OSError: + pass + + +def open_job_directory(directory): + directory = Path(directory) + if not directory.is_dir(): + raise RuntimeError(f"Job directory does not exist: {directory}") + result = bpy.ops.wm.path_open(filepath=str(directory)) + if "FINISHED" not in result: + raise RuntimeError(f"Unable to open job directory: {directory}") + + +def push_undo(message): + result = bpy.ops.ed.undo_push(message=message) + if "FINISHED" not in result: + raise RuntimeError(f"Unable to create the {message} Undo step") + + +def da3_parameters(settings, path, model, process_res, start_frame): + return { + "input": str(path), + "model": model, + "process_res": process_res, + "process_res_method": "upper_bound_resize", + "video_frames": settings.video_frames, + "start_frame": start_frame, + "device": production_device().lower(), + "ref_view_strategy": settings.ref_view_strategy, + } + + +def moge2_parameters(settings, path, model, resolution_level, start_frame): + return { + "input": str(path), + "model": model, + "resolution_level": int(resolution_level), + "video_frames": settings.video_frames, + "start_frame": start_frame, + "device": production_device().lower(), + } diff --git a/src/snapform/operators/debug.py b/src/snapform/operators/debug.py new file mode 100644 index 0000000..74537e1 --- /dev/null +++ b/src/snapform/operators/debug.py @@ -0,0 +1,91 @@ +import bpy + +from ..runtime import JobOperatorBase +from ..properties import BEN2_MODEL_KEY +from .common import ( + background_input_path, + da3_parameters, + input_path, + moge2_parameters, + open_job_directory, + require_environment, + require_model, +) + + +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.model + require_model(model) + return { + **da3_parameters( + settings, + input_path(settings.input_path), + model, + settings.process_res, + 1, + ), + "device": settings.device.lower(), + "inference_mode": "joint", + "remove_background_model": "", + } + + def response(self, _context, result): + open_job_directory(result.directory) + return f"Saved DA3 API prediction to {result.directory}" + + +class DebugMoge2(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.debug_moge2" + bl_label = "Debug MoGe-2" + job_type = "debug-moge2" + starting_message = "Starting MoGe-2 inference" + + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + require_model(settings.moge_model) + return { + **moge2_parameters( + settings, + input_path(settings.input_path), + settings.moge_model, + settings.moge_resolution_level, + 1, + ), + "device": settings.device.lower(), + } + + def response(self, _context, result): + open_job_directory(result.directory) + return f"Exported MoGe-2 data to {result.directory}" + + +class DebugBen2(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.debug_ben2" + bl_label = "Debug BEN2" + job_type = "debug-ben2" + starting_message = "Starting BEN2 background removal" + + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + require_model(BEN2_MODEL_KEY) + return { + "input": str(background_input_path(settings.background_input_path)), + "model": BEN2_MODEL_KEY, + "device": settings.device.lower(), + "video_frames": settings.video_frames, + "start_frame": 1, + } + + def response(self, _context, result): + open_job_directory(result.directory) + return f"Removed background to {result.directory}" diff --git a/src/snapform/operators/depth_plane/__init__.py b/src/snapform/operators/depth_plane/__init__.py new file mode 100644 index 0000000..8aea1ed --- /dev/null +++ b/src/snapform/operators/depth_plane/__init__.py @@ -0,0 +1,7 @@ +from .operator import ( + ConvertToDepthPlane, + ExtractOriginalFrames, + GenerateDa3Depth, + GenerateMoge2Depth, + GenerateOriginalBackground, +) diff --git a/src/snapform/operators/depth_plane/build.py b/src/snapform/operators/depth_plane/build.py new file mode 100644 index 0000000..77a5de5 --- /dev/null +++ b/src/snapform/operators/depth_plane/build.py @@ -0,0 +1,546 @@ +import bpy + +from ..image import ( + copy_depth_image, + create_sequence_image, + depth_data_name, + image_user_settings, + is_animated_image, + plane_bounds, + sorted_frame_paths, +) +from ..material import ( + DEPTH_PLANE_NODE_GROUP_NAME, + DISPLACEMENT_MIDLEVEL, + DISPLACEMENT_STRENGTH, + MOGE_DEPTH_PLANE_NODE_GROUP_NAME, + ORIGINAL_DISPLACEMENT_MIDLEVEL, + ORIGINAL_DISPLACEMENT_STRENGTH, + create_material, + create_moge_material, + drive_modifier_input, + load_depth_plane_node_group, + modifier_input_identifier, + set_modifier_input, +) +from .data import ( + camera_plane_bounds, + create_depth_image, + create_depth_sequence_image, + depth_displacement_values, + extract_depth_arrays, + extract_intrinsics_array, + prepare_depth_values, + scale_depth_values, +) +def plane_fit_transform(source_object, camera_bounds): + target_bounds = plane_bounds(source_object) + camera_width = camera_bounds[1] - camera_bounds[0] + camera_height = camera_bounds[3] - camera_bounds[2] + if camera_width <= 0.0 or camera_height <= 0.0: + raise ValueError("The camera plane has invalid dimensions") + + target_width = target_bounds[1] - target_bounds[0] + target_height = target_bounds[3] - target_bounds[2] + scale_x = target_width / camera_width + scale_y = target_height / camera_height + scale_z = (scale_x * scale_y) ** 0.5 + camera_center_x = (camera_bounds[0] + camera_bounds[1]) * 0.5 + camera_center_y = (camera_bounds[2] + camera_bounds[3]) * 0.5 + target_center_x = (target_bounds[0] + target_bounds[1]) * 0.5 + target_center_y = (target_bounds[2] + target_bounds[3]) * 0.5 + return ( + scale_x, + scale_y, + scale_z, + target_center_x - camera_center_x * scale_x, + target_center_y - camera_center_y * scale_y, + ) + + +def apply_plane_transform(plane_object, source_object, camera_bounds): + from mathutils import Matrix + + scale_x, scale_y, scale_z, offset_x, offset_y = plane_fit_transform( + source_object, + camera_bounds, + ) + fit_matrix = Matrix( + ( + (scale_x, 0.0, 0.0, offset_x), + (0.0, scale_y, 0.0, offset_y), + (0.0, 0.0, 1.0, 0.0), + (0.0, 0.0, 0.0, 1.0), + ) + ) + plane_object.matrix_world = source_object.matrix_world @ fit_matrix + return abs(float(plane_object.scale.z)) * abs(scale_z) + + +def apply_object_scale_to_mesh(plane_object): + """Bake object scale into the plane while keeping location and rotation.""" + from mathutils import Matrix + + scale = plane_object.scale.copy() + plane_object.data.transform( + Matrix.Diagonal((scale.x, scale.y, scale.z, 1.0)) + ) + plane_object.data.update() + plane_object.scale = (1.0, 1.0, 1.0) + + + +def grid_segments(bounds, subdivision_level): + if subdivision_level < 0: + raise ValueError("Subdivision must not be negative") + + width = bounds[1] - bounds[0] + height = bounds[3] - bounds[2] + if width <= 0.0 or height <= 0.0: + raise ValueError("The plane has invalid dimensions") + + short_segments = 2**subdivision_level + if width >= height: + return ( + max(short_segments, int(short_segments * width / height + 0.5)), + short_segments, + ) + return ( + short_segments, + max(short_segments, int(short_segments * height / width + 0.5)), + ) + + +def create_plane_mesh(source_object, bounds=None, subdivision_level=5): + x_min, x_max, y_min, y_max = ( + bounds if bounds is not None else plane_bounds(source_object) + ) + if subdivision_level is None: + x_segments, y_segments = 1, 1 + else: + x_segments, y_segments = grid_segments( + (x_min, x_max, y_min, y_max), + subdivision_level, + ) + vertices = [] + uv_coordinates = [] + for y_index in range(y_segments + 1): + v = y_index / y_segments + y = y_min + (y_max - y_min) * v + for x_index in range(x_segments + 1): + u = x_index / x_segments + x = x_min + (x_max - x_min) * u + vertices.append((x, y, 0.0)) + uv_coordinates.append((u, v)) + + faces = [] + row_width = x_segments + 1 + for y_index in range(y_segments): + for x_index in range(x_segments): + bottom_left = y_index * row_width + x_index + faces.append( + ( + bottom_left, + bottom_left + 1, + bottom_left + row_width + 1, + bottom_left + row_width, + ) + ) + + mesh = bpy.data.meshes.new(source_object.name) + mesh.from_pydata(vertices, (), faces) + mesh.update() + for polygon in mesh.polygons: + polygon.use_smooth = True + uv_layer = mesh.uv_layers.new(name="UVMap") + for loop in mesh.loops: + uv_layer.data[loop.index].uv = uv_coordinates[ + loop.vertex_index + ] + return mesh + + + +def add_depth_plane_processing( + plane_object, + depth_image, + color_image, + material, + frame_start=1, + frame_duration=1, + use_cyclic=False, + animate_frame=False, +): + node_group = load_depth_plane_node_group() + modifier = plane_object.modifiers.new( + name=DEPTH_PLANE_NODE_GROUP_NAME, + type="NODES", + ) + modifier.node_group = node_group + modifier.show_group_selector = False + for name, value in ( + ("Depth Image", depth_image), + ("Color Image", color_image), + ("Material", material), + ("Frame", 1), + ("Frame Start", frame_start), + ("Frames", max(int(frame_duration or 1), 1)), + ("Cyclic", use_cyclic), + ): + identifier = modifier_input_identifier(node_group, name) + set_modifier_input(modifier, identifier, value) + if animate_frame: + frame_identifier = modifier_input_identifier(node_group, "Frame") + drive_modifier_input(plane_object, modifier, frame_identifier) + plane_object.update_tag(refresh={"DATA"}) + return modifier + + +def add_moge_depth_plane_processing( + plane_object, + depth_image, + color_image, + material, + frame_start=1, + frame_duration=1, + use_cyclic=False, + animate_frame=False, +): + node_group = load_depth_plane_node_group( + MOGE_DEPTH_PLANE_NODE_GROUP_NAME + ) + modifier = plane_object.modifiers.new( + name=MOGE_DEPTH_PLANE_NODE_GROUP_NAME, + type="NODES", + ) + modifier.node_group = node_group + modifier.show_group_selector = False + for name, value in ( + ("Depth Image", depth_image), + ("Color Image", color_image), + ("Material", material), + ("Frame", 1), + ("Frame Start", frame_start), + ("Frames", max(int(frame_duration or 1), 1)), + ("Cyclic", use_cyclic), + ): + identifier = modifier_input_identifier(node_group, name) + set_modifier_input(modifier, identifier, value) + if animate_frame: + frame_identifier = modifier_input_identifier(node_group, "Frame") + drive_modifier_input(plane_object, modifier, frame_identifier) + plane_object.update_tag(refresh={"DATA"}) + return modifier + + +def finish_depth_plane( + context, + source_object, + plane_object, + delete_source, +): + source_name = source_object.name + for selected in context.selected_objects: + selected.select_set(False) + plane_object.select_set(True) + context.view_layer.objects.active = plane_object + if delete_source: + bpy.data.objects.remove(source_object, do_unlink=True) + plane_object.name = source_name + return plane_object + + +def build_image_depth_plane( + context, + source_object, + delete_source, + shadeless=False, + color_image=None, + depth_frames_directory=None, + depth_image=None, +): + source_name = source_object.name + frame_count = 1 + frames_image = None + if depth_frames_directory is not None: + frames_image = create_sequence_image( + depth_data_name(source_object), + depth_frames_directory, + frame_prefix="color.", + ) + frame_count = len( + sorted_frame_paths(depth_frames_directory, "color.") + ) + ( + source_frame_start, + source_frame_offset, + source_frame_duration, + source_use_cyclic, + ) = image_user_settings(source_object) + uses_extracted = frames_image is not None + texture_frame_count = ( + frame_count if uses_extracted else source_frame_duration + ) + texture_frame_offset = 0 if uses_extracted else source_frame_offset + depth_source_image = depth_image or source_object.data + if depth_image is None: + if frames_image is not None: + depth_source_image = frames_image + elif not is_animated_image(source_object.data): + depth_source_image = copy_depth_image(source_object) + material_image = color_image + if material_image is None: + raise RuntimeError("Depth Plane material image is required") + mesh = create_plane_mesh(source_object, subdivision_level=None) + plane_object = bpy.data.objects.new(source_name, mesh) + context.collection.objects.link(plane_object) + plane_object.matrix_world = source_object.matrix_world.copy() + + depth_scale = abs(float(plane_object.scale.z)) + material = create_material( + source_name, + material_image, + depth_scale, + displacement_image=depth_source_image, + displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL, + displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH, + shadeless=shadeless, + frame_count=frame_count, + frame_start=source_frame_start, + color_frame_offset=texture_frame_offset, + color_frame_duration=texture_frame_count, + depth_frame_offset=texture_frame_offset, + depth_frame_duration=texture_frame_count, + use_cyclic=source_use_cyclic, + scene=context.scene, + ) + mesh.materials.append(material) + add_depth_plane_processing( + plane_object, + depth_source_image, + material_image, + material, + frame_start=source_frame_start, + frame_duration=frame_count, + use_cyclic=source_use_cyclic, + animate_frame=is_animated_image(source_object.data), + ) + plane_object["snapform_frame_count"] = frame_count + return finish_depth_plane( + context, + source_object, + plane_object, + delete_source, + ) + + +def build_depth_plane( + context, + source_object, + npz_path, + delete_source, + shadeless=False, + depth_name="", + color_image=None, + depth_frames_directory=None, +): + source_name = source_object.name + raw_depth_arrays = extract_depth_arrays(npz_path) + raw_depth = raw_depth_arrays[0] + intrinsics = extract_intrinsics_array(npz_path) + raw_depth_values = prepare_depth_values(raw_depth) + bounds, reference_depth = camera_plane_bounds( + source_object, + raw_depth_values, + intrinsics, + DISPLACEMENT_STRENGTH, + ) + depth_arrays = [ + depth_displacement_values(frame) for frame in raw_depth_arrays + ] + depth_name = depth_name or depth_data_name(source_object) + frame_count = len(depth_arrays) + material_image = color_image + if material_image is None: + raise RuntimeError("Depth Plane material image is required") + ( + source_frame_start, + source_frame_offset, + source_frame_duration, + source_use_cyclic, + ) = image_user_settings(source_object) + mesh = create_plane_mesh(source_object, bounds, None) + plane_object = bpy.data.objects.new( + source_name, + mesh, + ) + context.collection.objects.link(plane_object) + depth_scale = apply_plane_transform( + plane_object, + source_object, + bounds, + ) + apply_object_scale_to_mesh(plane_object) + scaled_depth_arrays = [ + scale_depth_values(frame, depth_scale) for frame in depth_arrays + ] + depth_image = create_depth_image( + depth_name, + scaled_depth_arrays[0], + ) + if ( + is_animated_image(source_object.data) + and depth_frames_directory is not None + ): + bpy.data.images.remove(depth_image, do_unlink=True) + depth_image = create_depth_sequence_image( + depth_name, + scaled_depth_arrays, + context.scene, + depth_frames_directory, + ) + + material = create_material( + source_name, + material_image, + depth_scale, + displacement_image=depth_image, + displacement_midlevel=DISPLACEMENT_MIDLEVEL, + displacement_scale=DISPLACEMENT_STRENGTH, + shadeless=shadeless, + frame_count=frame_count, + frame_start=source_frame_start, + color_frame_offset=0 if color_image is not None else source_frame_offset, + color_frame_duration=( + frame_count if color_image is not None else source_frame_duration + ), + depth_frame_offset=0, + depth_frame_duration=frame_count, + use_cyclic=source_use_cyclic, + scene=context.scene, + ) + mesh.materials.append(material) + add_depth_plane_processing( + plane_object, + depth_image, + material_image, + material, + frame_start=source_frame_start, + frame_duration=frame_count, + use_cyclic=source_use_cyclic, + animate_frame=is_animated_image(source_object.data), + ) + plane_object["snapform_reference_depth"] = reference_depth + plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0] + plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2] + plane_object["snapform_depth_scale"] = depth_scale + plane_object["snapform_frame_count"] = frame_count + + return finish_depth_plane( + context, + source_object, + plane_object, + delete_source, + ) + + +def build_moge2_depth_plane( + context, + source_object, + npz_path, + delete_source, + shadeless=False, + color_image=None, + data_frames_directory=None, +): + import numpy as np + + try: + with np.load(npz_path, allow_pickle=False) as archive: + depth = np.asarray(archive["depth"], dtype=np.float32) + mask = np.asarray(archive["mask"], dtype=bool) + intrinsics = np.asarray(archive["intrinsics"], dtype=np.float32) + except (KeyError, OSError) as error: + raise ValueError(f"Invalid MoGe-2 NPZ: {npz_path}") from error + if depth.ndim == 2: + depth = depth[None, ...] + mask = mask[None, ...] + if intrinsics.ndim == 2: + intrinsics = intrinsics[None, ...] + if not (len(depth) == len(mask) == len(intrinsics)): + raise ValueError("MoGe-2 arrays have different frame counts") + + first_valid = mask[0] & np.isfinite(depth[0]) & (depth[0] > 0.0) + if not first_valid.any(): + raise ValueError("MoGe-2 mask contains no valid depth") + reference_depth = float(np.median(depth[0][first_valid])) + height, width = depth[0].shape + pixel_intrinsics = intrinsics[0].copy() + pixel_intrinsics[0] *= width + pixel_intrinsics[1] *= height + bounds, _ = camera_plane_bounds( + source_object, + depth[0], + pixel_intrinsics, + 1.0, + reference_depth=reference_depth, + ) + displacement = [reference_depth - frame for frame in depth] + source_frame_start, _, _, source_use_cyclic = image_user_settings( + source_object + ) + mesh = create_plane_mesh(source_object, bounds, None) + plane_object = bpy.data.objects.new(source_object.name, mesh) + context.collection.objects.link(plane_object) + depth_scale = apply_plane_transform(plane_object, source_object, bounds) + apply_object_scale_to_mesh(plane_object) + depth_image = create_depth_image( + depth_data_name(source_object), + displacement[0] * depth_scale, + ) + if ( + is_animated_image(source_object.data) + and data_frames_directory is not None + ): + bpy.data.images.remove(depth_image, do_unlink=True) + depth_image = create_depth_sequence_image( + depth_data_name(source_object), + [frame * depth_scale for frame in displacement], + context.scene, + data_frames_directory, + "moge_depth", + ) + material_image = color_image + if material_image is None: + raise RuntimeError("Depth Plane material image is required") + material = create_moge_material( + source_object.name, + material_image, + depth_image, + depth_scale, + shadeless=shadeless, + frame_count=len(depth), + frame_start=source_frame_start, + use_cyclic=source_use_cyclic, + scene=context.scene, + ) + mesh.materials.append(material) + modifier = add_moge_depth_plane_processing( + plane_object, + depth_image, + material_image, + material, + frame_start=source_frame_start, + frame_duration=len(depth), + use_cyclic=source_use_cyclic, + animate_frame=is_animated_image(source_object.data), + ) + plane_object["snapform_reference_depth"] = reference_depth + plane_object["snapform_depth_scale"] = depth_scale + plane_object["snapform_frame_count"] = len(depth) + return finish_depth_plane( + context, + source_object, + plane_object, + delete_source, + ) diff --git a/src/snapform/operators/depth_plane/data.py b/src/snapform/operators/depth_plane/data.py new file mode 100644 index 0000000..a3c18df --- /dev/null +++ b/src/snapform/operators/depth_plane/data.py @@ -0,0 +1,262 @@ +from pathlib import Path + +import bpy + +from ..image import force_sequence_source, pack_sequence_if_possible +DEPTH_KEY_CANDIDATES = ("depth", "depths") +INTRINSICS_KEY_CANDIDATES = ("intrinsics", "ixts") + +def create_depth_sequence_image( + name, + depth_frames, + scene, + frame_directory, + frame_prefix="depth", +): + """Write one EXR per depth frame and return a SEQUENCE image datablock.""" + import numpy as np + + frame_directory = Path(frame_directory) + frame_directory.mkdir(parents=True, exist_ok=True) + height, width = depth_frames[0].shape + image = bpy.data.images.new( + name, + width=width, + height=height, + alpha=False, + float_buffer=True, + ) + image.colorspace_settings.name = "Non-Color" + image_settings = scene.render.image_settings + original_file_format = image_settings.file_format + frame_paths = [] + try: + image_settings.file_format = "OPEN_EXR" + for frame_index, depth_values in enumerate(depth_frames): + # NPZ arrays use a top-left origin, while Blender image pixels start below. + depth_pixels = np.flipud(depth_values) + rgba = np.empty((height, width, 4), dtype=np.float32) + rgba[..., 0] = depth_pixels + rgba[..., 1] = depth_pixels + rgba[..., 2] = depth_pixels + rgba[..., 3] = 1.0 + image.pixels.foreach_set(rgba.ravel()) + frame_path = ( + frame_directory / f"{frame_prefix}.{frame_index + 1:04d}.exr" + ) + image.save_render(str(frame_path), scene=scene) + if not frame_path.is_file(): + raise RuntimeError(f"Blender did not write: {frame_path}") + frame_paths.append(frame_path) + finally: + image_settings.file_format = original_file_format + bpy.data.images.remove(image, do_unlink=True) + if not frame_paths: + raise RuntimeError("Unable to write the depth frame sequence") + try: + sequence_image = bpy.data.images.load( + str(frame_paths[0]), + check_existing=False, + ) + force_sequence_source(sequence_image) + except (OSError, RuntimeError) as error: + if "sequence_image" in locals(): + bpy.data.images.remove(sequence_image, do_unlink=True) + raise RuntimeError( + "Unable to load the depth frame sequence" + ) from error + sequence_image.name = name + sequence_image.colorspace_settings.name = "Non-Color" + pack_sequence_if_possible(sequence_image, len(frame_paths)) + return sequence_image + + + +def scale_depth_values(depth, depth_scale): + return prepare_depth_values(depth) * depth_scale + + +def extract_depth_arrays(npz_path): + try: + import numpy as np + except ImportError as error: + raise RuntimeError("Blender Python does not provide NumPy") from error + + try: + with np.load(npz_path, allow_pickle=False) as archive: + depth_key = next( + ( + key + for key in DEPTH_KEY_CANDIDATES + if key in archive.files + ), + None, + ) + if depth_key is None: + keys = ", ".join(archive.files) or "none" + raise ValueError( + f"Mini NPZ has no depth array; available keys: {keys}" + ) + depth = np.asarray(archive[depth_key], dtype=np.float32) + except OSError as error: + raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error + + if depth.ndim >= 3 and depth.shape[-1] == 1: + depth = depth[..., 0] + if depth.ndim == 2: + depth = depth[None, ...] + if depth.ndim != 3: + raise ValueError( + f"Unsupported depth array shape: {tuple(depth.shape)}" + ) + return [depth[index] for index in range(depth.shape[0])] + + +def extract_depth_array(npz_path): + return extract_depth_arrays(npz_path)[0] + + +def extract_intrinsics_array(npz_path): + try: + import numpy as np + except ImportError as error: + raise RuntimeError("Blender Python does not provide NumPy") from error + + try: + with np.load(npz_path, allow_pickle=False) as archive: + intrinsics_key = next( + ( + key + for key in INTRINSICS_KEY_CANDIDATES + if key in archive.files + ), + None, + ) + if intrinsics_key is None: + keys = ", ".join(archive.files) or "none" + raise ValueError( + "Mini NPZ has no camera intrinsics; " + f"available keys: {keys}" + ) + intrinsics = np.asarray( + archive[intrinsics_key], + dtype=np.float32, + ) + except OSError as error: + raise ValueError(f"Unable to read Mini NPZ: {npz_path}") from error + + if intrinsics.ndim == 3: + intrinsics = intrinsics[0] + if intrinsics.shape != (3, 3): + raise ValueError( + f"Unsupported intrinsics shape: {tuple(intrinsics.shape)}" + ) + return intrinsics + + +def prepare_depth_values(depth): + import numpy as np + + valid = np.isfinite(depth) + if not valid.any(): + raise ValueError("Depth array contains no finite values") + values = depth.astype(np.float32, copy=True) + values[~valid] = 0.0 + return values + + +def depth_displacement_values(depth): + """Convert raw DA3 camera distance to positive outward displacement.""" + import numpy as np + + return 1.0 - np.asarray(depth, dtype=np.float32) + + +def camera_plane_bounds( + source_object, + depth, + intrinsics, + displacement_strength, + reference_depth=None, +): + import numpy as np + + scale = abs(float(displacement_strength)) + if scale <= 1e-8: + raise ValueError("Displacement Strength must not be zero") + + if reference_depth is None: + valid = np.isfinite(depth) & (depth > 0.0) + if not valid.any(): + raise ValueError("Depth array contains no positive finite values") + reference_depth = float(np.median(depth[valid])) + reference_distance = reference_depth * scale + + height, width = depth.shape + fx = float(intrinsics[0, 0]) + fy = float(intrinsics[1, 1]) + cx = float(intrinsics[0, 2]) + cy = float(intrinsics[1, 2]) + if fx <= 0.0 or fy <= 0.0: + raise ValueError("Camera intrinsics contain an invalid focal length") + + source_width, source_height = source_object.data.size + if source_width <= 0 or source_height <= 0: + raise ValueError("The source image has no readable dimensions") + source_aspect = source_width / source_height + + # Predicted fx/fy may differ even for square pixels. Preserve their projected + # area while enforcing the source image aspect ratio to avoid image stretch. + raw_width = width * reference_distance / fx + raw_height = height * reference_distance / fy + plane_area = raw_width * raw_height + plane_width = float(np.sqrt(plane_area * source_aspect)) + plane_height = float(np.sqrt(plane_area / source_aspect)) + + x_min = -(cx / width) * plane_width + x_max = (1.0 - cx / width) * plane_width + y_min = -(1.0 - cy / height) * plane_height + y_max = (cy / height) * plane_height + + # Preserve an intentional Image Empty origin offset around the camera frame. + offset = source_object.empty_image_offset + x_offset = (float(offset[0]) + 0.5) * plane_width + y_offset = (float(offset[1]) + 0.5) * plane_height + return ( + ( + x_min + x_offset, + x_max + x_offset, + y_min + y_offset, + y_max + y_offset, + ), + reference_depth, + ) + + +def create_depth_image(name, depth_values): + import numpy as np + + height, width = depth_values.shape + image = bpy.data.images.new( + name, + width=width, + height=height, + alpha=False, + float_buffer=True, + ) + image.colorspace_settings.name = "Non-Color" + image.file_format = "OPEN_EXR" + + # NPZ arrays use a top-left origin, while Blender image pixels start below. + depth_pixels = np.flipud(depth_values) + rgba = np.empty((height, width, 4), dtype=np.float32) + rgba[..., 0] = depth_pixels + rgba[..., 1] = depth_pixels + rgba[..., 2] = depth_pixels + rgba[..., 3] = 1.0 + image.pixels.foreach_set(rgba.ravel()) + # Pack as EXR so values outside 0..1 survive saving and reopening the blend. + image.pack() + return image + + diff --git a/src/snapform/operators/depth_plane/operator.py b/src/snapform/operators/depth_plane/operator.py new file mode 100644 index 0000000..b7b6d63 --- /dev/null +++ b/src/snapform/operators/depth_plane/operator.py @@ -0,0 +1,481 @@ +import bpy + +from ...properties import BEN2_MODEL_KEY, DEFAULT_MOGE2_MODEL_KEY +from ...runtime import JobOperatorBase +from ..common import ( + HIDDEN, + background_input_path, + cleanup_temporary_input, + da3_parameters, + input_path, + moge2_parameters, + push_undo, + require_environment, + require_model, + source_start_frame, +) +from .build import build_depth_plane, build_image_depth_plane, build_moge2_depth_plane +from ..image import ( + default_extract_frames, + depth_data_name, + draw_process_resolution, + is_image_empty, + prepare_image_input, + process_resolution_scale_property, + remove_temporary_input, + scaled_process_resolution, + video_total_frames, +) + + +def image_depth_model_items(owner, context): + from ...properties import depth_plane_model_enum_items + + return depth_plane_model_enum_items(owner, context) + + + +class ConvertToDepthPlane(bpy.types.Operator): + bl_idname = "snapform.convert_to_depth_plane" + bl_label = "Convert to Depth Plane" + bl_description = "Convert an image to a displaced depth plane" + bl_options = {"REGISTER"} + + shadeless: bpy.props.BoolProperty( + name="Shadeless", + description="Display the source color without scene lighting", + default=False, + ) + remove_background: bpy.props.BoolProperty( + name="Remove Background", + description="Remove the background during conversion with BEN2", + default=False, + ) + model: bpy.props.EnumProperty( + name="Depth From", + description="Model or original image used to build depth", + items=image_depth_model_items, + default=0, + ) + process_res_factor: process_resolution_scale_property() + resolution_level: bpy.props.IntProperty( + name="Resolution Level", + description="MoGe-2 inference resolution level; higher is finer and slower", + default=5, + min=0, + max=9, + ) + keep_original: bpy.props.BoolProperty( + name="Keep Original", + description="Keep the source Image Empty after conversion", + default=False, + ) + + @classmethod + def poll(cls, context): + from ...runtime import runtime + from ...properties import required_models_ready + + return ( + is_image_empty(context.object) + and runtime.environment_ready() + and not runtime.server_busy() + and required_models_ready() + ) + + def invoke(self, context, _event): + return context.window_manager.invoke_props_dialog(self, width=340) + + def draw(self, context): + layout = self.layout + layout.prop(self, "model") + if self.model in {"SMALL", "BASE"}: + draw_process_resolution(layout, self, context.object.data) + elif self.model.startswith("MOGE2_"): + layout.prop(self, "resolution_level") + layout.prop(self, "remove_background") + layout.prop(self, "shadeless") + layout.separator(type="LINE") + layout.prop(self, "keep_original") + + def execute(self, context): + source_object = context.object + if video_total_frames(source_object.data) > 1: + context.scene.snapform_settings.video_frames = ( + default_extract_frames(source_object) + ) + if self.model == "ORIGINAL": + try: + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + run_original = ( + bpy.ops.snapform.generate_original_background + if self.remove_background + else bpy.ops.snapform.extract_original_frames + ) + result = run_original( + input_path=str(source_path), + source_object_name=source_object.name, + delete_source=not self.keep_original, + delete_input=delete_input, + shadeless=self.shadeless, + ) + if "CANCELLED" not in result: + return {"FINISHED"} + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + + if self.model.startswith("MOGE2_"): + try: + source_path, delete_input = prepare_image_input( + source_object.data, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + result = bpy.ops.snapform.generate_moge2_depth( + input_path=str(source_path), + source_object_name=source_object.name, + resolution_level=self.resolution_level, + model=self.model, + delete_source=not self.keep_original, + delete_input=delete_input, + shadeless=self.shadeless, + remove_background=self.remove_background, + ) + if "CANCELLED" not in result: + return {"FINISHED"} + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + + image = source_object.data + try: + process_res = scaled_process_resolution( + image, + self.process_res_factor, + ) + except ValueError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + try: + source_path, delete_input = prepare_image_input( + image, + context.scene, + ) + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + + result = bpy.ops.snapform.generate_da3_depth( + input_path=str(source_path), + source_object_name=source_object.name, + depth_data_name=depth_data_name(source_object), + model=self.model, + process_res=process_res, + delete_source=not self.keep_original, + delete_input=delete_input, + shadeless=self.shadeless, + remove_background=self.remove_background, + ) + if "CANCELLED" in result: + if delete_input: + remove_temporary_input(source_path) + return {"CANCELLED"} + return {"FINISHED"} + + + + + + +def build_da3_depth_response( + context, + result, + *, + source_object_name, + depth_data_name, + remove_background, + delete_source, + shadeless, +): + from .build import build_depth_plane + from ..image import ( + create_color_image, + is_animated_image, + is_image_empty, + ) + + directory = result.directory + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + animated = is_animated_image(source_object.data) + color_image = create_color_image( + source_object, + directory, + assign_to_source=False, + ) + build_depth_plane( + context, + source_object, + result.file("depth"), + delete_source, + shadeless=shadeless, + depth_name=depth_data_name, + color_image=color_image, + depth_frames_directory=directory if animated else None, + ) + push_undo("Convert to Depth Plane") + + +def build_moge2_depth_response( + context, + result, + *, + source_object_name, + remove_background, + delete_source, + shadeless, +): + from .build import build_moge2_depth_plane + from ..image import create_color_image, is_image_empty + + directory = result.directory + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + color_image = create_color_image( + source_object, + directory, + assign_to_source=False, + ) + build_moge2_depth_plane( + context, + source_object, + result.file("depth"), + delete_source, + shadeless=shadeless, + color_image=color_image, + data_frames_directory=directory, + ) + push_undo("Convert to Depth Plane") + + +def build_original_depth_response( + context, + result, + *, + source_object_name, + background, + delete_source, + shadeless, +): + from .build import build_image_depth_plane + from ..image import ( + create_color_image, + is_animated_image, + is_image_empty, + ) + + directory = result.directory + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + animated = is_animated_image(source_object.data) + color_image = create_color_image( + source_object, + directory, + assign_to_source=False, + ) + build_image_depth_plane( + context, + source_object, + delete_source, + shadeless=shadeless, + color_image=color_image, + depth_frames_directory=directory if animated else None, + ) + push_undo("Convert to Depth Plane") + +class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.generate_da3_depth" + bl_label = "Generate DA3 Depth" + job_type = "generate-depth" + starting_message = "Starting depth estimation" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + depth_data_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default="BASE", options=HIDDEN) + process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **da3_parameters( + settings, + input_path(self.input_path), + self.model, + self.process_res, + source_start_frame(self.source_object_name), + ), + "inference_mode": "per_frame", + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + } + + def response(self, context, result): + build_da3_depth_response( + context, + result, + source_object_name=self.source_object_name, + depth_data_name=self.depth_data_name, + remove_background=self.remove_background, + delete_source=self.delete_source, + shadeless=self.shadeless, + ) + return "Depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.generate_moge2_depth" + bl_label = "Generate MoGe-2 Depth" + job_type = "generate-depth" + starting_message = "Starting MoGe-2 inference" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) + resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **moge2_parameters( + settings, + input_path(self.input_path), + self.model, + self.resolution_level, + source_start_frame(self.source_object_name), + ), + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + } + + def response(self, context, result): + build_moge2_depth_response( + context, + result, + source_object_name=self.source_object_name, + remove_background=self.remove_background, + delete_source=self.delete_source, + shadeless=self.shadeless, + ) + return "MoGe-2 depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class GenerateOriginalBackground(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.generate_original_background" + bl_label = "Generate Original Background" + job_type = "remove-background" + starting_message = "Starting background removal" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + "input": str(background_input_path(self.input_path)), + "model": BEN2_MODEL_KEY, + "device": production_device().lower(), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + } + + def response(self, context, result): + build_original_depth_response( + context, + result, + source_object_name=self.source_object_name, + background=True, + delete_source=self.delete_source, + shadeless=self.shadeless, + ) + return "Depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class ExtractOriginalFrames(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.extract_original_frames" + bl_label = "Extract Original Frames" + job_type = "extract-frames" + starting_message = "Extracting source frames" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + return { + "input": str(input_path(self.input_path)), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + } + + def response(self, context, result): + build_original_depth_response( + context, + result, + source_object_name=self.source_object_name, + background=False, + delete_source=self.delete_source, + shadeless=self.shadeless, + ) + return "Depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) diff --git a/src/snapform/operators/image.py b/src/snapform/operators/image.py new file mode 100644 index 0000000..1a058a7 --- /dev/null +++ b/src/snapform/operators/image.py @@ -0,0 +1,451 @@ +import re +import shutil +import tempfile +from pathlib import Path + +import bpy +DA3_PATCH_SIZE = 14 +DA3_MIN_PROCESS_RES = 252 +DA3_MAX_PROCESS_RES = 1008 +IMAGE_NAME_SUFFIXES = ( + ".jpeg", + ".jpg", + ".png", + ".webp", + ".tiff", + ".tif", + ".exr", + ".hdr", + ".bmp", + ".tga", +) + + +VIDEO_NAME_SUFFIXES = ( + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", +) + + +def is_image_empty(candidate): + return bool( + candidate + and getattr(candidate, "type", None) == "EMPTY" + and getattr(candidate, "empty_display_type", None) == "IMAGE" + and getattr(candidate, "data", None) is not None + ) + + + +def is_animated_image(image): + return getattr(image, "source", None) in {"MOVIE", "SEQUENCE"} + + +def video_total_frames(image): + """Return the total animated frame count, or 1 for a still image.""" + source = getattr(image, "source", None) + if source == "MOVIE": + return int(getattr(image, "frame_duration", 1) or 1) + if source == "SEQUENCE": + return len(collect_sequence_frames(image)) or 1 + return 1 + + +def image_user_settings(source_object): + """Return the Empty's image user settings as a four-item tuple.""" + image_user = getattr(source_object, "image_user", None) + if image_user is None: + return 1, 0, 0, False + return ( + int(getattr(image_user, "frame_start", 1) or 1), + int(getattr(image_user, "frame_offset", 0) or 0), + int(getattr(image_user, "frame_duration", 0) or 0), + bool(getattr(image_user, "use_cyclic", False)), + ) + + +def extract_frame_range(source_object): + """Return the clamped (start_frame, frame_count) range extracted from the source.""" + total_frames = video_total_frames(getattr(source_object, "data", None)) + if total_frames <= 1: + return 1, total_frames + _frame_start, frame_offset, image_frames, _use_cyclic = ( + image_user_settings(source_object) + ) + frame_count = ( + min(image_frames, total_frames) + if image_frames >= 1 + else total_frames + ) + start_frame = min(1 + frame_offset, total_frames - frame_count + 1) + start_frame = max(start_frame, 1) + return start_frame, frame_count + + +def default_extract_frames(source_object): + """Return the Empty's Image Frames value, or the full frame count when unset.""" + return extract_frame_range(source_object)[1] + + +def image_path(image): + filepath_from_user = getattr(image, "filepath_from_user", None) + filepath = ( + filepath_from_user() + if callable(filepath_from_user) + else getattr(image, "filepath", "") + ) + if not filepath: + return Path() + return Path(bpy.path.abspath(filepath)).expanduser() + + +def image_is_packed(image): + if getattr(image, "packed_file", None) is not None: + return True + return bool(getattr(image, "packed_files", ())) + + +def remove_temporary_input(path): + """Remove a temporary input file or frame directory.""" + path = Path(path) + try: + if path.is_dir(): + shutil.rmtree(path, ignore_errors=True) + else: + path.unlink(missing_ok=True) + except OSError: + pass + + +def image_base_name(image): + name = image.name + lower_name = name.lower() + for suffix in IMAGE_NAME_SUFFIXES + VIDEO_NAME_SUFFIXES: + suffix_index = lower_name.rfind(suffix) + if suffix_index <= 0: + continue + trailing = name[suffix_index + len(suffix) :] + if trailing and not ( + trailing.startswith(".") and trailing[1:].isdigit() + ): + continue + return f"{name[:suffix_index]}{trailing}" + return name + + +def depth_data_name(source_object): + return f"{image_base_name(source_object.data)}.depth" + + +def normal_data_name(source_object): + return f"{image_base_name(source_object.data)}.normal" + + +def copy_depth_image(source_object): + """Return a renamed copy of the source image for use as the depth map.""" + depth_image = source_object.data.copy() + depth_image.name = depth_data_name(source_object) + return depth_image + + +def color_data_name(source_object): + return f"{image_base_name(source_object.data)}.color" + + +def scaled_process_resolution(image, scale_percent): + """Return the scaled long edge clamped and aligned for DA3 inference.""" + width, height = image.size + if width <= 0 or height <= 0: + raise ValueError("The source image has no readable dimensions") + if scale_percent < 1.0 or scale_percent > 100.0: + raise ValueError("Resolution Scale must be between 1% and 100%") + process_res = max(width, height) * scale_percent / 100.0 + process_res = max(DA3_MIN_PROCESS_RES, process_res) + process_res = min(DA3_MAX_PROCESS_RES, process_res) + return int(process_res / DA3_PATCH_SIZE + 0.5) * DA3_PATCH_SIZE + + +def process_resolution_scale_property(): + return bpy.props.FloatProperty( + name="Resolution Scale", + description=( + "Scale the source image long edge for DA3 processing; the result " + "is limited to 252-1008 px and aligned to a multiple of 14" + ), + default=100.0, + min=1.0, + max=100.0, + step=100, + precision=0, + subtype="PERCENTAGE", + ) + + +def draw_process_resolution(layout, owner, image): + process_res = scaled_process_resolution( + image, + owner.process_res_factor, + ) + layout.prop( + owner, + "process_res_factor", + text=f"Resolution Scale ({process_res} px)", + ) + + +def save_packed_image(image, scene): + output_path = Path(tempfile.mkdtemp(prefix="snapform-input-")) / "input.png" + image_settings = scene.render.image_settings + original_file_format = image_settings.file_format + try: + image_settings.file_format = "PNG" + image.save_render(str(output_path), scene=scene) + except (OSError, RuntimeError) as error: + output_path.unlink(missing_ok=True) + raise RuntimeError("Unable to export the packed image") from error + finally: + image_settings.file_format = original_file_format + if not output_path.is_file(): + raise RuntimeError("Blender did not export the packed image") + return output_path + + +def collect_sequence_frames(image): + """Return the numbered frame files of a SEQUENCE image, or [] when absent.""" + source_path = image_path(image) + if not is_animated_image(image) or not source_path.is_file(): + return [] + match = re.match( + r"^(?P.*?)(?P\d+)(?P\.[^.]+)$", + source_path.name, + ) + if match is None: + return [] + prefix = match.group("prefix") + suffix = match.group("suffix").lower() + frames = [] + for path in source_path.parent.iterdir(): + frame_match = re.match( + r"^(?P.*?)(?P\d+)(?P\.[^.]+)$", + path.name, + ) + if ( + frame_match is not None + and frame_match.group("prefix") == prefix + and frame_match.group("suffix").lower() == suffix + ): + frames.append((int(frame_match.group("digits")), path)) + return [path for _digits, path in sorted(frames)] + + +def prepare_image_input(image, scene): + source_path = image_path(image) + if image_is_packed(image): + return save_packed_image(image, scene), True + if getattr(image, "source", None) == "MOVIE": + if not source_path.is_file(): + raise RuntimeError( + "The image empty must reference an existing video file" + ) + return source_path, False + sequence_frames = collect_sequence_frames(image) + if sequence_frames: + temporary_directory = Path(tempfile.mkdtemp(prefix="snapform-input-")) + try: + for index, frame_path in enumerate(sequence_frames): + suffix = frame_path.suffix.lower() or ".jpg" + shutil.copy2( + frame_path, + temporary_directory / f"frame.{index + 1:04d}{suffix}", + ) + except Exception: + shutil.rmtree(temporary_directory, ignore_errors=True) + raise + return temporary_directory, True + if not source_path.is_file(): + raise RuntimeError( + "The image empty must reference an existing image file" + ) + return source_path, False + + +def configure_image_user( + image_user, + frame_count=None, + frame_start=1, + frame_offset=0, + use_cyclic=None, +): + """Map the image to scene frames and enable auto refresh.""" + if image_user is None: + return + if hasattr(image_user, "frame_start"): + image_user.frame_start = frame_start + if hasattr(image_user, "frame_offset"): + image_user.frame_offset = frame_offset + if hasattr(image_user, "use_auto_refresh"): + image_user.use_auto_refresh = True + if use_cyclic is not None and hasattr(image_user, "use_cyclic"): + image_user.use_cyclic = use_cyclic + if frame_count is not None and hasattr(image_user, "frame_duration"): + image_user.frame_duration = frame_count + + +def pack_sequence_if_possible(image, expected_frames): + """Pack a frame sequence when Blender supports it; otherwise keep external files.""" + if expected_frames <= 1: + return False + try: + image.pack() + except (OSError, RuntimeError): + return False + if len(getattr(image, "packed_files", ())) >= expected_frames: + return True + try: + image.unpack(method="REMOVE") + except (OSError, RuntimeError): + pass + return False + + +def sorted_frame_paths(directory, prefix=None): + """Return image files in a directory sorted by their trailing frame number.""" + def frame_key(path): + match = re.search(r"(\d+)(?=\.[^.]+$)", path.name) + return (int(match.group(1)) if match is not None else 0, path.name) + + return sorted( + ( + path + for path in Path(directory).iterdir() + if path.is_file() + and path.suffix.lower() in IMAGE_NAME_SUFFIXES + and (prefix is None or path.stem.startswith(prefix)) + ), + key=frame_key, + ) + + +def force_sequence_source(image): + """Ensure a loaded image is a SEQUENCE datablock for its numbered frames.""" + if getattr(image, "source", None) == "SEQUENCE": + return image + try: + image.source = "SEQUENCE" + except (OSError, RuntimeError, TypeError) as error: + raise RuntimeError( + "Blender cannot load the frame sequence" + ) from error + if getattr(image, "source", None) != "SEQUENCE": + raise RuntimeError("Blender cannot load the frame sequence") + return image + + +def create_sequence_image(name, frame_directory, frame_prefix="frame."): + """Load a numbered frame directory as a SEQUENCE image datablock.""" + frame_paths = sorted_frame_paths(frame_directory, frame_prefix) + if not frame_paths: + raise RuntimeError(f"No image frames found in: {frame_directory}") + try: + image = bpy.data.images.load( + str(frame_paths[0]), + check_existing=False, + ) + force_sequence_source(image) + except (OSError, RuntimeError) as error: + if "image" in locals(): + bpy.data.images.remove(image, do_unlink=True) + raise RuntimeError("Unable to load the frame sequence") from error + image.name = name + pack_sequence_if_possible(image, len(frame_paths)) + return image + + + +def create_color_image( + source_object, + output_directory, + assign_to_source=True, +): + output_directory = Path(output_directory) + color_paths = sorted_frame_paths(output_directory, "color.") + if not color_paths: + raise RuntimeError( + f"Color result has no images: {output_directory}" + ) + source_image = source_object.data + animated = is_animated_image(source_image) + try: + result_image = bpy.data.images.load( + str(color_paths[0]), + check_existing=False, + ) + result_image.alpha_mode = "CHANNEL_PACKED" + if animated: + force_sequence_source(result_image) + except (OSError, RuntimeError) as error: + if "result_image" in locals(): + bpy.data.images.remove(result_image, do_unlink=True) + raise RuntimeError("Unable to load the color result") from error + + original_fake_user = source_image.use_fake_user + try: + result_image.name = color_data_name(source_object) + if animated: + pack_sequence_if_possible(result_image, len(color_paths)) + else: + result_image.pack() + if assign_to_source: + image_user = getattr(source_object, "image_user", None) + ( + original_frame_start, + _original_frame_offset, + _original_frame_duration, + original_use_cyclic, + ) = image_user_settings(source_object) + source_image.use_fake_user = True + source_object.data = result_image + if animated: + configure_image_user( + image_user, + len(color_paths), + frame_start=original_frame_start, + use_cyclic=original_use_cyclic, + ) + except (OSError, RuntimeError) as error: + source_image.use_fake_user = original_fake_user + bpy.data.images.remove(result_image, do_unlink=True) + raise RuntimeError("Unable to prepare the color result") from error + return result_image + + +def plane_bounds(source_object): + image_width, image_height = source_object.data.size + if image_width <= 0 or image_height <= 0: + raise ValueError("The source image has no readable dimensions") + + aspect = image_width / image_height + display_size = float(source_object.empty_display_size) + if aspect >= 1.0: + width = display_size + height = display_size / aspect + else: + width = display_size * aspect + height = display_size + + offset = source_object.empty_image_offset + x_min = float(offset[0]) * width + y_min = float(offset[1]) * height + return ( + x_min, + x_min + width, + y_min, + y_min + height, + ) diff --git a/src/snapform/operators/material.py b/src/snapform/operators/material.py new file mode 100644 index 0000000..8349299 --- /dev/null +++ b/src/snapform/operators/material.py @@ -0,0 +1,693 @@ +from pathlib import Path + +import bpy + +from .image import configure_image_user, is_animated_image +DISPLACEMENT_STRENGTH = 1.0 +DISPLACEMENT_MIDLEVEL = 0.0 +ORIGINAL_DISPLACEMENT_STRENGTH = 1.0 +ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0 +BUMP_BASE_DISTANCE = 0.001 +DEPTH_PLANE_NODE_GROUP_NAME = "O DA3 Depth Plane" +MOGE_DEPTH_PLANE_NODE_GROUP_NAME = "O MoGe2 Depth Plane" +NODE_ASSET_NAME = "O_SnapForm.blend" +MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement" +BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth" +DEPTH_SCALE_ATTRIBUTE = "o_depth_scale" +SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer" +TEXTURE_MATERIAL_NODE_GROUP_NAME = "O Texture Layer" +MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Surface Layer" +DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer" +MOGE_MATERIAL_NODE_GROUP_NAME = "O MoGe2 Depth Plane Layer" +SHADELESS_NODE_GROUP_NAME = "O Shadeless" +VIEW_TRANSFORM_COLOR_SPACES = ( + ("ACES 1.3", "ACES 1.3 sRGB"), + ("ACES 2.0", "ACES 2.0 sRGB"), + ("AgX", "AgX Base sRGB"), + ("Filmic", "Filmic sRGB"), + ("Khronos PBR Neutral", "Khronos PBR Neutral sRGB"), +) + +def material_color_space_name(scene): + view_settings = getattr(scene, "view_settings", None) + view_transform = getattr(view_settings, "view_transform", "") + normalized = str(view_transform).strip().casefold() + for view_name, color_space in VIEW_TRANSFORM_COLOR_SPACES: + if normalized.startswith(view_name.casefold()): + return color_space + return "sRGB" + + +def set_material_color_space(image, scene=None): + if scene is None: + scene = getattr(bpy.context, "scene", None) + settings = image.colorspace_settings + preferred = material_color_space_name(scene) + original = settings.name + for color_space in dict.fromkeys((preferred, "sRGB", original)): + if not color_space: + continue + try: + settings.name = color_space + except (TypeError, ValueError): + continue + if settings.name == color_space: + return color_space + return settings.name + + + +def moge_material_node_group(): + existing = bpy.data.node_groups.get(MOGE_MATERIAL_NODE_GROUP_NAME) + if existing is not None: + return existing + group = bpy.data.node_groups.new( + MOGE_MATERIAL_NODE_GROUP_NAME, + "ShaderNodeTree", + ) + for name, socket_type in ( + ("Color", "NodeSocketColor"), + ("Alpha", "NodeSocketFloat"), + ("Depth", "NodeSocketFloat"), + ("Bump Scale", "NodeSocketFloat"), + ): + socket = group.interface.new_socket( + name=name, + in_out="INPUT", + socket_type=socket_type, + ) + if name == "Alpha": + socket.default_value = 1.0 + if name == "Bump Scale": + socket.default_value = 1.0 + for name, socket_type in ( + ("Color", "NodeSocketColor"), + ("Roughness", "NodeSocketFloat"), + ("Alpha", "NodeSocketFloat"), + ("Normal", "NodeSocketVector"), + ("Displacement", "NodeSocketVector"), + ): + group.interface.new_socket( + name=name, + in_out="OUTPUT", + socket_type=socket_type, + ) + nodes = group.nodes + links = group.links + group_input = nodes.new("NodeGroupInput") + group_input.location = (-900.0, 100.0) + bump = nodes.new("ShaderNodeBump") + bump.location = (250.0, -100.0) + bump_distance = nodes.new("ShaderNodeMath") + bump_distance.operation = "MULTIPLY" + bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE + bump_distance.location = (0.0, -100.0) + background = nodes.new("ShaderNodeAttribute") + background.attribute_type = "GEOMETRY" + background.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE + background.location = (-650.0, -400.0) + material_mode = nodes.new("ShaderNodeAttribute") + material_mode.attribute_type = "GEOMETRY" + material_mode.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE + material_mode.location = (-650.0, -600.0) + strength = nodes.new("ShaderNodeAttribute") + strength.attribute_type = "GEOMETRY" + strength.attribute_name = DEPTH_SCALE_ATTRIBUTE + strength.location = (-400.0, -600.0) + clamp_depth = nodes.new("ShaderNodeMath") + clamp_depth.operation = "MAXIMUM" + clamp_depth.location = (-350.0, -350.0) + enable_depth = nodes.new("ShaderNodeMath") + enable_depth.operation = "MULTIPLY" + enable_depth.location = (-100.0, -350.0) + scale_depth = nodes.new("ShaderNodeMath") + scale_depth.operation = "MULTIPLY" + scale_depth.location = (150.0, -350.0) + displacement = nodes.new("ShaderNodeDisplacement") + displacement.inputs["Midlevel"].default_value = 0.0 + displacement.inputs["Scale"].default_value = 1.0 + displacement.location = (400.0, -350.0) + group_output = nodes.new("NodeGroupOutput") + group_output.location = (300.0, 100.0) + links.new(group_input.outputs["Color"], group_output.inputs["Color"]) + links.new(group_input.outputs["Alpha"], group_output.inputs["Alpha"]) + links.new(group_input.outputs["Color"], bump.inputs["Height"]) + links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0]) + links.new(bump_distance.outputs[0], bump.inputs["Distance"]) + links.new(bump.outputs["Normal"], group_output.inputs["Normal"]) + links.new(group_input.outputs["Depth"], clamp_depth.inputs[0]) + links.new(background.outputs["Fac"], clamp_depth.inputs[1]) + links.new(clamp_depth.outputs[0], enable_depth.inputs[0]) + links.new(material_mode.outputs["Fac"], enable_depth.inputs[1]) + links.new(enable_depth.outputs[0], scale_depth.inputs[0]) + links.new(strength.outputs["Fac"], scale_depth.inputs[1]) + links.new(scale_depth.outputs[0], displacement.inputs["Height"]) + links.new(displacement.outputs["Displacement"], group_output.inputs["Displacement"]) + group_output.inputs["Roughness"].default_value = 0.5 + group.color_tag = "SHADER" + group.description = "Apply MoGe-2 color and derive bump detail from it." + group.use_fake_user = True + return group + + +def create_moge_material( + name, + color_image, + depth_image, + depth_scale, + shadeless=False, + frame_count=1, + frame_start=1, + use_cyclic=False, + scene=None, +): + material = bpy.data.materials.new(name) + material.use_nodes = True + if hasattr(material, "surface_render_method"): + material.surface_render_method = "DITHERED" + nodes = material.node_tree.nodes + nodes.clear() + links = material.node_tree.links + uv = nodes.new("ShaderNodeUVMap") + uv.uv_map = "UVMap" + uv.location = (-900.0, 200.0) + color = nodes.new("ShaderNodeTexImage") + set_material_color_space(color_image, scene) + color.image = color_image + color.interpolation = "Linear" + color.extension = "EXTEND" + color.location = (-650.0, 300.0) + depth = nodes.new("ShaderNodeTexImage") + depth.image = depth_image + depth.interpolation = "Linear" + depth.extension = "EXTEND" + depth.location = (-650.0, -100.0) + if frame_count > 1: + configure_image_user( + color.image_user, + frame_count, + frame_start=frame_start, + use_cyclic=use_cyclic, + ) + configure_image_user( + depth.image_user, + frame_count, + frame_start=frame_start, + use_cyclic=use_cyclic, + ) + layer = nodes.new("ShaderNodeGroup") + layer.node_tree = moge_material_node_group() + layer.name = MOGE_MATERIAL_NODE_GROUP_NAME + layer.location = (-300.0, 200.0) + output = nodes.new("ShaderNodeOutputMaterial") + output.location = (500.0, 200.0) + links.new(uv.outputs["UV"], color.inputs["Vector"]) + links.new(uv.outputs["UV"], depth.inputs["Vector"]) + links.new(color.outputs["Color"], layer.inputs["Color"]) + links.new(color.outputs["Alpha"], layer.inputs["Alpha"]) + links.new(depth.outputs["Color"], layer.inputs["Depth"]) + layer.inputs["Bump Scale"].default_value = depth_scale + if shadeless: + shader = nodes.new("ShaderNodeGroup") + shader.node_tree = load_shadeless_node_group() + shader.name = SHADELESS_NODE_GROUP_NAME + transparent = nodes.new("ShaderNodeBsdfTransparent") + mix = nodes.new("ShaderNodeMixShader") + links.new(layer.outputs["Color"], shader.inputs["Color"]) + links.new(layer.outputs["Alpha"], mix.inputs["Fac"]) + links.new(transparent.outputs["BSDF"], mix.inputs[1]) + links.new(shader.outputs["Shader"], mix.inputs[2]) + links.new(mix.outputs["Shader"], output.inputs["Surface"]) + 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) + return material + + + +def material_node_group(depth_plane=False, surface_group_name=None): + group_name = ( + DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME + if depth_plane + else surface_group_name or SURFACE_MATERIAL_NODE_GROUP_NAME + ) + existing = bpy.data.node_groups.get(group_name) + if existing is not None: + return existing + group = bpy.data.node_groups.new(group_name, "ShaderNodeTree") + interface = group.interface + input_sockets = [ + ("Color", "NodeSocketColor"), + ("Alpha", "NodeSocketFloat"), + ("Bump Scale", "NodeSocketFloat"), + ] + if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME: + input_sockets.insert(2, ("Normal", "NodeSocketColor")) + input_sockets.insert(3, ("Normal Strength", "NodeSocketFloat")) + if depth_plane: + input_sockets.insert(2, ("Depth", "NodeSocketFloat")) + for name, socket_type in input_sockets: + socket = interface.new_socket( + name=name, + in_out="INPUT", + socket_type=socket_type, + ) + if name == "Alpha": + socket.default_value = 1.0 + socket.min_value = 0.0 + socket.max_value = 1.0 + socket.subtype = "FACTOR" + elif name == "Bump Scale": + socket.default_value = 1.0 + socket.min_value = -10000.0 + socket.max_value = 10000.0 + elif name == "Normal Strength": + socket.default_value = 1.0 + socket.min_value = 0.0 + socket.max_value = 10.0 + output_sockets = [ + ("Color", "NodeSocketColor"), + ("Roughness", "NodeSocketFloat"), + ("Alpha", "NodeSocketFloat"), + ("Normal", "NodeSocketVector"), + ] + if depth_plane: + output_sockets.append(("Displacement", "NodeSocketVector")) + for name, socket_type in output_sockets: + interface.new_socket( + name=name, + in_out="OUTPUT", + socket_type=socket_type, + ) + + nodes = group.nodes + links = group.links + group_input = nodes.new("NodeGroupInput") + group_input.location = (-500.0, 100.0) + group_output = nodes.new("NodeGroupOutput") + group_output.location = (300.0, 100.0) + roughness_curve = nodes.new("ShaderNodeFloatCurve") + roughness_curve.mapping.initialize() + curve_points = roughness_curve.mapping.curves[0].points + curve_points[0].location = (0.0, 0.5) + curve_points[-1].location = (1.0, 0.1) + roughness_curve.mapping.update() + roughness_curve.location = (-250.0, 350.0) + bump = nodes.new("ShaderNodeBump") + bump.location = (-250.0, 0.0) + if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME: + normal_map = nodes.new("ShaderNodeNormalMap") + normal_map.space = "OBJECT" + normal_map.location = (-500.0, -150.0) + links.new(group_input.outputs["Normal"], normal_map.inputs["Color"]) + links.new( + group_input.outputs["Normal Strength"], + normal_map.inputs["Strength"], + ) + links.new(normal_map.outputs["Normal"], bump.inputs["Normal"]) + displacement = None + + processed_alpha = group_input.outputs["Alpha"] + bump_depth = group_input.outputs["Color"] + if depth_plane: + displacement = nodes.new("ShaderNodeDisplacement") + displacement.inputs["Midlevel"].default_value = 0.0 + displacement.inputs["Scale"].default_value = 1.0 + group_input.location = (-1200.0, 100.0) + group_output.location = (1050.0, 100.0) + roughness_curve.location = (-300.0, 450.0) + bump.location = (750.0, 50.0) + displacement.location = (750.0, -300.0) + background_attribute = nodes.new("ShaderNodeAttribute") + background_attribute.attribute_type = "GEOMETRY" + background_attribute.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE + background_attribute.location = (-950.0, -500.0) + material_attribute = nodes.new("ShaderNodeAttribute") + material_attribute.attribute_type = "GEOMETRY" + material_attribute.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE + material_attribute.location = (-50.0, -650.0) + strength_attribute = nodes.new("ShaderNodeAttribute") + strength_attribute.attribute_type = "GEOMETRY" + strength_attribute.attribute_name = DEPTH_SCALE_ATTRIBUTE + strength_attribute.location = (200.0, -650.0) + foreground_mask = nodes.new("ShaderNodeMath") + foreground_mask.operation = "GREATER_THAN" + foreground_mask.inputs[1].default_value = 0.5 + foreground_mask.location = (-950.0, -250.0) + clamp_foreground = nodes.new("ShaderNodeMath") + clamp_foreground.operation = "MAXIMUM" + clamp_foreground.location = (-650.0, -400.0) + masked_depth = nodes.new("ShaderNodeMix") + masked_depth.data_type = "FLOAT" + masked_depth.location = (-350.0, -350.0) + enabled_depth = nodes.new("ShaderNodeMath") + enabled_depth.operation = "MULTIPLY" + enabled_depth.location = (250.0, -350.0) + scaled_depth = nodes.new("ShaderNodeMath") + scaled_depth.operation = "MULTIPLY" + scaled_depth.location = (450.0, -350.0) + links.new(group_input.outputs["Alpha"], foreground_mask.inputs[0]) + links.new(group_input.outputs["Depth"], clamp_foreground.inputs[0]) + links.new(background_attribute.outputs["Fac"], clamp_foreground.inputs[1]) + links.new(foreground_mask.outputs[0], masked_depth.inputs[0]) + links.new(background_attribute.outputs["Fac"], masked_depth.inputs[2]) + links.new(clamp_foreground.outputs[0], masked_depth.inputs[3]) + links.new(masked_depth.outputs["Result"], enabled_depth.inputs[0]) + links.new(material_attribute.outputs["Fac"], enabled_depth.inputs[1]) + links.new(enabled_depth.outputs[0], scaled_depth.inputs[0]) + links.new(strength_attribute.outputs["Fac"], scaled_depth.inputs[1]) + processed_alpha = foreground_mask.outputs[0] + processed_depth = scaled_depth.outputs[0] + + links.new(group_input.outputs["Color"], group_output.inputs["Color"]) + links.new(group_input.outputs["Color"], roughness_curve.inputs["Value"]) + links.new( + roughness_curve.outputs["Value"], + group_output.inputs["Roughness"], + ) + links.new(processed_alpha, group_output.inputs["Alpha"]) + links.new(bump_depth, bump.inputs["Height"]) + bump_distance = nodes.new("ShaderNodeMath") + bump_distance.operation = "MULTIPLY" + bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE + bump_distance.location = (500.0, 100.0) if depth_plane else (-250.0, 150.0) + links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0]) + links.new(bump_distance.outputs[0], bump.inputs["Distance"]) + links.new(bump.outputs["Normal"], group_output.inputs["Normal"]) + if displacement is not None: + links.new(processed_depth, displacement.inputs["Height"]) + links.new( + displacement.outputs["Displacement"], + group_output.inputs["Displacement"], + ) + group.color_tag = "SHADER" + group.description = ( + "Prepare DA3 depth-plane material outputs." + if depth_plane + else f"Prepare {group_name} material outputs." + ) + group.use_fake_user = True + return group + + +def create_material( + name, + color_image, + depth_scale, + displacement_image=None, + displacement_midlevel=0.0, + displacement_scale=1.0, + shadeless=False, + frame_count=1, + frame_start=1, + color_frame_offset=0, + color_frame_duration=None, + depth_frame_offset=0, + depth_frame_duration=None, + use_cyclic=None, + surface_group_name=None, + normal_image=None, + scene=None, +): + material = bpy.data.materials.new(name) + material.use_nodes = True + if hasattr(material, "surface_render_method"): + material.surface_render_method = "DITHERED" + nodes = material.node_tree.nodes + nodes.clear() + + color_texture = nodes.new("ShaderNodeTexImage") + uv_map = nodes.new("ShaderNodeUVMap") + output = nodes.new("ShaderNodeOutputMaterial") + + uv_map.uv_map = "UVMap" + uv_map.location = (-1250.0, 350.0) + material_image = color_image + set_material_color_space(material_image, scene) + color_texture.image = material_image + color_texture.interpolation = "Linear" + color_texture.extension = "EXTEND" + color_texture.location = (-1000.0, 350.0) + output.location = (850.0, 350.0) + extracted_frames = frame_count if frame_count > 1 else None + color_frames = ( + color_frame_duration + if color_frame_duration is not None + else extracted_frames + ) + depth_frames = ( + depth_frame_duration + if depth_frame_duration is not None + else extracted_frames + ) + material_frames = color_frames + material_frame_offset = color_frame_offset + if frame_count > 1 or is_animated_image(material_image): + configure_image_user( + color_texture.image_user, + material_frames, + frame_start=frame_start, + frame_offset=material_frame_offset, + use_cyclic=use_cyclic, + ) + + links = material.node_tree.links + links.new(uv_map.outputs["UV"], color_texture.inputs["Vector"]) + alpha = color_texture.outputs["Alpha"] + material_inputs = nodes.new("ShaderNodeGroup") + is_depth_plane = displacement_image is not None + material_inputs.node_tree = material_node_group( + is_depth_plane, + surface_group_name=surface_group_name, + ) + material_inputs.name = material_inputs.node_tree.name + material_inputs.label = material_inputs.name + material_inputs.location = (-300.0, 350.0) + material_inputs.inputs["Bump Scale"].default_value = depth_scale + links.new(color_texture.outputs["Color"], material_inputs.inputs["Color"]) + links.new(alpha, material_inputs.inputs["Alpha"]) + if normal_image is not None and "Normal" in material_inputs.inputs: + normal_texture = nodes.new("ShaderNodeTexImage") + normal_texture.image = normal_image + normal_texture.interpolation = "Linear" + normal_texture.extension = "EXTEND" + normal_texture.location = (-1000.0, -50.0) + links.new(uv_map.outputs["UV"], normal_texture.inputs["Vector"]) + links.new(normal_texture.outputs["Color"], material_inputs.inputs["Normal"]) + if shadeless: + shader = nodes.new("ShaderNodeGroup") + shader.node_tree = load_shadeless_node_group() + shader.name = SHADELESS_NODE_GROUP_NAME + transparent = nodes.new("ShaderNodeBsdfTransparent") + mix_shader = nodes.new("ShaderNodeMixShader") + shader.location = (0.0, 550.0) + transparent.location = (0.0, 150.0) + mix_shader.location = (300.0, 350.0) + links.new( + material_inputs.outputs["Color"], + shader.inputs["Color"], + ) + links.new( + material_inputs.outputs["Alpha"], + mix_shader.inputs["Fac"], + ) + links.new( + transparent.outputs["BSDF"], + mix_shader.inputs[1], + ) + links.new( + shader.outputs["Shader"], + mix_shader.inputs[2], + ) + links.new( + mix_shader.outputs["Shader"], + output.inputs["Surface"], + ) + else: + shader = nodes.new("ShaderNodeBsdfPrincipled") + shader.location = (250.0, 400.0) + + links.new( + 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"], + ) + links.new(shader.outputs["BSDF"], output.inputs["Surface"]) + + if displacement_image is not None: + depth_texture = nodes.new("ShaderNodeTexImage") + depth_texture.image = displacement_image + depth_texture.interpolation = "Linear" + depth_texture.extension = "EXTEND" + if frame_count > 1 or is_animated_image(displacement_image): + configure_image_user( + depth_texture.image_user, + depth_frames, + frame_start=frame_start, + frame_offset=depth_frame_offset, + use_cyclic=use_cyclic, + ) + depth_texture.location = (-1000.0, -400.0) + links.new(uv_map.outputs["UV"], depth_texture.inputs["Vector"]) + displacement_height = depth_texture.outputs["Color"] + if displacement_midlevel != 0.0 or displacement_scale != 1.0: + depth_adjust = nodes.new("ShaderNodeMath") + depth_adjust.operation = "SUBTRACT" + depth_adjust.inputs[1].default_value = displacement_midlevel + depth_adjust.location = (500.0, -550.0) + depth_scale_node = nodes.new("ShaderNodeMath") + depth_scale_node.operation = "MULTIPLY" + depth_scale_node.inputs[1].default_value = displacement_scale + depth_scale_node.location = (700.0, -550.0) + links.new(displacement_height, depth_adjust.inputs[0]) + links.new(depth_adjust.outputs[0], depth_scale_node.inputs[0]) + links.new(depth_scale_node.outputs[0], material_inputs.inputs["Depth"]) + else: + links.new(displacement_height, material_inputs.inputs["Depth"]) + links.new( + material_inputs.outputs["Displacement"], + output.inputs["Displacement"], + ) + set_displacement_only(material) + return material + + +def set_displacement_only(material): + for owner, property_name in ( + (material, "surface_displacement_method"), + (material, "displacement_method"), + (getattr(material, "cycles", None), "displacement_method"), + ): + if owner is not None and hasattr(owner, property_name): + setattr(owner, property_name, "DISPLACEMENT") + return + + +def node_asset_path(): + return Path(__file__).resolve().parents[1] / "assets" / NODE_ASSET_NAME + + +def load_asset_node_group(group_name): + existing = bpy.data.node_groups.get(group_name) + if existing is not None: + return existing + asset_path = node_asset_path() + if not asset_path.is_file(): + raise RuntimeError(f"Missing node asset: {asset_path}") + use_pack = bpy.app.version >= (5, 0, 0) + if not use_pack: + directory = str(asset_path) + "/NodeTree/" + result = bpy.ops.wm.append( + filepath=directory + group_name, + directory=directory, + filename=group_name, + link=False, + do_reuse_local_id=True, + ) + if "FINISHED" not in result: + raise RuntimeError("Unable to append the node asset") + node_group = bpy.data.node_groups.get(group_name) + if node_group is None: + raise RuntimeError("Appended node asset is unavailable") + node_group.use_fake_user = True + return node_group + with bpy.data.libraries.load(str(asset_path), link=True) as (source, target): + if group_name not in source.node_groups: + raise RuntimeError(f"Node asset is missing {group_name}") + target.node_groups = [group_name] + linked_group = target.node_groups[0] + node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group) + if linked_group != node_group and linked_group.users == 0: + bpy.data.node_groups.remove(linked_group) + node_group.use_fake_user = True + return node_group + + +def load_shadeless_node_group(): + return load_asset_node_group(SHADELESS_NODE_GROUP_NAME) + + +def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME): + return load_asset_node_group(group_name) + + +def modifier_input_identifier(node_group, name): + for item in node_group.interface.items_tree: + if item.item_type != "SOCKET" or item.in_out != "INPUT": + continue + if item.name == name: + return item.identifier + raise RuntimeError(f"Geometry Nodes input is missing: {name}") + + +def modifier_input_slots(modifier): + """Return the runtime input collection, or None for the legacy API.""" + properties = getattr(modifier, "properties", None) + if properties is None: + return None + return getattr(properties, "inputs", None) + + +def set_modifier_input(modifier, identifier, value): + """Assign a value to a Geometry Nodes modifier input socket.""" + inputs = modifier_input_slots(modifier) + if inputs is not None: + # Blender 5.2+ exposes inputs as runtime RNA keyed by socket identifier. + getattr(inputs, identifier).value = value + else: + # Blender 4.x exposes inputs as modifier ID properties. + modifier[identifier] = value + + +def get_modifier_input(modifier, identifier): + """Read a Geometry Nodes modifier input socket value.""" + inputs = modifier_input_slots(modifier) + if inputs is not None: + return getattr(inputs, identifier).value + return modifier[identifier] + + +def modifier_input_driver_path(modifier, identifier): + """Return the driver data path for a modifier input socket.""" + if modifier_input_slots(modifier) is not None: + return ( + f'modifiers["{modifier.name}"]' + f".properties.inputs.{identifier}.value" + ) + return f'modifiers["{modifier.name}"]["{identifier}"]' + + +def drive_modifier_input(owner_object, modifier, identifier, expression="frame"): + """Add a scripted driver to a modifier input so it follows the scene frame.""" + driver = owner_object.driver_add( + modifier_input_driver_path(modifier, identifier) + ) + driver.driver.type = "SCRIPTED" + driver.driver.expression = expression + return driver diff --git a/src/snapform/operators/models.py b/src/snapform/operators/models.py new file mode 100644 index 0000000..d0deb10 --- /dev/null +++ b/src/snapform/operators/models.py @@ -0,0 +1,48 @@ +import bpy + +from ..runtime import JobOperatorBase, runtime +from ..properties import REQUIRED_MODEL_KEYS, model_record +from .common import HIDDEN, require_environment + + +class DownloadModel(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.download_model" + bl_label = "Download Model" + job_type = "download-model" + + model: bpy.props.StringProperty(options=HIDDEN) + + def request(self, _context): + require_environment() + if not bpy.app.online_access: + raise RuntimeError("Online Access is disabled in Blender Preferences") + model = model_record(self.model) + label = model["label"] if model else self.model + self.starting_message = f"Starting {label} download" + return {"model": self.model} + + +class DownloadRequiredModels(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.download_required_models" + bl_label = "Download Required Models" + job_type = "download-required-models" + starting_message = "Starting required model downloads" + + def request(self, _context): + require_environment() + if not bpy.app.online_access: + raise RuntimeError("Online Access is disabled in Blender Preferences") + return {"models": list(REQUIRED_MODEL_KEYS)} + + +class ClearModels(bpy.types.Operator): + bl_idname = "snapform.clear_models" + bl_label = "Clear Models" + + def execute(self, _context): + try: + runtime.clear_resource("model_manager") + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} diff --git a/src/snapform/blender/surface.py b/src/snapform/operators/surface.py similarity index 79% rename from src/snapform/blender/surface.py rename to src/snapform/operators/surface.py index 9c13cee..66f5636 100644 --- a/src/snapform/blender/surface.py +++ b/src/snapform/operators/surface.py @@ -2,9 +2,21 @@ from pathlib import Path import bpy -from .depth_plane import ( - MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, - create_material, +from ..properties import BEN2_MODEL_KEY, DEFAULT_MOGE2_MODEL_KEY +from ..runtime import JobOperatorBase +from .common import ( + HIDDEN, + cleanup_temporary_input, + da3_parameters, + input_path, + moge2_parameters, + push_undo, + require_environment, + require_model, + source_start_frame, +) + +from .image import ( draw_process_resolution, is_image_empty, normal_data_name, @@ -13,9 +25,13 @@ from .depth_plane import ( process_resolution_scale_property, remove_temporary_input, scaled_process_resolution, - set_modifier_input, video_total_frames, ) +from .material import ( + MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, + create_material, + set_modifier_input, +) POINT_ATTRIBUTES = { @@ -530,7 +546,7 @@ class ConvertToSurface(bpy.types.Operator): @classmethod def poll(cls, context): - from ..job_runtime import runtime + from ..runtime import runtime from ..properties import required_models_ready return ( @@ -601,3 +617,146 @@ class ConvertToSurface(bpy.types.Operator): remove_temporary_input(source_path) return {"CANCELLED"} return {"FINISHED"} + + + + +def build_surface_response( + context, + result, + *, + source_object_name, + delete_source, + create_camera, + shadeless, + moge2, +): + from .image import create_color_image, is_image_empty + + + 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") + color_image = create_color_image( + source_object, + result.directory, + assign_to_source=False, + ) + build_image_mesh( + context, + source_object, + result.file("surface"), + delete_source, + create_camera, + shadeless, + moge2, + color_image=color_image, + normal_path=( + result.directory / "normal.png" + if moge2 + else None + ), + ) + push_undo("Convert to Surface") + + +class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.generate_da3_surface" + bl_label = "Generate DA3 Surface" + job_type = "generate-surface" + starting_message = "Starting perspective surface generation" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default="BASE", options=HIDDEN) + process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) + shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **da3_parameters( + settings, + input_path(self.input_path), + self.model, + self.process_res, + source_start_frame(self.source_object_name), + ), + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + } + + def response(self, context, result): + build_surface_response( + context, + result, + source_object_name=self.source_object_name, + delete_source=self.delete_source, + create_camera=self.create_camera, + shadeless=self.shadeless, + moge2=False, + ) + return "Perspective surface is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.generate_moge2_surface" + bl_label = "Generate MoGe-2 Surface" + job_type = "generate-surface" + starting_message = "Starting MoGe-2 inference" + + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) + resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) + shadeless: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **moge2_parameters( + settings, + input_path(self.input_path), + self.model, + self.resolution_level, + source_start_frame(self.source_object_name), + ), + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + } + + def response(self, context, result): + build_surface_response( + context, + result, + source_object_name=self.source_object_name, + delete_source=self.delete_source, + create_camera=self.create_camera, + shadeless=self.shadeless, + moge2=True, + ) + return "MoGe-2 surface is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 4d1adf3..9d47afa 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -1,6 +1,6 @@ import bpy -from .job_runtime import runtime +from .runtime import runtime from .properties import BEN2_MODEL_KEY, model_catalog, model_ready class ServerPanel(bpy.types.Panel): diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index 8951d92..15575b1 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -1,6 +1,6 @@ import bpy -from .job_runtime import DEFAULT_STORAGE_ROOT, runtime +from .runtime import DEFAULT_STORAGE_ROOT, runtime from .properties import device_enum_items diff --git a/src/snapform/properties.py b/src/snapform/properties.py index faae138..bb79583 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -6,7 +6,7 @@ import sys import bpy -from .job_runtime import runtime +from .runtime import runtime from . import model_catalog as shared_model_catalog diff --git a/src/snapform/job_runtime.py b/src/snapform/runtime.py similarity index 94% rename from src/snapform/job_runtime.py rename to src/snapform/runtime.py index ac178ce..7815a23 100644 --- a/src/snapform/job_runtime.py +++ b/src/snapform/runtime.py @@ -9,7 +9,7 @@ DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" def configured_storage_root(): - addon = bpy.context.preferences.addons.get(__package__) + addon = bpy.context.preferences.addons.get(__package__.partition(".")[0]) configured = getattr(getattr(addon, "preferences", None), "storage_root", "") if configured: return Path(bpy.path.abspath(configured)) diff --git a/tests/test_addon.py b/tests/test_addon.py index fead88c..6a9dc39 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -71,11 +71,32 @@ 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.job_runtime = importlib.import_module("snapform.runtime") + cls.model_operators = importlib.import_module("snapform.operators.models") cls.image_depth = importlib.import_module( - "snapform.blender.depth_plane" + "snapform.operators.material" + ) + cls.image_data = importlib.import_module( + "snapform.operators.image" + ) + cls.depth_operator = importlib.import_module( + "snapform.operators.depth_plane.operator" + ) + cls.image_background = importlib.import_module( + "snapform.operators.background" ) cls.image_mesh = importlib.import_module( - "snapform.blender.surface" + "snapform.operators.surface" + ) + cls.balloon = importlib.import_module("snapform.operators.balloon.operator") + cls.balloon_geometry = importlib.import_module( + "snapform.operators.balloon.geometry" + ) + cls.balloon_poisson = importlib.import_module( + "snapform.operators.balloon.poisson" + ) + cls.balloon_teddy = importlib.import_module( + "snapform.operators.balloon.teddy" ) cls.paths_after_import = tuple(sys.path) cls.fake_bpy = fake_bpy @@ -113,7 +134,7 @@ class AddonTest(unittest.TestCase): if hasattr(self.fake_bpy.types.Scene, "snapform_settings"): del self.fake_bpy.types.Scene.snapform_settings - def test_snapform_has_no_runtime_boundary_modules(self): + def test_runtime_and_operator_support_modules_have_explicit_boundaries(self): for name in ( "client.py", "environment.py", @@ -122,7 +143,12 @@ class AddonTest(unittest.TestCase): "status_bar.py", ): self.assertFalse((SOURCE_ROOT / "snapform" / name).exists(), name) - self.assertFalse(hasattr(self.addon, "client")) + self.assertTrue((SOURCE_ROOT / "snapform" / "runtime.py").is_file()) + self.assertFalse((SOURCE_ROOT / "snapform" / "job_runtime.py").exists()) + self.assertFalse((SOURCE_ROOT / "snapform" / "client").exists()) + self.assertFalse((SOURCE_ROOT / "snapform" / "ui").exists()) + self.assertFalse((SOURCE_ROOT / "snapform" / "operators.py").exists()) + self.assertFalse((SOURCE_ROOT / "snapform" / "blender").exists()) def test_mps_requires_m4_or_newer(self): properties = self.addon.properties @@ -184,7 +210,7 @@ class AddonTest(unittest.TestCase): blendjob = importlib.import_module("blendjob") self.assertFalse(hasattr(blendjob, "JobRequest")) - def test_job_operators_are_flat_and_use_plain_properties(self): + def test_operator_modules_use_plain_properties(self): operators = self.addon.operators for name in ( "SnapFormJobMixin", @@ -195,18 +221,18 @@ class AddonTest(unittest.TestCase): "BackgroundRequestMixin", ): self.assertFalse(hasattr(operators, name), name) - for operator in operators.JOB_OPERATOR_CLASSES: + for operator in operators.CLASSES: for name in getattr(operator, "__annotations__", {}): self.assertFalse(name.startswith("snapform_"), (operator, name)) def test_concrete_job_operator_is_the_blender_operator(self): operator = self.addon.operators.GenerateMoge2Depth self.assertEqual(operator.bl_idname, "snapform.generate_moge2_depth") - self.assertIn(self.addon.job_runtime.JobOperatorBase, operator.__mro__) + self.assertIn(self.job_runtime.JobOperatorBase, operator.__mro__) self.assertIn(self.fake_bpy.types.Operator, operator.__mro__) def test_default_request_sends_every_operator_property(self): - class Example(self.addon.job_runtime.JobOperatorBase): + class Example(self.job_runtime.JobOperatorBase): value: object() enabled: object() @@ -219,7 +245,7 @@ class AddonTest(unittest.TestCase): ) def test_overridden_request_purely_replaces_default_properties(self): - class Example(self.addon.job_runtime.JobOperatorBase): + class Example(self.job_runtime.JobOperatorBase): value: object() def request(self, context): @@ -248,7 +274,7 @@ class AddonTest(unittest.TestCase): def mark_job_complete(self, _job_id): pass - class Example(self.addon.job_runtime.JobOperatorBase): + class Example(self.job_runtime.JobOperatorBase): job_type = "example" value: object() @@ -335,7 +361,7 @@ class AddonTest(unittest.TestCase): (job_type, parameters) ) ) - self.addon.job_runtime.post_install(fake_runtime) + self.job_runtime.post_install(fake_runtime) self.assertEqual( calls, [ @@ -355,7 +381,7 @@ class AddonTest(unittest.TestCase): RuntimeError, "MoGe-2 ViT-S Normal: network unavailable; BEN2 Base: network unavailable", ): - self.addon.job_runtime.post_install( + self.job_runtime.post_install( SimpleNamespace(request=request) ) @@ -452,7 +478,7 @@ class AddonTest(unittest.TestCase): def test_required_model_operator_requests_both_models(self): operator = self.addon.operators.DownloadRequiredModels() - with patch.object(self.addon.operators, "require_environment"): + with patch.object(self.model_operators, "require_environment"): parameters = operator.request(None) self.assertEqual( @@ -494,7 +520,11 @@ class AddonTest(unittest.TestCase): self.assertEqual( calls[0], + ("snapform.balloon_lasso", True), + ) + self.assertIn( ("snapform.download_required_models", True), + calls, ) self.assertIn(("snapform.convert_to_depth_plane", False), calls) self.assertIn(("snapform.convert_to_surface", False), calls) @@ -532,24 +562,336 @@ class AddonTest(unittest.TestCase): ): menu.draw(None) - self.assertEqual(calls[0], ("snapform.install_environment", True)) self.assertEqual( - calls[1], - ("snapform.install_environment_mirror", True), + calls[0], + ("snapform.balloon_lasso", True), ) + self.assertIn(("snapform.install_environment", True), calls) + self.assertIn(("snapform.install_environment_mirror", True), calls) self.assertNotIn(("snapform.download_required_models", True), calls) self.assertIn(("snapform.convert_to_depth_plane", False), calls) self.assertIn(("snapform.convert_to_surface", False), calls) self.assertIn(("snapform.remove_image_background", False), calls) + def test_balloon_contour_is_smooth_and_evenly_resampled(self): + import numpy as np + + contour = self.balloon_geometry.smooth_lasso_contour( + ((0, 0), (3, 0), (4, 0), (4, 2), (0, 2)), + 0.25, + ) + closed = np.vstack((contour, contour[0])) + lengths = np.linalg.norm(np.diff(closed, axis=0), axis=1) + + self.assertGreater(len(contour), 20) + self.assertGreater(self.balloon_geometry.polygon_area(contour), 0.0) + self.assertLess(float(lengths.std() / lengths.mean()), 0.15) + + def test_balloon_interior_points_have_uniform_spacing(self): + import numpy as np + + contour = np.asarray(((0, 0), (4, 0), (4, 4), (0, 4)), dtype=float) + points = self.balloon_geometry.uniform_interior_points(contour, 0.5) + distances = np.linalg.norm(points[:, None] - points[None, :], axis=2) + np.fill_diagonal(distances, np.inf) + nearest = distances.min(axis=1) + + self.assertGreater(len(points), 20) + self.assertAlmostEqual(float(np.median(nearest)), 0.5, delta=0.02) + self.assertGreater(float(nearest.min()), 0.49) + + def test_balloon_default_sampling_stays_coarse(self): + import numpy as np + + polygon = np.asarray( + ((100, 100), (900, 100), (950, 500), (800, 900), + (200, 900), (50, 500)), + dtype=float, + ) + contour, interior, _spacing = ( + self.balloon_geometry.prepare_planar_points( + polygon, + (1000, 1000), + (0.0, 1.0, 0.0, 1.0), + ) + ) + + self.assertLess(len(contour), 150) + self.assertLess(len(interior), 400) + def test_balloon_sampling_uses_world_distance(self): + import numpy as np + + polygon = np.asarray( + ((0, 0), (1000, 0), (1000, 1000), (0, 1000)), + dtype=float, + ) + unit_contour, unit_interior, unit_spacing = ( + self.balloon_geometry.prepare_planar_points( + polygon, + (1000, 1000), + (0.0, 1.0, 0.0, 1.0), + ) + ) + scaled_contour, scaled_interior, scaled_spacing = ( + self.balloon_geometry.prepare_planar_points( + polygon, + (1000, 1000), + (0.0, 2.0, 0.0, 2.0), + world_metric=((0.25, 0.0), (0.0, 0.25)), + ) + ) + + self.assertEqual(unit_spacing, 0.1) + self.assertEqual(scaled_spacing, 0.1) + self.assertEqual(len(unit_contour), len(scaled_contour)) + self.assertEqual(len(unit_interior), len(scaled_interior)) + + def test_poisson_balloon_is_zero_on_boundary_and_positive_inside(self): + import numpy as np + + points = np.asarray( + ((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)), + dtype=float, + ) + faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4)) + boundary = np.asarray((True, True, True, True, False)) + heights = self.balloon_poisson.poisson_balloon_heights( + points, faces, boundary, 0.45 + ) + + self.assertTrue((heights[:4] == 0.0).all()) + self.assertGreater(heights[4], 0.0) + + def test_poisson_balloon_scales_in_world_units(self): + import numpy as np + + points = np.asarray( + ((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)), + dtype=float, + ) + faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4)) + boundary = np.asarray((True, True, True, True, False)) + original = self.balloon_poisson.poisson_balloon_heights( + points, faces, boundary, 0.45 + ) + scaled = self.balloon_poisson.poisson_balloon_heights( + points * 2.0, faces, boundary, 0.45 + ) + + self.assertAlmostEqual(scaled[4], original[4] * 2.0, places=6) + + def test_poisson_balloon_makes_wide_regions_thicker(self): + import numpy as np + + faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4)) + boundary = np.asarray((True, True, True, True, False)) + wide = np.asarray( + ((0, 0), (4, 0), (4, 2), (0, 2), (2, 1)), + dtype=float, + ) + narrow = np.asarray( + ((0, 0), (4, 0), (4, 0.6), (0, 0.6), (2, 0.3)), + dtype=float, + ) + wide_height = self.balloon_poisson.poisson_balloon_heights( + wide, faces, boundary, 0.45 + )[4] + narrow_height = self.balloon_poisson.poisson_balloon_heights( + narrow, faces, boundary, 0.45 + )[4] + + self.assertGreater(wide_height, narrow_height) + + def test_teddy_balloon_uses_chordal_axis_radius(self): + import numpy as np + + points = np.asarray( + ((0, 0), (4, 0), (4, 2), (0, 2), (2, 1)), + dtype=float, + ) + faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4)) + boundary = np.asarray((True, True, True, True, False)) + nodes = np.asarray(((1, 1), (3, 1)), dtype=float) + with patch.object( + self.balloon_teddy, + "chordal_axis", + return_value=(nodes, [(0, 1)], np.asarray((1.0, 1.0))), + ): + heights = self.balloon_teddy.teddy_balloon_heights( + points, faces, points[:4], boundary, 0.45 + ) + + self.assertTrue((heights[:4] == 0.0).all()) + self.assertAlmostEqual(heights[4], 0.45, places=6) + + def test_balloon_dispatches_distinct_height_algorithms(self): + import numpy as np + + points = np.asarray(((0, 0),), dtype=float) + boundary = np.asarray((True,)) + with patch.object( + self.balloon_geometry, + "poisson_balloon_heights", + return_value=np.asarray((1.0,)), + ) as poisson, patch.object( + self.balloon_geometry, + "teddy_balloon_heights", + return_value=np.asarray((2.0,)), + ) as teddy: + poisson_result = self.balloon_geometry.balloon_heights( + "POISSON", points, (), points, boundary, 0.45 + ) + teddy_result = self.balloon_geometry.balloon_heights( + "TEDDY", points, (), points, boundary, 0.45 + ) + + self.assertEqual(poisson_result[0], 1.0) + self.assertEqual(teddy_result[0], 2.0) + poisson.assert_called_once() + teddy.assert_called_once() + def test_balloon_geometry_has_a_shared_closed_seam(self): + from collections import Counter + + import numpy as np + + planar = np.asarray( + ((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)), + dtype=float, + ) + faces_2d = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4)) + boundary = np.asarray((True, True, True, True, False)) + heights = np.asarray((0, 0, 0, 0, 1), dtype=float) + vertices, faces, uv = self.balloon_geometry.build_closed_geometry( + planar, + faces_2d, + boundary, + heights, + (0.0, 2.0, 0.0, 2.0), + ) + edges = Counter() + for face in faces: + for start, end in zip(face, face[1:] + face[:1]): + edges[tuple(sorted((start, end)))] += 1 + + self.assertEqual(len(vertices), 6) + self.assertEqual(len(vertices), len(uv)) + self.assertEqual(uv.dtype, np.float32) + self.assertTrue(edges) + self.assertTrue(all(users == 2 for users in edges.values())) + + def test_balloon_controller_and_extract_operator_are_registered(self): + self.assertIn(self.balloon.ExtractBalloon, self.addon.CLASSES) + self.assertIn(self.balloon.BalloonLasso, self.addon.CLASSES) + self.assertLess( + self.addon.CLASSES.index(self.balloon.ExtractBalloon), + self.addon.CLASSES.index(self.balloon.BalloonLasso), + ) + def test_balloon_redo_properties_use_mode_and_world_resolution(self): + import inspect + + spacing, inflation = self.balloon.balloon_geometry_parameters(2.0, 0.2) + + self.assertEqual(spacing, 0.2) + self.assertEqual(inflation, 0.9) + operator = self.balloon.ExtractBalloon + self.assertIn("REGISTER", operator.bl_options) + self.assertIn("mode", operator.__annotations__) + self.assertIn("thickness", operator.__annotations__) + self.assertIn("resolution", operator.__annotations__) + self.assertIn("source_matrix_data", operator.__annotations__) + resolution_source = inspect.getsource(self.balloon.resolution_property) + thickness_source = inspect.getsource(self.balloon.thickness_property) + self.assertIn('subtype="DISTANCE"', resolution_source) + self.assertIn('unit="LENGTH"', resolution_source) + self.assertIn("soft_max=5.0", thickness_source) + self.assertNotIn(" max=", thickness_source) + + def test_balloon_lasso_passes_hud_and_navigation_events(self): + regions = ( + SimpleNamespace(type="WINDOW", x=0, y=0, width=800, height=600), + SimpleNamespace(type="HUD", x=20, y=20, width=240, height=180), + ) + context = SimpleNamespace(area=SimpleNamespace(regions=regions)) + + self.assertTrue( + self.balloon.event_over_interface_region( + context, + SimpleNamespace(mouse_x=40, mouse_y=40), + ) + ) + self.assertFalse( + self.balloon.event_over_interface_region( + context, + SimpleNamespace(mouse_x=500, mouse_y=400), + ) + ) + for event_type in ( + "MIDDLEMOUSE", + "WHEELUPMOUSE", + "TRACKPADPAN", + "NDOF_MOTION", + ): + self.assertIn(event_type, self.balloon.VIEW_NAVIGATION_EVENTS) + self.assertEqual(self.balloon.BalloonLasso.bl_options, {"INTERNAL"}) + def test_balloon_lasso_dispatches_individual_redo_operator(self): + import inspect + + execute_source = inspect.getsource(self.balloon.ExtractBalloon.execute) + invoke_source = inspect.getsource(self.balloon.BalloonLasso.invoke) + modal_source = inspect.getsource(self.balloon.BalloonLasso.modal) + + self.assertIn("json.loads(self.selection_data)", execute_source) + self.assertIn("deserialize_matrix(self.source_matrix_data)", execute_source) + self.assertIn("select_result=not self.keep_source_selected", execute_source) + self.assertIn("serialize_matrix(source_object.matrix_world)", invoke_source) + self.assertIn("bpy.ops.snapform.extract_balloon(", modal_source) + self.assertIn('"EXEC_DEFAULT"', modal_source) + self.assertIn("keep_source_selected=True", modal_source) + self.assertIn('event.type == "ESC"', modal_source) + self.assertIn('return {"RUNNING_MODAL"}', modal_source) + self.assertIn('return {"PASS_THROUGH"}', modal_source) + self.assertIn("event_over_interface_region(context, event)", modal_source) + self.assertNotIn("extract_balloon_object(", modal_source) + def test_balloon_matrix_serialization_keeps_all_components(self): + import json + + matrix = tuple( + tuple(row * 4 + column for column in range(4)) + for row in range(4) + ) + values = json.loads(self.balloon.serialize_matrix(matrix)) + + self.assertEqual(values, list(range(16))) + + def test_extract_balloon_preserves_the_source_object(self): + import inspect + + source = inspect.getsource(self.balloon.extract_balloon_object) + self.assertNotIn("objects.remove(source_object", source) + self.assertIn("balloon_name = source_name", source) + self.assertIn("selected_object = balloon if select_result else source_object", source) + self.assertIn("balloon.matrix_world = matrix_world.copy()", source) + self.assertIn("create_material(", source) + self.assertIn("color_image = image.copy()", source) + self.assertIn("color_image.name = color_data_name(source_object)", source) + self.assertIn("balloon_name,\n color_image,\n 0.0,", source) + self.assertIn( + "surface_group_name=TEXTURE_MATERIAL_NODE_GROUP_NAME", + source, + ) + self.assertEqual( + self.image_depth.TEXTURE_MATERIAL_NODE_GROUP_NAME, + "O Texture Layer", + ) + def test_image_operators_require_environment_and_required_models(self): context = SimpleNamespace( object=SimpleNamespace(data=SimpleNamespace()), ) operators = ( - (self.image_depth, self.addon.ConvertToDepthPlane), - (self.image_mesh, self.addon.ConvertToSurface), - (self.image_depth, self.addon.RemoveImageBackground), + (self.depth_operator, self.depth_operator.ConvertToDepthPlane), + (self.image_mesh, self.image_mesh.ConvertToSurface), + (self.image_background, self.image_background.RemoveImageBackground), ) states = ( @@ -719,8 +1061,8 @@ class AddonTest(unittest.TestCase): ) with patch.object( - self.addon.operators, "require_environment" - ), patch.object(self.addon.operators, "require_model"): + self.depth_operator, "require_environment" + ), patch.object(self.depth_operator, "require_model"): parameters = operator.request(context) self.assertEqual(parameters["input"], __file__) @@ -738,7 +1080,7 @@ class AddonTest(unittest.TestCase): context = object() with patch.object( - self.addon.operators, "build_moge2_depth_response" + self.depth_operator, "build_moge2_depth_response" ) as build: operator.response(context, result) @@ -752,7 +1094,7 @@ class AddonTest(unittest.TestCase): ) def test_entry_operator_dispatches_plain_properties(self): - operator = self.addon.ConvertToDepthPlane() + operator = self.depth_operator.ConvertToDepthPlane() operator.model = "MOGE2_VITS_NORMAL" operator.resolution_level = 7 operator.keep_original = False @@ -776,7 +1118,7 @@ class AddonTest(unittest.TestCase): ) with patch.object( - self.image_depth, + self.depth_operator, "prepare_image_input", return_value=(Path(__file__), False), ): diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 0d470f2..0afed99 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -23,25 +23,66 @@ def load_image_depth(): ) fake_bpy.types = SimpleNamespace(Menu=object, Operator=object) fake_bpy.path = SimpleNamespace(abspath=lambda path: path) - with patch.dict(sys.modules, {"bpy": fake_bpy}): - specification = importlib.util.spec_from_file_location( - "snapform_image_depth_test", - PROJECT_ROOT - / "src" - / "snapform" - / "blender" - / "depth_plane.py", - ) - module = importlib.util.module_from_spec(specification) - specification.loader.exec_module(module) - return module + fake_bpy.context = SimpleNamespace(scene=None) + module_paths = ( + ( + "snapform.operators.image", + PROJECT_ROOT / "src" / "snapform" / "operators" / "image.py", + ), + ( + "snapform.operators.material", + PROJECT_ROOT / "src" / "snapform" / "operators" / "material.py", + ), + ( + "snapform.operators.depth_plane.data", + PROJECT_ROOT / "src" / "snapform" / "operators" / "depth_plane" / "data.py", + ), + ( + "snapform.operators.depth_plane.build", + PROJECT_ROOT / "src" / "snapform" / "operators" / "depth_plane" / "build.py", + ), + ) + packages = { + name: ModuleType(name) + for name in ( + "snapform", + "snapform.operators", + "snapform.operators.depth_plane", + ) + } + for package in packages.values(): + package.__path__ = [] + + loaded = [] + with patch.dict(sys.modules, {"bpy": fake_bpy, **packages}): + for module_name, module_path in module_paths: + specification = importlib.util.spec_from_file_location( + module_name, + module_path, + ) + module = importlib.util.module_from_spec(specification) + sys.modules[module_name] = module + specification.loader.exec_module(module) + loaded.append(module) + + namespace = SimpleNamespace() + for module in loaded: + for name, value in vars(module).items(): + if not name.startswith("__"): + setattr(namespace, name, value) + return namespace class ImageDepthTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.image_depth = load_image_depth() + def test_material_asset_path_survives_package_split(self): + self.assertEqual( + self.image_depth.node_asset_path(), + PROJECT_ROOT / "src" / "snapform" / "assets" / "O_SnapForm.blend", + ) def test_material_color_space_matches_scene_view_transform(self): mappings = { "ACES 1.3": "ACES 1.3 sRGB", @@ -651,7 +692,7 @@ class ImageDepthTest(unittest.TestCase): "img.001.depth", ) - def test_moge_map_names_use_source_image_name(self): + def test_moge_normal_name_uses_source_image_name(self): source_object = SimpleNamespace( name="obj_img", data=SimpleNamespace(name="img.png"), @@ -661,10 +702,6 @@ class ImageDepthTest(unittest.TestCase): self.image_depth.normal_data_name(source_object), "img.normal", ) - self.assertEqual( - self.image_depth.mask_data_name(source_object), - "img.mask", - ) def test_copy_depth_image_renames_still_source(self): class FakeImage: diff --git a/tests/test_pack.py b/tests/test_pack.py index 658e7b0..7fb739d 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -61,8 +61,10 @@ class PackTest(unittest.TestCase): self.assertIn("__init__.py", names) self.assertIn("blender_manifest.toml", names) - self.assertIn("blender/depth_plane.py", names) - self.assertIn("blender/surface.py", names) + self.assertIn("operators/depth_plane/operator.py", names) + self.assertIn("operators/image.py", names) + self.assertIn("operators/balloon/operator.py", names) + self.assertIn("operators/surface.py", names) self.assertIn( "assets/O_SnapForm.blend", names, diff --git a/tools/validate_depth_plane_asset.py b/tools/validate_depth_plane_asset.py index de5ab8a..dc132f8 100644 --- a/tools/validate_depth_plane_asset.py +++ b/tools/validate_depth_plane_asset.py @@ -11,20 +11,20 @@ sys.path.insert( str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), ) -from snapform.blender.depth_plane import ( +from snapform.operators.material import ( DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME, - material_color_space_name, SHADELESS_NODE_GROUP_NAME, SURFACE_MATERIAL_NODE_GROUP_NAME, - add_depth_plane_processing, create_material, get_modifier_input, load_depth_plane_node_group, + material_color_space_name, modifier_input_driver_path, modifier_input_identifier, set_modifier_input, ) -from snapform.blender.surface import load_geometry_node_group +from snapform.operators.depth_plane.build import add_depth_plane_processing +from snapform.operators.surface import load_geometry_node_group def main(): diff --git a/tools/validate_geometry_nodes_asset.py b/tools/validate_geometry_nodes_asset.py index 9846838..e0d21c5 100644 --- a/tools/validate_geometry_nodes_asset.py +++ b/tools/validate_geometry_nodes_asset.py @@ -12,11 +12,8 @@ sys.path.insert( str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), ) -from snapform.blender.depth_plane import ( - get_modifier_input, - set_modifier_input, -) -from snapform.blender.surface import create_matching_camera, load_geometry_node_group +from snapform.operators.material import get_modifier_input, set_modifier_input +from snapform.operators.surface import create_matching_camera, load_geometry_node_group def set_input(modifier, group, name, value): diff --git a/tools/validate_moge_depth_plane.py b/tools/validate_moge_depth_plane.py index cde2db8..175c32c 100644 --- a/tools/validate_moge_depth_plane.py +++ b/tools/validate_moge_depth_plane.py @@ -15,16 +15,14 @@ sys.path.insert( str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), ) -import snapform.blender.depth_plane as depth_plane - -from snapform.blender.depth_plane import ( +from snapform.operators.image import create_color_image +from snapform.operators.material import ( MOGE_DEPTH_PLANE_NODE_GROUP_NAME, MOGE_MATERIAL_NODE_GROUP_NAME, - build_moge2_depth_plane, - create_color_image, get_modifier_input, modifier_input_identifier, ) +from snapform.operators.depth_plane.build import build_moge2_depth_plane def validate_single_frame_source_types(): diff --git a/tools/validate_moge_surface.py b/tools/validate_moge_surface.py index ccde3c6..e4b2055 100644 --- a/tools/validate_moge_surface.py +++ b/tools/validate_moge_surface.py @@ -13,12 +13,11 @@ sys.path.insert( 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.depth_plane import material_color_space_name -from snapform.blender.surface import ( - MOGE_NODE_GROUP_NAME, - build_image_mesh, +from snapform.operators.material import ( + MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, + material_color_space_name, ) +from snapform.operators.surface import MOGE_NODE_GROUP_NAME, build_image_mesh from snapform.server.geometry import GeometryFrame, build_surface from snapform.server.outputs.surface import write_surface_archive