diff --git a/README.md b/README.md index 83a34b8..db7ed44 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 ## 它能做什么 -- 从 Image Empty 套索提取不规则区域，并通过 Mode 选择 Poisson 或 Teddy / Chordal Axis 生成封闭 Balloon 网格；源 Empty 保持不变 +- 从 Image Empty 套索提取不规则区域，通过 Poisson 或 Teddy / Chordal Axis 生成封闭 Balloon，并可用 BEN2 分析轮廓、MoGe-2 添加可调深度；源 Empty 保持不变 - 一键把 Image Empty 转换为立体面片：让 AI 估算深度，或直接用原图作为深度 - 把单张 Image Empty 转换为透视纹理网格，并自动匹配原对象的 Transform 与 Offset - 一键去除图片背景，生成透明前景图 @@ -33,7 +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 本地几何功能，在上述状态下仍可使用；只有开启 `AI Refine` 时才需要 Environment 与 BEN2。 +`Balloon Mesh (Lasso)` 默认是 Blender 本地几何功能，在上述状态下仍可使用；开启 `Contour` 或 `Depth` 后，才需要 Environment 与对应的 BEN2 或 MoGe-2 模型。 这些 UI 状态直接检查 Storage Root 的 `models/` 文件，不会仅因打开菜单或面板而启动 Job Server。 ## 快速上手 @@ -56,11 +56,16 @@ Environment 未安装时，同一位置显示两个安装按钮。Environment 按住左键沿目标轮廓绘制，松开后生成一个带原图 UV 的封闭网格。每次拖动都是一个独立完成的顶层 `Extract Balloon` Operator，工具本身继续保持激活，下一次拖动会创建另一个 Balloon；因此 `Ctrl+Z` 只撤销最近一次生成，`Ctrl+Shift+Z` 重做。每次生成后源 Image Empty 都保持选中和 Active。未绘制时没有持续 Modal 拦截事件，视图导航和其他原生快捷键由 Blender 正常处理；绘制中可用右键或 `Esc` 取消当前轨迹。选择左侧工具栏中的其他工具即可离开 Balloon Lasso。 -- `Mode`、`Thickness`、`Resolution (m)`、`Reuse Material` 和 `AI Refine` 显示在 Tool Header / Tool Settings +- `Mode`、`Thickness`、`Resolution (m)`、`Reuse Material`、`Contour`、`Depth` 和启用 `Depth` 后出现的 `Depth Model` 显示在 Tool Header / Tool Settings - `S`：调用 Blender 原生 Radial Control 调节 Thickness；通过 `zoom_path` 把默认 `1.0` 映射为约 100 px 的交互圆，软上限为 `5.0`，不设硬上限 - `D`：调用 Blender 原生 Radial Control 调节 Resolution；通过独立 `zoom_path` 把默认 `0.1 m` 映射为约 100 px 的交互圆，数值越小网格越密 -`Resolution (m)` 会根据 Scene Unit Scale 换算为网格使用的世界距离，因此不同单位设置下仍代表相同的物理采样间距。两种模式共用轮廓平滑、`.color` 图片副本和 `O Texture Layer` 材质。开启 `Reuse Material` 后，连续创建会按图片基名复用同一张 `.color` 与兼容的 Balloon 材质；关闭后每次新建。Balloon、Surface 与 Depth Plane 的材质名均取自图片数据名，而不是 Image Empty 对象名。`AI Refine` 默认关闭，此时 Balloon 完全在 Blender 本地生成，不需要 Environment；开启后会把套索外接范围裁切为图片，并将不规则选区外的 RGB 与 Alpha 全部清零后送入 BEN2。BEN2 Alpha 严格大于 `0.5` 的最大连通前景会成为新的套索轮廓，再生成 Balloon。当前只支持静态图片，不处理 Movie 或 Image Sequence。 +`Resolution (m)` 会根据 Scene Unit Scale 换算为网格使用的世界距离，因此不同单位设置下仍代表相同的物理采样间距。Poisson 模式使用尺度无关的求解判断，支持 `0.02 m` 等较密采样。两个 AI 选项默认关闭；`Contour` 使用 BEN2 Alpha 严格大于 `0.5` 的最大连通前景重建轮廓，`Depth` 使用所选 MoGe-2 模型。MoGe 的物理距离会先减去稳健拟合的基面，再换算为不受绝对拍摄距离影响的相对起伏；近侧与远侧信号都保留，不削顶、不截平。 + +Depth Balloon 使用局部 X 作为深度轴，并创建 ` Base Plane` Single Arrow Empty。Empty 的旋转定义基面朝向，沿自身局部 Z 移动则调整 middle level。默认深度校准为 `修正深度 = (0.5 × 相对基面深度 + 原始 Balloon 最大半厚度) × Depth Strength`；原始最大半厚度来自生成后的固定网格，不受修改器 `Thickness` 影响。这在 `Thickness = 1` 时等价于旧参数的 `Depth Strength = 0.5` 和两倍半厚度基面偏移。节点再以 Smoother Step 中心程度插值修正厚度与修正深度；`Thickness` 只控制厚度分支，`Depth Strength` 只控制深度分支。`Double Sided Depth` 默认关闭，使背面维持稳定厚度。 + +MoGe normal 会转换成适配 Balloon UV 的 tangent-space `normal.png`，由专用 `O Balloon Layer` 材质补回网格分辨率未能表达的细节；`Normal Strength` 默认 `0.5`，可在材质节点中调整。Balloon 材质通过 Custom Properties 直接保存源 Image datablock 引用，不依赖材质名或节点组名识别；勾选 `Reuse Material` 时优先复用同一 Image、同一材质类型的标记材质。Depth Balloon 共用一张原图尺寸的 normal，后续套索会把自己的有效区域合并进去，重叠区域采用最新结果。旧材质仍保留节点识别回退。当前只支持静态图片。 + ### 去除背景 选中 Image Empty 后，右键打开 `SnapForm` 子菜单并选择 `Remove Background`，完成后 Empty 会自动 diff --git a/docs/development.md b/docs/development.md index 24a016f..6be57c1 100644 --- a/docs/development.md +++ b/docs/development.md @@ -106,6 +106,7 @@ BLENDER_BIN="/c/Program Files/Blender Foundation/Blender 4.5/blender.exe" "$BLENDER_BIN" --background --python tools/build_geometry_nodes_asset.py "$BLENDER_BIN" --background --python tools/validate_depth_plane_asset.py "$BLENDER_BIN" --background --python tools/validate_geometry_nodes_asset.py +"$BLENDER_BIN" --background --python tools/validate_balloon_asset.py "$BLENDER_BIN" --background --python tools/validate_moge_depth_plane.py "$BLENDER_BIN" --background --python tools/validate_moge_surface.py ``` diff --git a/docs/features.md b/docs/features.md index 0bc02d7..4b25256 100644 --- a/docs/features.md +++ b/docs/features.md @@ -20,7 +20,11 @@ Depth Plane 与 Surface 可使用 DA3 或 MoGe-2。DA3 提供相对深度和相 转换任务完成后，SnapForm 在一个 Undo 步骤中创建网格、图片、材质和修改器。源 Image Empty 是否保留由转换选项决定。 -Balloon 默认在 Blender 内部生成，无需 Environment 或模型。它支持连续绘制，每次 Stroke 都生成一个可单独 Undo/Redo 的结果。Tool Settings 提供模式、厚度、采样距离、材质复用与 `AI Refine`。开启 `AI Refine` 后，Server 会裁切套索外接范围，把不规则选区外的 RGB 与 Alpha 清零，再用 BEN2 Alpha 严格大于 `0.5` 的最大连通前景重建轮廓；此模式需要 Environment 与 BEN2。 +Balloon 默认在 Blender 内部生成，无需 Environment 或模型。Tool Settings 提供模式、厚度、采样距离、材质复用、`Contour`、`Depth` 与 MoGe-2 `Depth Model`。Depth 会把物理距离相对化并稳健拟合初始基面，保留未截断的正负起伏。Base Plane 子 Empty 的旋转定义基面朝向，沿自身局部 Z 移动定义 middle level。 + +`O Balloon Shape` 计算 `修正厚度 = abs(原 X) × Thickness`，并将修正深度校准为 `(0.5 × 相对基面深度 + 原始 Balloon 最大半厚度) × Depth Strength`。固定的原始半厚度不乘 `Thickness`，因此 Thickness 与 Depth Strength 分别只控制各自分支。节点以归一化原 X 距离经 Smoother Step 得到中心程度，再插值厚度与深度。`Double Sided Depth` 默认关闭，背面保持 Balloon 厚度。MoGe normal 会输出为原图尺寸的 tangent-space 纹理，并通过 `O Balloon Layer` 的 `Normal Strength`（默认 `0.5`）补偿网格无法表达的细节。 + +新建 Balloon 材质会用 Custom Properties 保存源 Image datablock 和可选 normal Image 的直接引用。`Reuse Material` 优先按这些标记匹配，不依赖材质或节点组名称；普通与 Depth 材质分别复用。Depth Balloon 的共享 normal 会逐次合并新套索的有效区域，避免复用材质时丢失其他套索的法线。 ## 动画 diff --git a/docs/outputs.md b/docs/outputs.md index cd1c1ec..5fcf7cd 100644 --- a/docs/outputs.md +++ b/docs/outputs.md @@ -9,7 +9,7 @@ | `generate-depth` | `results.npz`、深度与颜色序列 | Depth Plane 所需深度、相机和贴图 | | `generate-surface` | `mesh.npz`、`normal.png`、颜色层 | Surface 网格、相机与材质输入 | | `remove-background` | `color.0001.png` 起 | 带 Alpha 的颜色层 | -| `refine-balloon` | `lasso-input.png`、`foreground.png`、`contour.json` | 套索约束输入、BEN2 前景与重建轮廓 | +| `analyze-balloon` | `lasso-input.png`，按选项输出 `foreground.png`、`contour.json`、`depth.npz`、`normal.png` | 套索约束输入、BEN2 轮廓与 MoGe-2 深度、法线 | | `debug-da3` | `prediction.npz` | DA3 API 返回 | | `debug-moge2` | `prediction*.npz` | MoGe-2 API 返回 | | `debug-ben2` | `foreground*.png` | BEN2 API 返回 | @@ -18,6 +18,8 @@ DA3 Depth 结果包含 `depth`、`conf`、`extrinsics` 与 `intrinsics`。MoGe-2 Depth 结果包含 `depth`、`normal`、`mask` 与 `intrinsics`。 +Balloon Depth 的 `depth.npz` 包含相对 `depth_signal`、归一化拟合平面 `baseline_vector`、`valid`、`reference_depth`、裁切原点与输入尺寸。Geometry Nodes 将归一化原 X 距离经 Smoother Step 后作为 Thickness 与 Depth 的混合系数。`normal.png` 是铺回原图尺寸的 tangent-space 法线；套索和有效区域外为中性法线。 + Surface 的 `mesh.npz` 包含 `vertices`、`faces`、`uv`、`intrinsics`、`image_size` 与 `reference_depth`，以及 Geometry Nodes 使用的网格属性。 Debug 文件保持模型 API 的数据形态，用于比较推理结果；正式 Job 则提供 Blender 导入所需的派生数据。 diff --git a/docs/server.md b/docs/server.md index 634862e..1f592c1 100644 --- a/docs/server.md +++ b/docs/server.md @@ -57,7 +57,7 @@ BlendJob 提供 Runtime、HTTP transport、队列、状态和取消协议。Snap Environment 使用 Python 3.12、ONNX Runtime、NumPy、Pillow、PyAV、trimesh、FastAPI 和 Uvicorn。Windows 使用 DirectML 或 CPU；其他平台根据支持情况使用 CoreML、OpenVINO 或 CPU。 -Balloon 的 `refine-balloon` Job 会在 BEN2 推理前裁切套索外接范围，并将不规则选区外的 RGBA 清零；推理后只回传 Alpha 大于 `0.5` 的最大连通轮廓，由 Blender 主线程生成网格。 +Balloon 的 `analyze-balloon` Job 会裁切套索输入并清零选区外 RGBA。`Contour` 用 BEN2 回传最大连通轮廓；`Depth` 用 MoGe-2 对有效相机距离稳健拟合二维基面，并把物理距离换算为不受绝对拍摄距离影响、未截断的相对深度信号。Job 同时把 MoGe normal 转换为 tangent-space 纹理，铺回原图尺寸并以中性法线填充无效区域。Blender 主线程负责采样深度属性、创建 Base Plane Empty、组装 `O Balloon Shape` 修改器与 `O Balloon Layer` 材质。 正式功能支持： diff --git a/src/snapform/operators/__init__.py b/src/snapform/operators/__init__.py index 298e639..e10c12d 100644 --- a/src/snapform/operators/__init__.py +++ b/src/snapform/operators/__init__.py @@ -2,7 +2,7 @@ from .background import ProcessImageBackground, RemoveImageBackground from .balloon import ( ActivateBalloonTool, ExtractBalloon, - RefineBalloon, + AnalyzeBalloon, ) from .debug import DebugBen2, DebugDa3, DebugMoge2 from .depth_plane import ( @@ -31,7 +31,7 @@ CLASSES = ( ExtractOriginalFrames, ClearModels, ExtractBalloon, - RefineBalloon, + AnalyzeBalloon, ActivateBalloonTool, RemoveImageBackground, ConvertToDepthPlane, diff --git a/src/snapform/operators/balloon/__init__.py b/src/snapform/operators/balloon/__init__.py index b60a402..30cf68e 100644 --- a/src/snapform/operators/balloon/__init__.py +++ b/src/snapform/operators/balloon/__init__.py @@ -2,5 +2,5 @@ from .operator import ( ActivateBalloonTool, BalloonTool, ExtractBalloon, - RefineBalloon, + AnalyzeBalloon, ) diff --git a/src/snapform/operators/balloon/depth.py b/src/snapform/operators/balloon/depth.py new file mode 100644 index 0000000..43e184b --- /dev/null +++ b/src/snapform/operators/balloon/depth.py @@ -0,0 +1,622 @@ +"""MoGe-2 depth sampling and adjustable Balloon deformation.""" + +from pathlib import Path + +import bpy + + +BALLOON_DEPTH_ATTRIBUTE = "o_balloon_depth" +BALLOON_SHAPE_NODE_GROUP_NAME = "O Balloon Shape" +BALLOON_SHAPE_NODE_GROUP_VERSION = 16 +BALLOON_SHAPE_VERSION_PROPERTY = "snapform_balloon_shape_version" + + +def load_balloon_depth_data(path): + import numpy as np + + try: + with np.load(Path(path), allow_pickle=False) as archive: + required = { + "depth_signal", + "baseline_vector", + "valid", + "origin", + "reference_depth", + "crop_size", + "source_size", + } + missing = required.difference(archive.files) + if missing: + raise ValueError( + f"Balloon depth data is missing: {', '.join(sorted(missing))}" + ) + data = { + "depth_signal": np.asarray( + archive["depth_signal"], + dtype=np.float32, + ), + "baseline_vector": np.asarray( + archive["baseline_vector"], + dtype=np.float32, + ), + "valid": np.asarray(archive["valid"], dtype=bool), + "origin": tuple(int(value) for value in archive["origin"]), + "reference_depth": float(archive["reference_depth"]), + "crop_size": tuple(int(value) for value in archive["crop_size"]), + "source_size": tuple(int(value) for value in archive["source_size"]), + } + except OSError as error: + raise ValueError(f"Unable to read Balloon depth data: {path}") from error + if data["depth_signal"].ndim != 2: + raise ValueError("Balloon depth signal must be two-dimensional") + if data["valid"].shape != data["depth_signal"].shape: + raise ValueError("Balloon depth mask dimensions do not match") + if data["baseline_vector"].shape != (3,): + raise ValueError("Balloon base plane vector must have three components") + if any(value <= 0 for value in (*data["crop_size"], *data["source_size"])): + raise ValueError("Balloon depth image dimensions must be positive") + return data + + +def _sample_valid_bilinear(values, valid, x, y): + import numpy as np + + height, width = values.shape + x = np.clip(x, 0.0, max(width - 1, 0)) + y = np.clip(y, 0.0, max(height - 1, 0)) + x0 = np.floor(x).astype(np.int64) + y0 = np.floor(y).astype(np.int64) + x1 = np.minimum(x0 + 1, width - 1) + y1 = np.minimum(y0 + 1, height - 1) + fx = x - x0 + fy = y - y0 + samples = np.stack( + ( + values[y0, x0], + values[y0, x1], + values[y1, x0], + values[y1, x1], + ), + axis=1, + ) + weights = np.stack( + ( + (1.0 - fx) * (1.0 - fy), + fx * (1.0 - fy), + (1.0 - fx) * fy, + fx * fy, + ), + axis=1, + ) + sample_valid = np.stack( + ( + valid[y0, x0], + valid[y0, x1], + valid[y1, x0], + valid[y1, x1], + ), + axis=1, + ) + weights *= sample_valid + totals = weights.sum(axis=1) + result = np.zeros(len(x), dtype=np.float32) + usable = totals > 1e-8 + result[usable] = (samples[usable] * weights[usable]).sum(axis=1) / totals[usable] + return result, usable + + +def balloon_depth_fields(vertices, vertex_uv, depth_data): + """Sample camera depth and convert the fitted base plane to local space.""" + import numpy as np + + vertices = np.asarray(vertices, dtype=np.float32) + uv = np.asarray(vertex_uv, dtype=np.float32) + if vertices.ndim != 2 or vertices.shape[1] != 3: + raise ValueError("Balloon vertices must have three components") + if uv.shape != (len(vertices), 2): + raise ValueError("Balloon UV coordinates do not match its vertices") + depth_signal = np.asarray(depth_data["depth_signal"], dtype=np.float32) + valid = np.asarray(depth_data["valid"], dtype=bool) + origin_x, origin_y = depth_data["origin"] + crop_width, crop_height = depth_data["crop_size"] + source_width, source_height = depth_data["source_size"] + depth_height, depth_width = depth_signal.shape + + crop_x = uv[:, 0] * source_width - origin_x + crop_y = (1.0 - uv[:, 1]) * source_height - origin_y + sample_x = crop_x * max(depth_width - 1, 0) / max(crop_width - 1, 1) + sample_y = crop_y * max(depth_height - 1, 0) / max(crop_height - 1, 1) + sampled, usable = _sample_valid_bilinear( + depth_signal, + valid, + sample_x, + sample_y, + ) + inside_crop = ( + (crop_x >= 0.0) + & (crop_x <= crop_width) + & (crop_y >= 0.0) + & (crop_y <= crop_height) + ) + usable &= inside_crop + sampled[~usable] = 0.0 + coordinates = np.stack( + ( + sample_x / max(depth_width - 1, 1) * 2.0 - 1.0, + sample_y / max(depth_height - 1, 1) * 2.0 - 1.0, + np.ones(len(vertices), dtype=np.float32), + ), + axis=1, + ).astype(np.float32) + + base_height = float(np.max(np.abs(vertices[:, 2]))) + if not base_height > 0.0: + return np.zeros(len(vertices), dtype=np.float32), np.zeros(3, dtype=np.float32) + depth_scale = base_height + normalized_plane = np.asarray( + depth_data["baseline_vector"], + dtype=np.float32, + ) + baseline_values = coordinates @ normalized_plane * depth_scale + local_design = np.stack( + ( + vertices[:, 0], + vertices[:, 1], + np.ones(len(vertices), dtype=np.float32), + ), + axis=1, + ) + if np.count_nonzero(usable) >= 3: + local_plane = np.linalg.lstsq( + local_design[usable], + baseline_values[usable], + rcond=None, + )[0].astype(np.float32) + else: + local_plane = np.zeros(3, dtype=np.float32) + depth_values = sampled * depth_scale + if usable.any(): + plane_without_offset = local_design[:, :2] @ local_plane[:2] + residual = depth_values - plane_without_offset + depth_values[usable] -= float(np.median(residual[usable])) + local_plane[2] = 0.0 + return depth_values, local_plane + + +def _assign_point_attribute(mesh, name, values): + import numpy as np + + values = np.asarray(values, dtype=np.float32) + if len(values) != len(mesh.vertices): + raise ValueError(f"Balloon attribute {name} does not match its vertices") + attribute = mesh.attributes.new(name=name, type="FLOAT", domain="POINT") + attribute.data.foreach_set("value", values) + + +def assign_balloon_shape_attributes(mesh, _vertices, depth_values=None): + if depth_values is not None: + _assign_point_attribute(mesh, BALLOON_DEPTH_ATTRIBUTE, depth_values) + + +def create_balloon_base_plane(context, balloon, vertices, local_plane): + import numpy as np + from mathutils import Vector + + vertices = np.asarray(vertices, dtype=np.float32) + slope_y, slope_z, offset = (float(value) for value in local_plane) + extent = vertices[:, 1:].max(axis=0) - vertices[:, 1:].min(axis=0) + normal = Vector((1.0, -slope_y, -slope_z)).normalized() + + controller = bpy.data.objects.new(f"{balloon.name} Base Plane", None) + context.collection.objects.link(controller) + controller.parent = balloon + controller.location = (0.0, 0.0, 0.0) + controller.rotation_euler = normal.to_track_quat("Z", "Y").to_euler() + controller.empty_display_type = "SINGLE_ARROW" + controller.empty_display_size = max(float(extent.max()) * 0.15, 0.01) + controller.show_in_front = True + controller.hide_render = True + controller["snapform_source_type"] = "BALLOON_BASE_PLANE" + controller["snapform_base_plane"] = (slope_y, slope_z, offset) + return controller + + +def _interface_socket(group, name, in_out, socket_type): + for item in group.interface.items_tree: + if item.item_type == "SOCKET" and item.name == name and item.in_out == in_out: + return item + return group.interface.new_socket( + name=name, + in_out=in_out, + socket_type=socket_type, + ) + + +def _remove_obsolete_interface_sockets(group): + obsolete = {"Base X", "Base Y", "Base Offset", "Far Clamp"} + for item in list(group.interface.items_tree): + if ( + item.item_type == "SOCKET" + and item.in_out == "INPUT" + and item.name in obsolete + ): + group.interface.remove(item) + + +def _float_input(group, name, default, minimum, maximum, description): + socket = _interface_socket(group, name, "INPUT", "NodeSocketFloat") + socket.default_value = default + socket.min_value = minimum + socket.max_value = maximum + socket.description = description + if hasattr(socket, "structure_type"): + socket.structure_type = "SINGLE" + return socket + + +def _bool_input(group, name, default, description): + socket = _interface_socket(group, name, "INPUT", "NodeSocketBool") + socket.default_value = default + socket.description = description + if hasattr(socket, "structure_type"): + socket.structure_type = "SINGLE" + return socket + + +def balloon_shape_node_group(): + group = bpy.data.node_groups.get(BALLOON_SHAPE_NODE_GROUP_NAME) + if ( + group is not None + and group.get(BALLOON_SHAPE_VERSION_PROPERTY) + == BALLOON_SHAPE_NODE_GROUP_VERSION + ): + return group + if group is None: + group = bpy.data.node_groups.new( + BALLOON_SHAPE_NODE_GROUP_NAME, + "GeometryNodeTree", + ) + else: + group.nodes.clear() + _remove_obsolete_interface_sockets(group) + + group.is_modifier = True + geometry_input = _interface_socket( + group, + "Geometry", + "INPUT", + "NodeSocketGeometry", + ) + if hasattr(geometry_input, "structure_type"): + geometry_input.structure_type = "SINGLE" + _float_input( + group, + "Thickness", + 1.0, + 0.0, + 10.0, + "Scale the symmetric base Balloon thickness along local X.", + ) + _float_input( + group, + "Depth Strength", + 1.0, + 0.0, + 10.0, + "Scale the signed MoGe-2 depth variation.", + ) + _bool_input( + group, + "Double Sided Depth", + False, + "Apply AI depth to both sides; keep disabled for a stable Balloon back.", + ) + base_plane_input = _interface_socket( + group, + "Base Plane", + "INPUT", + "NodeSocketObject", + ) + base_plane_input.description = ( + "Rotate the child Empty for middle-plane orientation; move it along local Z " + "to set the neutral depth level." + ) + _interface_socket(group, "Geometry", "OUTPUT", "NodeSocketGeometry") + + nodes = group.nodes + links = group.links + group_input = nodes.new("NodeGroupInput") + group_input.location = (-1500.0, 500.0) + position = nodes.new("GeometryNodeInputPosition") + position.location = (-1500.0, 200.0) + separate = nodes.new("ShaderNodeSeparateXYZ") + separate.location = (-1300.0, 200.0) + base_scale = nodes.new("ShaderNodeMath") + base_scale.operation = "MULTIPLY" + base_scale.location = (-1050.0, 350.0) + base_magnitude = nodes.new("ShaderNodeMath") + base_magnitude.operation = "ABSOLUTE" + base_magnitude.location = (-800.0, 350.0) + + side = nodes.new("ShaderNodeMath") + side.operation = "SIGN" + side.location = (-1050.0, 100.0) + front = nodes.new("ShaderNodeMath") + front.operation = "LESS_THAN" + front.inputs[1].default_value = 0.0 + front.location = (750.0, -350.0) + + depth_attribute = nodes.new("GeometryNodeInputNamedAttribute") + depth_attribute.data_type = "FLOAT" + depth_attribute.inputs["Name"].default_value = BALLOON_DEPTH_ATTRIBUTE + depth_attribute.location = (-1500.0, -150.0) + profile_factor = nodes.new("ShaderNodeMath") + profile_factor.operation = "ABSOLUTE" + profile_factor.location = (-800.0, 100.0) + profile_statistics = nodes.new("GeometryNodeAttributeStatistic") + profile_statistics.data_type = "FLOAT" + profile_statistics.domain = "POINT" + profile_statistics.location = (-600.0, 50.0) + profile_max_safe = nodes.new("ShaderNodeMath") + profile_max_safe.operation = "MAXIMUM" + profile_max_safe.inputs[1].default_value = 1e-6 + profile_max_safe.location = (-350.0, 100.0) + profile_normalized = nodes.new("ShaderNodeMath") + profile_normalized.operation = "DIVIDE" + profile_normalized.location = (-100.0, 100.0) + profile_smooth = nodes.new("ShaderNodeMapRange") + profile_smooth.data_type = "FLOAT" + profile_smooth.interpolation_type = "SMOOTHERSTEP" + profile_smooth.clamp = True + profile_smooth.inputs["From Min"].default_value = 0.0 + profile_smooth.inputs["From Max"].default_value = 1.0 + profile_smooth.inputs["To Min"].default_value = 0.0 + profile_smooth.inputs["To Max"].default_value = 1.0 + profile_smooth.location = (100.0, 50.0) + profile_active = nodes.new("ShaderNodeMath") + profile_active.operation = "MULTIPLY" + profile_active.location = (350.0, 50.0) + + object_info = nodes.new("GeometryNodeObjectInfo") + object_info.transform_space = "RELATIVE" + object_info.location = (-1500.0, -450.0) + plane_normal = nodes.new("FunctionNodeRotateVector") + plane_normal.inputs["Vector"].default_value = (0.0, 0.0, 1.0) + plane_normal.location = (-1250.0, -500.0) + effective_normal = nodes.new("GeometryNodeSwitch") + effective_normal.input_type = "VECTOR" + effective_normal.inputs["False"].default_value = (1.0, 0.0, 0.0) + effective_normal.location = (-1050.0, -500.0) + normal_separate = nodes.new("ShaderNodeSeparateXYZ") + normal_separate.location = (-1000.0, -550.0) + zero_position = nodes.new("ShaderNodeCombineXYZ") + zero_position.location = (-1050.0, -250.0) + plane_dot = nodes.new("ShaderNodeVectorMath") + plane_dot.operation = "DOT_PRODUCT" + plane_dot.location = (-800.0, -300.0) + plane_offset = nodes.new("ShaderNodeVectorMath") + plane_offset.operation = "DOT_PRODUCT" + plane_offset.location = (-800.0, -500.0) + plane_numerator = nodes.new("ShaderNodeMath") + plane_numerator.operation = "SUBTRACT" + plane_numerator.location = (-550.0, -300.0) + plane_divide = nodes.new("ShaderNodeMath") + plane_divide.operation = "DIVIDE" + plane_divide.location = (-300.0, -300.0) + + depth_residual = nodes.new("ShaderNodeMath") + depth_residual.operation = "ADD" + depth_residual.location = (-50.0, -200.0) + depth_calibration = nodes.new("ShaderNodeMath") + depth_calibration.operation = "MULTIPLY" + depth_calibration.inputs[1].default_value = 0.5 + depth_calibration.label = "Calibrate Relative Depth" + depth_calibration.location = (200.0, -200.0) + corrected_depth = nodes.new("ShaderNodeMath") + corrected_depth.operation = "ADD" + corrected_depth.location = (450.0, -200.0) + depth_scale = nodes.new("ShaderNodeMath") + depth_scale.operation = "MULTIPLY" + depth_scale.location = (700.0, -150.0) + depth_mix = nodes.new("ShaderNodeMix") + depth_mix.data_type = "FLOAT" + depth_mix.clamp_factor = True + depth_mix.location = (450.0, -50.0) + signed_depth = nodes.new("ShaderNodeMath") + signed_depth.operation = "MULTIPLY" + signed_depth.location = (850.0, 50.0) + front_choice = nodes.new("GeometryNodeSwitch") + front_choice.input_type = "FLOAT" + front_choice.location = (1100.0, 0.0) + double_choice = nodes.new("GeometryNodeSwitch") + double_choice.input_type = "FLOAT" + double_choice.location = (1350.0, 100.0) + combine = nodes.new("ShaderNodeCombineXYZ") + combine.location = (1600.0, 250.0) + set_position = nodes.new("GeometryNodeSetPosition") + set_position.location = (1850.0, 450.0) + group_output = nodes.new("NodeGroupOutput") + group_output.location = (2450.0, 450.0) + + base_frame = nodes.new("NodeFrame") + base_frame.label = "Base Thickness" + plane_frame = nodes.new("NodeFrame") + plane_frame.label = "Middle Plane" + mix_frame = nodes.new("NodeFrame") + mix_frame.label = "Depth Mix" + output_frame = nodes.new("NodeFrame") + output_frame.label = "Side Output" + for node in ( + separate, + base_scale, + base_magnitude, + profile_factor, + profile_statistics, + profile_max_safe, + profile_normalized, + profile_smooth, + ): + node.parent = base_frame + for node in ( + object_info, + plane_normal, + effective_normal, + normal_separate, + zero_position, + plane_dot, + plane_offset, + plane_numerator, + plane_divide, + ): + node.parent = plane_frame + for node in ( + depth_residual, + depth_calibration, + corrected_depth, + depth_scale, + profile_active, + depth_mix, + ): + node.parent = mix_frame + for node in ( + side, + front, + signed_depth, + front_choice, + double_choice, + combine, + set_position, + ): + node.parent = output_frame + + base_frame.location = (-1800.0, 500.0) + separate.location = (0.0, 200.0) + base_scale.location = (220.0, 350.0) + base_magnitude.location = (440.0, 350.0) + profile_factor.location = (220.0, 50.0) + profile_statistics.location = (440.0, 0.0) + profile_max_safe.location = (660.0, 0.0) + profile_normalized.location = (880.0, 0.0) + profile_smooth.location = (1100.0, 0.0) + + plane_frame.location = (-1800.0, -600.0) + object_info.location = (0.0, -300.0) + plane_normal.location = (220.0, -300.0) + effective_normal.location = (440.0, -300.0) + normal_separate.location = (660.0, -300.0) + zero_position.location = (220.0, 0.0) + plane_dot.location = (440.0, 0.0) + plane_offset.location = (660.0, -100.0) + plane_numerator.location = (880.0, 0.0) + plane_divide.location = (1100.0, 0.0) + + mix_frame.location = (-200.0, 200.0) + profile_active.location = (0.0, 300.0) + depth_residual.location = (0.0, 0.0) + depth_calibration.location = (220.0, 0.0) + corrected_depth.location = (440.0, 0.0) + depth_scale.location = (660.0, 0.0) + depth_mix.location = (880.0, 150.0) + + output_frame.location = (900.0, 200.0) + side.location = (0.0, 0.0) + front.location = (0.0, -250.0) + signed_depth.location = (220.0, 200.0) + front_choice.location = (440.0, 100.0) + double_choice.location = (660.0, 150.0) + combine.location = (880.0, 150.0) + set_position.location = (1100.0, 150.0) + group_input.location = (-2200.0, 500.0) + position.location = (-2200.0, 200.0) + depth_attribute.location = (-2200.0, -200.0) + group_output.location = (2500.0, 350.0) + + links.new(position.outputs["Position"], separate.inputs["Vector"]) + links.new(separate.outputs["X"], base_scale.inputs[0]) + links.new(group_input.outputs["Thickness"], base_scale.inputs[1]) + links.new(base_scale.outputs[0], base_magnitude.inputs[0]) + links.new(separate.outputs["X"], side.inputs[0]) + links.new(separate.outputs["X"], front.inputs[0]) + links.new(separate.outputs["X"], profile_factor.inputs[0]) + links.new(group_input.outputs["Geometry"], profile_statistics.inputs["Geometry"]) + links.new(profile_factor.outputs[0], profile_statistics.inputs["Attribute"]) + links.new(profile_statistics.outputs["Max"], profile_max_safe.inputs[0]) + links.new(profile_factor.outputs[0], profile_normalized.inputs[0]) + links.new(profile_max_safe.outputs[0], profile_normalized.inputs[1]) + links.new(profile_normalized.outputs[0], profile_smooth.inputs["Value"]) + links.new(profile_smooth.outputs["Result"], profile_active.inputs[0]) + links.new(depth_attribute.outputs["Exists"], profile_active.inputs[1]) + + links.new(group_input.outputs["Base Plane"], object_info.inputs["Object"]) + links.new(object_info.outputs["Rotation"], plane_normal.inputs["Rotation"]) + links.new(depth_attribute.outputs["Exists"], effective_normal.inputs["Switch"]) + links.new(plane_normal.outputs["Vector"], effective_normal.inputs["True"]) + links.new(effective_normal.outputs["Output"], normal_separate.inputs["Vector"]) + links.new(separate.outputs["Y"], zero_position.inputs["Y"]) + links.new(separate.outputs["Z"], zero_position.inputs["Z"]) + links.new(zero_position.outputs["Vector"], plane_dot.inputs[0]) + links.new(effective_normal.outputs["Output"], plane_dot.inputs[1]) + links.new(object_info.outputs["Location"], plane_offset.inputs[0]) + links.new(effective_normal.outputs["Output"], plane_offset.inputs[1]) + links.new(plane_offset.outputs["Value"], plane_numerator.inputs[0]) + links.new(plane_dot.outputs["Value"], plane_numerator.inputs[1]) + links.new(plane_numerator.outputs[0], plane_divide.inputs[0]) + links.new(normal_separate.outputs["X"], plane_divide.inputs[1]) + + links.new(depth_attribute.outputs["Attribute"], depth_residual.inputs[0]) + links.new(plane_divide.outputs[0], depth_residual.inputs[1]) + links.new(depth_residual.outputs[0], depth_calibration.inputs[0]) + links.new(depth_calibration.outputs[0], corrected_depth.inputs[0]) + links.new(profile_statistics.outputs["Max"], corrected_depth.inputs[1]) + links.new(corrected_depth.outputs[0], depth_scale.inputs[0]) + links.new(group_input.outputs["Depth Strength"], depth_scale.inputs[1]) + links.new(profile_active.outputs[0], depth_mix.inputs["Factor"]) + links.new(base_magnitude.outputs[0], depth_mix.inputs["A"]) + links.new(depth_scale.outputs[0], depth_mix.inputs["B"]) + links.new(side.outputs[0], signed_depth.inputs[0]) + links.new(depth_mix.outputs["Result"], signed_depth.inputs[1]) + links.new(front.outputs[0], front_choice.inputs["Switch"]) + links.new(base_scale.outputs[0], front_choice.inputs["False"]) + links.new(signed_depth.outputs[0], front_choice.inputs["True"]) + links.new( + group_input.outputs["Double Sided Depth"], + double_choice.inputs["Switch"], + ) + links.new(front_choice.outputs["Output"], double_choice.inputs["False"]) + links.new(signed_depth.outputs[0], double_choice.inputs["True"]) + links.new(double_choice.outputs["Output"], combine.inputs["X"]) + links.new(separate.outputs["Y"], combine.inputs["Y"]) + links.new(separate.outputs["Z"], combine.inputs["Z"]) + links.new(group_input.outputs["Geometry"], set_position.inputs["Geometry"]) + links.new(combine.outputs["Vector"], set_position.inputs["Position"]) + links.new(set_position.outputs["Geometry"], group_output.inputs["Geometry"]) + group.color_tag = "GEOMETRY" + group.description = ( + "Mix Balloon thickness with depth around an adjustable child Empty middle plane." + ) + group[BALLOON_SHAPE_VERSION_PROPERTY] = BALLOON_SHAPE_NODE_GROUP_VERSION + group.use_fake_user = True + return group + + +def add_balloon_shape_modifier( + balloon, + thickness, + base_plane=None, +): + from ..material import modifier_input_identifier, set_modifier_input + + node_group = balloon_shape_node_group() + modifier = balloon.modifiers.new(name=node_group.name, type="NODES") + modifier.node_group = node_group + modifier.show_group_selector = False + values = { + "Thickness": thickness, + "Depth Strength": 1.0, + "Double Sided Depth": False, + "Base Plane": base_plane, + } + for name, value in values.items(): + identifier = modifier_input_identifier(node_group, name) + set_modifier_input(modifier, identifier, value) + return modifier diff --git a/src/snapform/operators/balloon/operator.py b/src/snapform/operators/balloon/operator.py index 5a9c009..6770b38 100644 --- a/src/snapform/operators/balloon/operator.py +++ b/src/snapform/operators/balloon/operator.py @@ -3,10 +3,19 @@ import math import bpy +from .depth import ( + add_balloon_shape_modifier, + assign_balloon_shape_attributes, + balloon_depth_fields, + create_balloon_base_plane, + load_balloon_depth_data, +) from .geometry import INFLATION, POINT_SPACING, triangulate_balloon from ...properties import ( BEN2_MODEL_KEY, - balloon_ai_refine_property as ai_refine_property, + balloon_contour_analysis_property as contour_analysis_property, + balloon_depth_analysis_property as depth_analysis_property, + balloon_depth_model_property as depth_model_property, balloon_mode_property as mode_property, balloon_resolution_property as resolution_property, balloon_reuse_material_property as reuse_material_property, @@ -25,16 +34,24 @@ from ..image import ( color_data_name, is_image_empty, material_data_name, + normal_data_name, plane_bounds, prepare_image_input, remove_temporary_input, video_total_frames, ) -from ..material import TEXTURE_MATERIAL_NODE_GROUP_NAME, create_material +from ..material import ( + BALLOON_MATERIAL_NODE_GROUP_NAME, + TEXTURE_MATERIAL_NODE_GROUP_NAME, + create_material, +) from ..surface import assign_vertex_uv MIN_PATH_DISTANCE = 2.0 +BALLOON_MATERIAL_PROPERTY = "snapform_balloon_material" +BALLOON_SOURCE_IMAGE_PROPERTY = "snapform_source_image" +BALLOON_NORMAL_IMAGE_PROPERTY = "snapform_balloon_normal_image" def simplify_screen_path(path): @@ -150,7 +167,45 @@ def material_uses_balloon_texture(material, color_name): return has_texture_layer and has_color_image -def reusable_balloon_material(material_name, color_name): +def marked_balloon_material(source_image, normal_required=False): + return next( + ( + material + for material in bpy.data.materials + if callable(getattr(material, "get", None)) + and bool(material.get(BALLOON_MATERIAL_PROPERTY)) + and material.get(BALLOON_SOURCE_IMAGE_PROPERTY) is source_image + and ( + material.get(BALLOON_NORMAL_IMAGE_PROPERTY) is not None + ) + == normal_required + ), + None, + ) + + +def mark_balloon_material(material, source_image, normal_image=None): + material[BALLOON_MATERIAL_PROPERTY] = True + material[BALLOON_SOURCE_IMAGE_PROPERTY] = source_image + if normal_image is not None: + material[BALLOON_NORMAL_IMAGE_PROPERTY] = normal_image + + +def reusable_balloon_material( + material_name, + color_name, + source_image=None, + normal_required=False, +): + marked = ( + marked_balloon_material(source_image, normal_required) + if source_image is not None + else None + ) + if marked is not None: + return marked + if normal_required: + return None return next( ( material @@ -162,6 +217,35 @@ def reusable_balloon_material(material_name, color_name): ) +def load_balloon_normal_image(path, source_object): + image = bpy.data.images.load( + str(path), + check_existing=False, + ) + image.name = normal_data_name(source_object) + image.colorspace_settings.name = "Non-Color" + return image + + +def merge_balloon_normal_image(target, update): + import numpy as np + + if tuple(target.size) != tuple(update.size): + raise ValueError("Reusable Balloon normal image dimensions do not match") + target_pixels = np.empty(len(target.pixels), dtype=np.float32) + update_pixels = np.empty(len(update.pixels), dtype=np.float32) + target.pixels.foreach_get(target_pixels) + update.pixels.foreach_get(update_pixels) + target_rgba = target_pixels.reshape((-1, 4)) + update_rgba = update_pixels.reshape((-1, 4)) + neutral = np.asarray((128.0 / 255.0, 128.0 / 255.0, 1.0)) + changed = np.max(np.abs(update_rgba[:, :3] - neutral), axis=1) > (1.5 / 255.0) + target_rgba[changed] = update_rgba[changed] + target.pixels.foreach_set(target_pixels) + target.pack() + target.update() + + def reusable_color_image(color_name): return next( ( @@ -173,6 +257,27 @@ def reusable_color_image(color_name): ) +BALLOON_AXIS_TO_SOURCE = ( + (0.0, -1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (-1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), +) + +def orient_balloon_vertices(vertices): + """Rotate legacy XY/Z Balloon coordinates into a YZ/-X frame.""" + return [(-z, -x, y) for x, y, z in vertices] + + +def orient_balloon_world_matrix(source_matrix): + from mathutils import Matrix + + return source_matrix @ Matrix(BALLOON_AXIS_TO_SOURCE) + +def orient_balloon_base_plane(local_plane): + slope_x, slope_y, offset = local_plane + return (slope_x, -slope_y, -offset) + def extract_balloon_object( context, source_object, @@ -183,6 +288,8 @@ def extract_balloon_object( resolution=POINT_SPACING, select_result=True, reuse_material=True, + depth_data=None, + normal_path=None, ): image = source_object.data image_size = tuple(image.size) @@ -199,45 +306,87 @@ def extract_balloon_object( spacing=point_spacing, world_metric=world_metric, world_z_scale=world_z_scale, - inflation=inflation, + inflation=INFLATION, algorithm=algorithm, ) + depth_values = None + local_base_plane = (0.0, 0.0, 0.0) + if depth_data is not None: + depth_values, local_base_plane = balloon_depth_fields( + vertices, + vertex_uv, + depth_data, + ) + local_base_plane = orient_balloon_base_plane(local_base_plane) + vertices = orient_balloon_vertices(vertices) source_name = source_object.name balloon_name = source_name material_name = material_data_name(source_object) color_name = color_data_name(source_object) color_image = None + uses_normal = normal_path is not None material = ( - reusable_balloon_material(material_name, color_name) + reusable_balloon_material( + material_name, + color_name, + image, + normal_required=uses_normal, + ) if reuse_material else None ) + normal_image = None owns_color_image = False + owns_normal_image = False owns_material = False mesh = None balloon = None + base_plane = None try: + if material is not None and not uses_normal: + mark_balloon_material(material, image) + if material is not None and uses_normal: + normal_image = material.get(BALLOON_NORMAL_IMAGE_PROPERTY) + update_normal = load_balloon_normal_image(normal_path, source_object) + try: + merge_balloon_normal_image(normal_image, update_normal) + finally: + bpy.data.images.remove(update_normal) if material is None: - color_image = ( - reusable_color_image(color_name) - if reuse_material - else None - ) + color_image = reusable_color_image(color_name) if reuse_material else None if color_image is None: color_image = image.copy() color_image.name = color_name owns_color_image = True + if normal_path is not None: + normal_image = load_balloon_normal_image( + normal_path, + source_object, + ) + normal_image.pack() + owns_normal_image = True material = create_material( material_name, color_image, - 0.0, - surface_group_name=TEXTURE_MATERIAL_NODE_GROUP_NAME, + 1.0, + surface_group_name=( + BALLOON_MATERIAL_NODE_GROUP_NAME + if normal_image is not None + else TEXTURE_MATERIAL_NODE_GROUP_NAME + ), + normal_image=normal_image, scene=context.scene, ) + mark_balloon_material(material, image, normal_image) owns_material = True mesh = bpy.data.meshes.new(balloon_name) mesh.from_pydata(vertices, [], faces) + assign_balloon_shape_attributes( + mesh, + vertices, + depth_values, + ) mesh.update() assign_vertex_uv(mesh, vertex_uv) for face in mesh.polygons: @@ -245,26 +394,37 @@ def extract_balloon_object( balloon = bpy.data.objects.new(balloon_name, mesh) context.collection.objects.link(balloon) - balloon.matrix_world = matrix_world.copy() + balloon.matrix_world = orient_balloon_world_matrix(matrix_world) 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) + if depth_data is not None: + base_plane = create_balloon_base_plane( + context, + balloon, + vertices, + local_base_plane, + ) + add_balloon_shape_modifier(balloon, thickness, base_plane) + if depth_data is not None: + balloon["snapform_balloon_depth_analysis"] = True + balloon["snapform_reference_depth"] = depth_data["reference_depth"] except Exception: + if base_plane is not None: + bpy.data.objects.remove(base_plane, do_unlink=True) 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 owns_material and material is not None and material.users == 0: bpy.data.materials.remove(material) - if ( - owns_color_image - and color_image is not None - and color_image.users == 0 - ): + if owns_color_image and color_image is not None and color_image.users == 0: bpy.data.images.remove(color_image) + if owns_normal_image and normal_image is not None and normal_image.users == 0: + bpy.data.images.remove(normal_image) raise bpy.ops.object.select_all(action="DESELECT") @@ -323,7 +483,9 @@ class ExtractBalloon(bpy.types.Operator): resolution: resolution_property() source_object_name: hidden_string_property() reuse_material: reuse_material_property() - ai_refine: ai_refine_property() + contour_analysis: contour_analysis_property() + depth_analysis: depth_analysis_property() + depth_model: depth_model_property() source_matrix_data: hidden_string_property() selection_data: hidden_string_property() @@ -338,7 +500,10 @@ class ExtractBalloon(bpy.types.Operator): layout.prop(self, "thickness") layout.prop(self, "resolution", text="Resolution (m)") layout.prop(self, "reuse_material") - layout.prop(self, "ai_refine") + layout.prop(self, "contour_analysis") + layout.prop(self, "depth_analysis") + if self.depth_analysis: + layout.prop(self, "depth_model") @classmethod def poll(cls, context): @@ -349,9 +514,8 @@ class ExtractBalloon(bpy.types.Operator): ) def _draw_lasso(self): - if ( - not self._path - or not drawing_in_region(self._area_pointer, self._region_pointer) + if not self._path or not drawing_in_region( + self._area_pointer, self._region_pointer ): return vertices = list(self._path) @@ -377,7 +541,9 @@ class ExtractBalloon(bpy.types.Operator): self.thickness = settings.balloon_thickness self.resolution = settings.balloon_resolution self.reuse_material = settings.balloon_reuse_material - self.ai_refine = settings.balloon_ai_refine + self.contour_analysis = settings.balloon_contour_analysis + self.depth_analysis = settings.balloon_depth_analysis + self.depth_model = settings.balloon_depth_model self.source_object_name = source_object.name self.source_matrix_data = serialize_matrix(source_object.matrix_world) self.selection_data = "" @@ -449,14 +615,14 @@ class ExtractBalloon(bpy.types.Operator): polygon = json.loads(self.selection_data) if len(polygon) < 3: raise ValueError("The stored lasso region is invalid") - if self.ai_refine: + if self.contour_analysis or self.depth_analysis: if video_total_frames(source_object.data) > 1: - raise RuntimeError("AI Refine only supports a still image") + raise RuntimeError("AI Analysis only supports a still image") source_path, delete_input = prepare_image_input( source_object.data, context.scene, ) - result = bpy.ops.snapform.refine_balloon( + result = bpy.ops.snapform.analyze_balloon( input_path=str(source_path), delete_input=delete_input, source_object_name=self.source_object_name, @@ -466,6 +632,9 @@ class ExtractBalloon(bpy.types.Operator): thickness=self.thickness, resolution=self.resolution, reuse_material=self.reuse_material, + contour_analysis=self.contour_analysis, + depth_analysis=self.depth_analysis, + depth_model=self.depth_model, ) if "CANCELLED" in result: if delete_input: @@ -495,12 +664,11 @@ class ExtractBalloon(bpy.types.Operator): return {"FINISHED"} - -class RefineBalloon(JobOperatorBase, bpy.types.Operator): - bl_idname = "snapform.refine_balloon" - bl_label = "Refine Balloon" - job_type = "refine-balloon" - starting_message = "Starting Balloon AI refinement" +class AnalyzeBalloon(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.analyze_balloon" + bl_label = "Analyze Balloon" + job_type = "analyze-balloon" + starting_message = "Starting Balloon AI analysis" input_path: bpy.props.StringProperty(options=HIDDEN) delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) @@ -511,17 +679,35 @@ class RefineBalloon(JobOperatorBase, bpy.types.Operator): thickness: thickness_property() resolution: resolution_property() reuse_material: reuse_material_property() + contour_analysis: bpy.props.BoolProperty(default=False, options=HIDDEN) + depth_analysis: bpy.props.BoolProperty(default=False, options=HIDDEN) + depth_model: depth_model_property() + resolution_level: bpy.props.IntProperty( + default=5, + min=0, + max=9, + options=HIDDEN, + ) def request(self, _context): require_environment() - require_model(BEN2_MODEL_KEY) + if self.contour_analysis: + require_model(BEN2_MODEL_KEY) + if self.depth_analysis: + require_model(self.depth_model) + if not self.contour_analysis and not self.depth_analysis: + raise ValueError("Select Contour or Depth") polygon = json.loads(self.selection_data) if len(polygon) < 3: raise ValueError("The stored lasso region is invalid") return { "input": str(self.input_path), - "model": BEN2_MODEL_KEY, + "contour_analysis": self.contour_analysis, + "depth_analysis": self.depth_analysis, + "contour_model": BEN2_MODEL_KEY, + "depth_model": self.depth_model, "device": production_device().lower(), + "resolution_level": self.resolution_level, "polygon": polygon, } @@ -529,12 +715,17 @@ class RefineBalloon(JobOperatorBase, bpy.types.Operator): source_object = bpy.data.objects.get(self.source_object_name) if not is_image_empty(source_object): raise RuntimeError("The source Image Empty is no longer available") - payload = json.loads( - result.file("contour").read_text(encoding="utf-8") + polygon = json.loads(self.selection_data) + if self.contour_analysis: + payload = json.loads(result.file("contour").read_text(encoding="utf-8")) + polygon = payload.get("points", ()) + if len(polygon) < 3: + raise RuntimeError("BEN2 did not return a valid Balloon contour") + depth_data = ( + load_balloon_depth_data(result.file("depth")) + if self.depth_analysis + else None ) - polygon = payload.get("points", ()) - if len(polygon) < 3: - raise RuntimeError("BEN2 did not return a valid Balloon contour") matrix_world = deserialize_matrix(self.source_matrix_data) resolution_world = meters_to_world_distance( context.scene, @@ -550,15 +741,21 @@ class RefineBalloon(JobOperatorBase, bpy.types.Operator): resolution=resolution_world, select_result=False, reuse_material=self.reuse_material, + depth_data=depth_data, + normal_path=(result.file("normal") if self.depth_analysis else None), ) balloon["snapform_resolution_m"] = self.resolution - balloon["snapform_balloon_ai_refined"] = True + if self.contour_analysis: + balloon["snapform_balloon_contour_analysis"] = True + if self.depth_analysis: + balloon["snapform_balloon_depth_model"] = self.depth_model push_undo("Extract Balloon") - return "AI-refined Balloon is ready" + return "AI-analyzed Balloon is ready" def cleanup(self): cleanup_temporary_input(self.input_path, self.delete_input) + class ActivateBalloonTool(bpy.types.Operator): bl_idname = "snapform.activate_balloon_tool" bl_label = "Balloon Lasso Tool" @@ -610,7 +807,7 @@ class BalloonTool(bpy.types.WorkSpaceTool): ( "zoom_path", "scene.snapform_settings.balloon_thickness_radial_zoom", - ) + ), ] }, ), @@ -626,10 +823,11 @@ class BalloonTool(bpy.types.WorkSpaceTool): ( "zoom_path", "scene.snapform_settings.balloon_resolution_radial_zoom", - ) + ), ] }, - ), ) + ), + ) def draw_settings(context, layout, tool): settings = context.scene.snapform_settings @@ -637,4 +835,7 @@ class BalloonTool(bpy.types.WorkSpaceTool): layout.prop(settings, "balloon_thickness") layout.prop(settings, "balloon_resolution", text="Resolution (m)") layout.prop(settings, "balloon_reuse_material") - layout.prop(settings, "balloon_ai_refine") + layout.prop(settings, "balloon_contour_analysis") + layout.prop(settings, "balloon_depth_analysis") + if settings.balloon_depth_analysis: + layout.prop(settings, "balloon_depth_model", text="Depth Model") diff --git a/src/snapform/operators/balloon/poisson.py b/src/snapform/operators/balloon/poisson.py index 86381ca..5a19a03 100644 --- a/src/snapform/operators/balloon/poisson.py +++ b/src/snapform/operators/balloon/poisson.py @@ -46,7 +46,11 @@ def cotangent_poisson_field( edge_starts = [] edge_ends = [] weights = [] + # 约束剖分的瘦三角形可能产生负 cotangent 权重；钳制后保持 Dirichlet 系统正定。 for (start, end), weight in edge_weights.items(): + weight = max(float(weight), 0.0) + if weight <= 1e-14: + continue start_compact = compact[start] end_compact = compact[end] if start_compact >= 0: @@ -87,7 +91,7 @@ def cotangent_poisson_field( break multiplied = multiply(direction) denominator = float(direction.dot(multiplied)) - if denominator <= 1e-20 or not math.isfinite(denominator): + if denominator <= 0.0 or not math.isfinite(denominator): raise ValueError("Unable to solve the Poisson Balloon field") step = residual_product / denominator solution += step * direction diff --git a/src/snapform/operators/material.py b/src/snapform/operators/material.py index 8349299..af43f0a 100644 --- a/src/snapform/operators/material.py +++ b/src/snapform/operators/material.py @@ -17,6 +17,7 @@ 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" +BALLOON_MATERIAL_NODE_GROUP_NAME = "O Balloon 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" @@ -250,7 +251,11 @@ def material_node_group(depth_plane=False, surface_group_name=None): ("Alpha", "NodeSocketFloat"), ("Bump Scale", "NodeSocketFloat"), ] - if group_name == MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME: + normal_group = group_name in { + MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, + BALLOON_MATERIAL_NODE_GROUP_NAME, + } + if normal_group: input_sockets.insert(2, ("Normal", "NodeSocketColor")) input_sockets.insert(3, ("Normal Strength", "NodeSocketFloat")) if depth_plane: @@ -271,7 +276,9 @@ def material_node_group(depth_plane=False, surface_group_name=None): socket.min_value = -10000.0 socket.max_value = 10000.0 elif name == "Normal Strength": - socket.default_value = 1.0 + socket.default_value = ( + 0.5 if group_name == BALLOON_MATERIAL_NODE_GROUP_NAME else 1.0 + ) socket.min_value = 0.0 socket.max_value = 10.0 output_sockets = [ @@ -304,9 +311,13 @@ def material_node_group(depth_plane=False, surface_group_name=None): 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: + if normal_group: normal_map = nodes.new("ShaderNodeNormalMap") - normal_map.space = "OBJECT" + normal_map.space = ( + "TANGENT" + if group_name == BALLOON_MATERIAL_NODE_GROUP_NAME + else "OBJECT" + ) normal_map.location = (-500.0, -150.0) links.new(group_input.outputs["Normal"], normal_map.inputs["Color"]) links.new( diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 938be2b..7c0c491 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -93,10 +93,7 @@ def models_status(refresh=False): def _model_files_ready(model): directory = runtime.storage_root() / "models" / model.directory_name - return all( - any(directory.glob(pattern)) - for pattern in model.ready_patterns - ) + return all(any(directory.glob(pattern)) for pattern in model.ready_patterns) def model_ready(model_key): @@ -147,6 +144,15 @@ def moge2_model_enum_items(_owner, _context): ) +def balloon_depth_model_enum_items(_owner, _context): + models = [ + model + for model in model_catalog() + if model.get("family") == "moge2" and model.get("enabled") + ] + models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY) + return _model_items(models) + def depth_plane_model_enum_items(_owner, _context): models = [ model @@ -165,8 +171,7 @@ def surface_model_enum_items(_owner, _context): models = [ model for model in model_catalog() - if model.get("family") in {"moge2", "da3"} - and model.get("enabled") + if model.get("family") in {"moge2", "da3"} and model.get("enabled") ] models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY) return _model_items(models) @@ -228,17 +233,33 @@ def balloon_reuse_material_property(): ) -def balloon_ai_refine_property(): +def balloon_contour_analysis_property(): return bpy.props.BoolProperty( - name="AI Refine", + name="Contour", description=( - "Use BEN2 Alpha above 0.5 to refine the lasso before building " - "the Balloon" + "Use BEN2 Alpha above 0.5 to analyze the lasso contour before " + "building the Balloon" ), default=False, ) +def balloon_depth_analysis_property(): + return bpy.props.BoolProperty( + name="Depth", + description=("Use MoGe-2 to add adjustable image depth to the Balloon mesh"), + default=False, + ) + + +def balloon_depth_model_property(): + return bpy.props.EnumProperty( + name="Depth Model", + description="MoGe-2 model used for Balloon depth", + items=balloon_depth_model_enum_items, + default=0, + ) + def balloon_radial_zoom_property(scale): return bpy.props.FloatVectorProperty( size=2, @@ -247,12 +268,15 @@ def balloon_radial_zoom_property(scale): options={"HIDDEN", "SKIP_SAVE"}, ) + class SnapFormSettings(bpy.types.PropertyGroup): balloon_mode: balloon_mode_property() balloon_thickness: balloon_thickness_property() balloon_resolution: balloon_resolution_property() balloon_reuse_material: balloon_reuse_material_property() - balloon_ai_refine: balloon_ai_refine_property() + balloon_contour_analysis: balloon_contour_analysis_property() + balloon_depth_analysis: balloon_depth_analysis_property() + balloon_depth_model: balloon_depth_model_property() balloon_thickness_radial_zoom: balloon_radial_zoom_property(100.0) balloon_resolution_radial_zoom: balloon_radial_zoom_property(1000.0) input_path: bpy.props.StringProperty( diff --git a/src/snapform/server/__init__.py b/src/snapform/server/__init__.py index 42e1219..9c7970b 100644 --- a/src/snapform/server/__init__.py +++ b/src/snapform/server/__init__.py @@ -5,7 +5,7 @@ from blendjob import JobServer from .jobs import generate_depth as _generate_depth from .jobs import generate_surface as _generate_surface from .jobs.background import run as _remove_background -from .jobs.balloon import run as _refine_balloon +from .jobs.balloon import run as _analyze_balloon from .jobs.debug_ben2 import run as _debug_ben2 from .jobs.debug_da3 import run as _debug_da3 from .jobs.debug_moge2 import run as _debug_moge2 @@ -33,9 +33,9 @@ def download_required_models(context, parameters): return manager.download_missing(context, parameters["models"]) -@server.job("refine-balloon") -def refine_balloon(context, parameters): - return _refine_balloon(context, parameters) +@server.job("analyze-balloon") +def analyze_balloon(context, parameters): + return _analyze_balloon(context, parameters) @server.job("remove-background") diff --git a/src/snapform/server/jobs/balloon.py b/src/snapform/server/jobs/balloon.py index 4618603..a164453 100644 --- a/src/snapform/server/jobs/balloon.py +++ b/src/snapform/server/jobs/balloon.py @@ -3,6 +3,7 @@ import math import time from pathlib import Path +from ..models import moge2 from ..models.ben2 import infer_foreground @@ -43,7 +44,7 @@ def create_lasso_input(input_path, polygon, output_path): masked = Image.fromarray(pixels, mode="RGBA") output_path = Path(output_path) masked.save(output_path, format="PNG", compress_level=1) - return output_path, selection_mask, (left, top) + return output_path, selection_mask, (left, top), (width, height) def largest_component(mask): @@ -178,15 +179,165 @@ def contour_from_alpha(alpha, selection_mask, origin=(0, 0)): return [[x + offset_x, y + offset_y] for x, y in contour] +def resized_selection_mask(selection_mask, shape): + import numpy as np + from PIL import Image + + selection_mask = np.asarray(selection_mask, dtype=bool) + if selection_mask.shape == tuple(shape): + return selection_mask + resized = Image.fromarray(selection_mask.astype(np.uint8) * 255, mode="L") + resized = resized.resize( + (int(shape[1]), int(shape[0])), + Image.Resampling.NEAREST, + ) + return np.asarray(resized, dtype=np.uint8) > 0 + + +def normalized_moge_depth(depth, model_mask, selection_mask): + """Normalize MoGe depth while preserving its adjustable fitted base plane.""" + import numpy as np + + depth = np.asarray(depth, dtype=np.float32) + model_mask = np.asarray(model_mask, dtype=bool) + selection_mask = resized_selection_mask(selection_mask, depth.shape) + if model_mask.shape != depth.shape: + raise ValueError("MoGe-2 depth and mask dimensions do not match") + valid = selection_mask & model_mask & np.isfinite(depth) & (depth > 0.0) + if not valid.any(): + raise ValueError("MoGe-2 found no valid depth inside the lasso") + + height, width = depth.shape + y, x = np.indices(depth.shape, dtype=np.float64) + x = x / max(width - 1, 1) * 2.0 - 1.0 + y = y / max(height - 1, 1) * 2.0 - 1.0 + design = np.stack((x[valid], y[valid], np.ones(valid.sum())), axis=1) + selected = depth[valid].astype(np.float64) + weights = np.ones_like(selected) + coefficients = np.zeros(3, dtype=np.float64) + for _iteration in range(4): + weighted_design = design * np.sqrt(weights)[:, None] + weighted_depth = selected * np.sqrt(weights) + coefficients = np.linalg.lstsq( + weighted_design, + weighted_depth, + rcond=None, + )[0] + residual = selected - design @ coefficients + centered = residual - np.median(residual) + scale = max(float(np.median(np.abs(centered))) * 1.4826, 1e-6) + weights = np.minimum( + 1.0, + (2.5 * scale) / np.maximum(np.abs(centered), 1e-12), + ) + + baseline = coefficients[0] * x + coefficients[1] * y + coefficients[2] + reference = float(np.median(baseline[valid])) + residual = baseline[valid] - selected + near_residual = np.maximum(residual, 0.0) + near_total = float(np.sum(near_residual)) + relative_span = ( + float(np.sum(np.square(near_residual))) / near_total + if near_total > 1e-12 + else float(np.sqrt(np.mean(np.square(residual)))) + ) + relative_span = max(relative_span, 1e-6) + depth_signal = np.zeros(depth.shape, dtype=np.float32) + depth_signal[valid] = (reference - selected) / relative_span + baseline_vector = np.asarray( + ( + -coefficients[0] / relative_span, + -coefficients[1] / relative_span, + (reference - coefficients[2]) / relative_span, + ), + dtype=np.float32, + ) + return depth_signal, valid, reference, baseline_vector + + +def write_depth_analysis( + prediction, + selection_mask, + origin, + source_size, + output_path, +): + import numpy as np + + if prediction.depth.shape[0] != 1 or prediction.mask.shape[0] != 1: + raise ValueError("Balloon Depth requires one MoGe-2 frame") + depth_signal, valid, reference, baseline_vector = normalized_moge_depth( + prediction.depth[0], + prediction.mask[0], + selection_mask, + ) + crop_size = (selection_mask.shape[1], selection_mask.shape[0]) + np.savez_compressed( + output_path, + depth_signal=depth_signal, + baseline_vector=baseline_vector, + valid=valid, + reference_depth=np.asarray(reference, dtype=np.float32), + origin=np.asarray(origin, dtype=np.int32), + crop_size=np.asarray(crop_size, dtype=np.int32), + source_size=np.asarray(source_size, dtype=np.int32), + ) + return Path(output_path) + + +def write_balloon_normal_texture( + normal, + valid, + origin, + crop_size, + source_size, + output_path, +): + """Write MoGe camera normals as a full-image tangent-space texture.""" + import numpy as np + from PIL import Image + + normal = np.asarray(normal, dtype=np.float32) + valid = np.asarray(valid, dtype=bool) + if normal.ndim != 3 or normal.shape[2] != 3: + raise ValueError("Balloon normal data must have three components") + if valid.shape != normal.shape[:2]: + valid = resized_selection_mask(valid, normal.shape[:2]) + encoded = np.clip( + (normal * np.asarray((1.0, -1.0, -1.0), dtype=np.float32) + 1.0) + * 127.5, + 0.0, + 255.0, + ).astype(np.uint8) + encoded[~valid] = (128, 128, 255) + crop_width, crop_height = (int(value) for value in crop_size) + crop = Image.fromarray(encoded, mode="RGB") + if crop.size != (crop_width, crop_height): + crop = crop.resize( + (crop_width, crop_height), + Image.Resampling.BILINEAR, + ) + source_width, source_height = (int(value) for value in source_size) + full = Image.new("RGB", (source_width, source_height), (128, 128, 255)) + full.paste(crop, tuple(int(value) for value in origin)) + output_path = Path(output_path) + full.save(output_path) + return output_path + + def run(context, parameters): - """Refine a Balloon lasso with BEN2 and return its largest contour.""" + """Analyze a Balloon lasso with optional BEN2 contour and MoGe-2 depth.""" total_started = time.perf_counter() + contour_analysis = bool(parameters.get("contour_analysis")) + depth_analysis = bool(parameters.get("depth_analysis")) + if not contour_analysis and not depth_analysis: + raise ValueError("Select Contour or Depth") + context.progress(0.02, "Preparing lasso input") context.check_cancelled() - output_directory = context.directory output_directory.mkdir(parents=True, exist_ok=True) - input_path, selection_mask, origin = create_lasso_input( + input_path, selection_mask, origin, source_size = create_lasso_input( parameters["input"], parameters["polygon"], output_directory / "lasso-input.png", @@ -195,41 +346,83 @@ def run(context, parameters): device = parameters["device"] model_manager = context.resource("model_manager") - model_directory = model_manager.directory(parameters["model"]) - context.progress(0.18, f"Loading BEN2 for {device.upper()}") - foreground, load_ms, inference_ms = infer_foreground( - input_path, - model_directory, - device, - model_cache=model_manager, - cancel_check=context.check_cancelled, - ) - context.check_cancelled() + result = {} + if contour_analysis: + model_directory = model_manager.directory(parameters["contour_model"]) + context.progress(0.15, f"Loading BEN2 for {device.upper()}") + foreground, _load_ms, _inference_ms = infer_foreground( + input_path, + model_directory, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + context.progress( + 0.55 if depth_analysis else 0.9, + "Tracing BEN2 foreground", + ) + import numpy as np + from PIL import Image - context.progress(0.9, "Tracing BEN2 foreground") - import numpy as np - from PIL import Image + alpha = np.asarray(foreground.getchannel("A"), dtype=np.uint8) + contour = contour_from_alpha(alpha, selection_mask, origin) + gated_alpha = np.where(selection_mask, alpha, 0).astype(np.uint8) + foreground_pixels = foreground.convert("RGBA") + foreground_pixels.putalpha(Image.fromarray(gated_alpha, mode="L")) + foreground_pixels.save( + output_directory / "foreground.png", + format="PNG", + compress_level=1, + ) + contour_path = output_directory / "contour.json" + contour_path.write_text( + json.dumps({"points": contour}, separators=(",", ":")), + encoding="utf-8", + ) + result["contour"] = "contour.json" - alpha = np.asarray(foreground.getchannel("A"), dtype=np.uint8) - contour = contour_from_alpha(alpha, selection_mask, origin) - gated_alpha = np.where(selection_mask, alpha, 0).astype(np.uint8) - foreground_pixels = foreground.convert("RGBA") - foreground_pixels.putalpha(Image.fromarray(gated_alpha, mode="L")) - foreground_pixels.save( - output_directory / "foreground.png", - format="PNG", - compress_level=1, - ) - contour_path = output_directory / "contour.json" - contour_path.write_text( - json.dumps({"points": contour}, separators=(",", ":")), - encoding="utf-8", - ) + if depth_analysis: + context.progress( + 0.12 if not contour_analysis else 0.58, + "Estimating Balloon depth with MoGe-2", + ) + prediction = moge2.infer( + model_manager, + { + "model": parameters["depth_model"], + "device": device, + "resolution_level": parameters.get("resolution_level", 5), + }, + [input_path], + include_points=False, + progress_callback=None, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + context.progress(0.9, "Writing Balloon depth") + depth_path = write_depth_analysis( + prediction, + selection_mask, + origin, + source_size, + output_directory / "depth.npz", + ) + result["depth"] = depth_path.name + import numpy as np + + with np.load(depth_path, allow_pickle=False) as data: + normal_path = write_balloon_normal_texture( + prediction.normal[0], + data["valid"], + origin, + (selection_mask.shape[1], selection_mask.shape[0]), + source_size, + output_directory / "normal.png", + ) + result["normal"] = normal_path.name total_ms = (time.perf_counter() - total_started) * 1000.0 - message = ( - f"Balloon contour refined in {inference_ms:.0f} ms " - f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" - ) + message = f"Balloon AI analysis finished in {total_ms:.0f} ms" context.progress(1.0, message) - return {"contour": "contour.json"} + return result diff --git a/tests/test_addon.py b/tests/test_addon.py index 814518d..aed0861 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -82,34 +82,25 @@ class AddonTest(unittest.TestCase): 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.operators.material" - ) - cls.image_data = importlib.import_module( - "snapform.operators.image" - ) + cls.image_depth = importlib.import_module("snapform.operators.material") + cls.image_data = importlib.import_module("snapform.operators.image") cls.depth_operator = importlib.import_module( "snapform.operators.depth_plane.operator" ) cls.depth_build = importlib.import_module( "snapform.operators.depth_plane.build" ) - cls.image_background = importlib.import_module( - "snapform.operators.background" - ) - cls.image_mesh = importlib.import_module( - "snapform.operators.surface" - ) + cls.image_background = importlib.import_module("snapform.operators.background") + cls.image_mesh = importlib.import_module("snapform.operators.surface") cls.balloon = importlib.import_module("snapform.operators.balloon.operator") cls.balloon_geometry = importlib.import_module( "snapform.operators.balloon.geometry" ) + cls.balloon_depth = importlib.import_module("snapform.operators.balloon.depth") cls.balloon_poisson = importlib.import_module( "snapform.operators.balloon.poisson" ) - cls.balloon_teddy = importlib.import_module( - "snapform.operators.balloon.teddy" - ) + cls.balloon_teddy = importlib.import_module("snapform.operators.balloon.teddy") cls.paths_after_import = tuple(sys.path) cls.fake_bpy = fake_bpy @@ -298,7 +289,9 @@ class AddonTest(unittest.TestCase): window_manager = SimpleNamespace( progress_begin=lambda *_args: None, progress_update=lambda _value: None, - event_timer_add=lambda *_args, **_kwargs: timers.append(object()) or timers[-1], + event_timer_add=lambda *_args, **_kwargs: ( + timers.append(object()) or timers[-1] + ), event_timer_remove=lambda _timer: None, modal_handler_add=lambda _operator: None, ) @@ -320,7 +313,9 @@ class AddonTest(unittest.TestCase): with patch.object(runtime, "storage_root_factory", root): self.assertEqual(runtime.storage_root(), root.resolve()) self.assertEqual(runtime.environment_directory(), root / ".venv") - self.assertEqual(runtime.install_status_path(), root / "install-status.json") + self.assertEqual( + runtime.install_status_path(), root / "install-status.json" + ) self.assertEqual(runtime.install_log_path(), root / "install.log") self.assertEqual(runtime.server_log_path(), root / "server.log") @@ -328,8 +323,10 @@ class AddonTest(unittest.TestCase): runtime = self.addon.runtime with tempfile.TemporaryDirectory() as directory: root = Path(directory) - python = root / ".venv" / ( - "Scripts/python.exe" if sys.platform == "win32" else "bin/python" + python = ( + root + / ".venv" + / ("Scripts/python.exe" if sys.platform == "win32" else "bin/python") ) python.parent.mkdir(parents=True) python.touch() @@ -370,9 +367,7 @@ class AddonTest(unittest.TestCase): def test_snapform_post_install_is_ordinary_code_using_runtime_request(self): calls = [] fake_runtime = SimpleNamespace( - request=lambda job_type, parameters: calls.append( - (job_type, parameters) - ) + request=lambda job_type, parameters: calls.append((job_type, parameters)) ) self.job_runtime.post_install(fake_runtime) self.assertEqual( @@ -394,9 +389,7 @@ class AddonTest(unittest.TestCase): RuntimeError, "MoGe-2 ViT-S Normal: network unavailable; BEN2 Base: network unavailable", ): - self.job_runtime.post_install( - SimpleNamespace(request=request) - ) + self.job_runtime.post_install(SimpleNamespace(request=request)) self.assertEqual(calls, ["MOGE2_VITS_NORMAL", "BEN2_BASE"]) @@ -414,14 +407,18 @@ class AddonTest(unittest.TestCase): controller = runtime_module.EnvironmentController(runtime) controller.process = SimpleNamespace(poll=lambda: 0, returncode=0) - with tempfile.TemporaryDirectory() as directory, patch.object( - runtime, - "post_install", - post_install, - ), patch.object( - runtime, - "install_status_path", - return_value=Path(directory) / "missing.json", + with ( + tempfile.TemporaryDirectory() as directory, + patch.object( + runtime, + "post_install", + post_install, + ), + patch.object( + runtime, + "install_status_path", + return_value=Path(directory) / "missing.json", + ), ): self.assertEqual(controller.status("environment")["state"], "running") self.assertTrue(started.wait(1.0)) @@ -441,14 +438,18 @@ class AddonTest(unittest.TestCase): controller = runtime_module.EnvironmentController(runtime) controller.process = SimpleNamespace(poll=lambda: 0, returncode=0) - with tempfile.TemporaryDirectory() as directory, patch.object( - runtime, - "post_install", - post_install, - ), patch.object( - runtime, - "install_status_path", - return_value=Path(directory) / "missing.json", + with ( + tempfile.TemporaryDirectory() as directory, + patch.object( + runtime, + "post_install", + post_install, + ), + patch.object( + runtime, + "install_status_path", + return_value=Path(directory) / "missing.json", + ), ): controller.status("environment") controller.post_thread.join(1.0) @@ -477,17 +478,31 @@ class AddonTest(unittest.TestCase): def test_conversion_model_lists_default_to_moge2_vits(self): properties = importlib.import_module("snapform.properties") catalog = ( - {"key": "SMALL", "family": "da3", "enabled": True, - "label": "DA3 Small", "description": "DA3", "license": "Apache"}, - {"key": "MOGE2_VITS_NORMAL", "family": "moge2", "enabled": True, - "label": "MoGe-2 ViT-S", "description": "MoGe", "license": "MIT"}, + { + "key": "SMALL", + "family": "da3", + "enabled": True, + "label": "DA3 Small", + "description": "DA3", + "license": "Apache", + }, + { + "key": "MOGE2_VITS_NORMAL", + "family": "moge2", + "enabled": True, + "label": "MoGe-2 ViT-S", + "description": "MoGe", + "license": "MIT", + }, ) with patch.object(properties, "model_catalog", return_value=catalog): depth_items = properties.depth_plane_model_enum_items(None, None) surface_items = properties.surface_model_enum_items(None, None) + balloon_items = properties.balloon_depth_model_enum_items(None, None) self.assertEqual(depth_items[0][0], "MOGE2_VITS_NORMAL") self.assertEqual(surface_items[0][0], "MOGE2_VITS_NORMAL") + self.assertEqual([item[0] for item in balloon_items], ["MOGE2_VITS_NORMAL"]) def test_required_model_operator_requests_both_models(self): operator = self.addon.operators.DownloadRequiredModels() @@ -521,13 +536,16 @@ class AddonTest(unittest.TestCase): menu = self.addon.SnapFormImageMenu() menu.layout = Layout() - with patch.object( - self.addon.runtime, - "environment_ready", - return_value=True, - ), patch( - "snapform.properties.required_models_ready", - return_value=False, + with ( + patch.object( + self.addon.runtime, + "environment_ready", + return_value=True, + ), + patch( + "snapform.properties.required_models_ready", + return_value=False, + ), ): menu.draw(None) @@ -565,13 +583,16 @@ class AddonTest(unittest.TestCase): menu = self.addon.SnapFormImageMenu() menu.layout = Layout() - with patch.object( - self.addon.runtime, - "environment_ready", - return_value=False, - ), patch( - "snapform.properties.required_models_ready", - return_value=False, + with ( + patch.object( + self.addon.runtime, + "environment_ready", + return_value=False, + ), + patch( + "snapform.properties.required_models_ready", + return_value=False, + ), ): menu.draw(None) @@ -617,20 +638,18 @@ class AddonTest(unittest.TestCase): import numpy as np polygon = np.asarray( - ((100, 100), (900, 100), (950, 500), (800, 900), - (200, 900), (50, 500)), + ((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), - ) + 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 @@ -715,6 +734,66 @@ class AddonTest(unittest.TestCase): self.assertGreater(wide_height, narrow_height) + def test_poisson_balloon_supports_resolution_002(self): + import numpy as np + + spacing = 0.02 + size = 51 + axis = np.arange(size, dtype=float) * spacing + points = np.asarray( + [(x, y) for y in axis for x in axis], + dtype=float, + ) + faces = [] + for row in range(size - 1): + for column in range(size - 1): + bottom_left = row * size + column + bottom_right = bottom_left + 1 + top_left = bottom_left + size + top_right = top_left + 1 + faces.append((bottom_left, bottom_right, top_right)) + faces.append((bottom_left, top_right, top_left)) + boundary = np.zeros(len(points), dtype=bool) + boundary[:size] = True + boundary[-size:] = True + boundary[::size] = True + boundary[size - 1 :: size] = True + + heights = self.balloon_poisson.poisson_balloon_heights( + points, + faces, + boundary, + 0.45, + ) + + self.assertTrue(np.isfinite(heights).all()) + self.assertTrue((heights[boundary] == 0.0).all()) + self.assertGreater(heights[(size * size) // 2], 0.0) + + def test_poisson_solver_has_no_absolute_scale_cutoff(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, + ) + tiny = self.balloon_poisson.poisson_balloon_heights( + points * 1e-5, + faces, + boundary, + 0.45, + ) + + self.assertAlmostEqual(tiny[4], original[4] * 1e-5, places=12) + def test_teddy_balloon_uses_chordal_axis_radius(self): import numpy as np @@ -742,15 +821,18 @@ class AddonTest(unittest.TestCase): 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: + 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 ) @@ -762,6 +844,7 @@ class AddonTest(unittest.TestCase): 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 @@ -801,29 +884,204 @@ class AddonTest(unittest.TestCase): self.addon.CLASSES.index(self.balloon.ActivateBalloonTool), ) - def test_balloon_ai_refine_is_registered_and_forwarded_to_ben2(self): + def test_balloon_ai_analysis_is_registered_and_forwarded(self): import inspect - self.assertIn(self.balloon.RefineBalloon, self.addon.CLASSES) + self.assertIn(self.balloon.AnalyzeBalloon, self.addon.CLASSES) operator = self.balloon.ExtractBalloon - self.assertIn("ai_refine", operator.__annotations__) + self.assertIn("contour_analysis", operator.__annotations__) + self.assertIn("depth_analysis", operator.__annotations__) + self.assertIn("depth_model", operator.__annotations__) settings = self.addon.properties.SnapFormSettings.__annotations__ - self.assertIn("balloon_ai_refine", settings) - self.assertFalse(settings["balloon_ai_refine"]["default"]) + for name in ( + "balloon_contour_analysis", + "balloon_depth_analysis", + ): + self.assertIn(name, settings) + self.assertFalse(settings[name]["default"]) + self.assertIn("balloon_depth_model", settings) + self.assertEqual(settings["balloon_contour_analysis"]["name"], "Contour") + self.assertEqual(settings["balloon_depth_analysis"]["name"], "Depth") execute_source = inspect.getsource(operator.execute) invoke_source = inspect.getsource(operator.invoke) - request_source = inspect.getsource(self.balloon.RefineBalloon.request) - response_source = inspect.getsource(self.balloon.RefineBalloon.response) + request_source = inspect.getsource(self.balloon.AnalyzeBalloon.request) + response_source = inspect.getsource(self.balloon.AnalyzeBalloon.response) tool_source = inspect.getsource(self.balloon.BalloonTool.draw_settings) - self.assertIn("bpy.ops.snapform.refine_balloon", execute_source) + self.assertIn("bpy.ops.snapform.analyze_balloon", execute_source) self.assertIn("prepare_image_input(", execute_source) - self.assertIn("settings.balloon_ai_refine", invoke_source) + self.assertIn( + "settings.balloon_contour_analysis", + invoke_source, + ) + self.assertIn( + "settings.balloon_depth_analysis", + invoke_source, + ) + self.assertIn("self.depth_model = settings.balloon_depth_model", invoke_source) + self.assertIn("require_model(BEN2_MODEL_KEY)", request_source) + self.assertIn( + "require_model(self.depth_model)", + request_source, + ) self.assertIn('"polygon": polygon', request_source) self.assertIn('result.file("contour")', response_source) - self.assertIn("extract_balloon_object(", response_source) + self.assertIn('result.file("depth")', response_source) + self.assertIn("depth_data=depth_data", response_source) self.assertIn('push_undo("Extract Balloon")', response_source) - self.assertIn('"balloon_ai_refine"', tool_source) + self.assertIn('"balloon_contour_analysis"', tool_source) + self.assertIn('"balloon_depth_analysis"', tool_source) + self.assertIn('"balloon_depth_model"', tool_source) + + def test_balloon_analysis_requires_only_enabled_models(self): + operator = self.balloon.AnalyzeBalloon() + operator.input_path = "input.png" + operator.selection_data = "[[0,0],[2,0],[0,2]]" + operator.resolution_level = 5 + operator.depth_model = "MOGE2_VITB_NORMAL" + required = [] + + with ( + patch.object(self.balloon, "require_environment"), + patch.object( + self.balloon, + "require_model", + side_effect=required.append, + ), + patch.object( + self.balloon, + "production_device", + return_value="CUDA", + ), + ): + operator.contour_analysis = True + operator.depth_analysis = False + contour_request = operator.request(None) + operator.contour_analysis = False + operator.depth_analysis = True + depth_request = operator.request(None) + + self.assertEqual( + required, + [ + self.addon.properties.BEN2_MODEL_KEY, + "MOGE2_VITB_NORMAL", + ], + ) + self.assertTrue(contour_request["contour_analysis"]) + self.assertFalse(contour_request["depth_analysis"]) + self.assertFalse(depth_request["contour_analysis"]) + self.assertTrue(depth_request["depth_analysis"]) + self.assertEqual(depth_request["depth_model"], "MOGE2_VITB_NORMAL") + + def test_balloon_depth_fields_and_shape_nodes(self): + import inspect + import numpy as np + + depth_data = { + "depth_signal": np.asarray( + ((1.0, -2.0), (0.5, 0.0)), + dtype=np.float32, + ), + "baseline_vector": np.zeros(3, dtype=np.float32), + "valid": np.ones((2, 2), dtype=bool), + "origin": (0, 0), + "crop_size": (2, 2), + "source_size": (2, 2), + } + vertices = ( + (0.0, 0.0, 2.0), + (0.0, 0.0, 2.0), + (0.0, 0.0, -2.0), + (0.0, 0.0, -2.0), + ) + uv = ((0.0, 1.0), (1.0, 1.0), (0.0, 0.0), (1.0, 0.0)) + + values, local_plane = self.balloon_depth.balloon_depth_fields( + vertices, + uv, + depth_data, + ) + + np.testing.assert_allclose(values, (1.5, -4.5, 0.5, -0.5)) + self.assertAlmostEqual(float(np.median(values)), 0.0) + np.testing.assert_allclose(local_plane, (0.0, 0.0, 0.0)) + self.assertEqual( + self.balloon.orient_balloon_vertices(((1.0, 2.0, 3.0),)), + [(-3.0, -1.0, 2.0)], + ) + restored = np.asarray(self.balloon.BALLOON_AXIS_TO_SOURCE) @ np.asarray( + (-3.0, -1.0, 2.0, 1.0) + ) + np.testing.assert_allclose(restored, (1.0, 2.0, 3.0, 1.0)) + self.assertEqual( + self.balloon.orient_balloon_base_plane((0.2, 0.3, 0.4)), + (0.2, -0.3, -0.4), + ) + source = inspect.getsource(self.balloon_depth.balloon_shape_node_group) + controller_source = inspect.getsource( + self.balloon_depth.create_balloon_base_plane + ) + self.assertIn('"Thickness"', source) + self.assertIn('"Depth Strength"', source) + self.assertNotIn('"Far Clamp"', source) + self.assertIn('"Double Sided Depth"', source) + self.assertIn('"Base Plane"', source) + self.assertIn('nodes.new("GeometryNodeObjectInfo")', source) + self.assertIn("controller.parent = balloon", controller_source) + self.assertIn( + 'controller.empty_display_type = "SINGLE_ARROW"', controller_source + ) + self.assertIn('profile_factor.operation = "ABSOLUTE"', source) + self.assertIn('profile_statistics.data_type = "FLOAT"', source) + self.assertIn('profile_max_safe.operation = "MAXIMUM"', source) + self.assertIn('profile_normalized.operation = "DIVIDE"', source) + self.assertIn('profile_smooth.interpolation_type = "SMOOTHERSTEP"', source) + self.assertIn('profile_active.operation = "MULTIPLY"', source) + self.assertIn('effective_normal.input_type = "VECTOR"', source) + self.assertNotIn("o_balloon_depth_mix", source) + self.assertIn('front.operation = "LESS_THAN"', source) + self.assertIn('depth_residual.operation = "ADD"', source) + self.assertIn('depth_calibration.inputs[1].default_value = 0.5', source) + self.assertIn('corrected_depth.operation = "ADD"', source) + self.assertIn( + 'profile_statistics.outputs["Max"], corrected_depth.inputs[1]', + source, + ) + self.assertNotIn("baseline_thickness", source) + self.assertIn('separate.outputs["X"], base_scale.inputs[0]', source) + self.assertIn('double_choice.outputs["Output"], combine.inputs["X"]', source) + self.assertIn('depth_mix = nodes.new("ShaderNodeMix")', source) + self.assertIn('depth_mix.data_type = "FLOAT"', source) + self.assertIn('depth_mix.clamp_factor = True', source) + self.assertIn('base_magnitude.outputs[0], depth_mix.inputs["A"]', source) + self.assertIn('depth_scale.outputs[0], depth_mix.inputs["B"]', source) + self.assertNotIn("profile_inverse", source) + self.assertNotIn("thickness_component", source) + self.assertNotIn("depth_component", source) + self.assertNotIn("blended_magnitude", source) + self.assertIn('plane_offset.operation = "DOT_PRODUCT"', source) + self.assertIn("controller.location = (0.0, 0.0, 0.0)", controller_source) + self.assertNotIn("clamp_distance", source) + self.assertNotIn("safety_clamped", source) + self.assertNotIn("crossing_choice", source) + self.assertIn('front_choice.input_type = "FLOAT"', source) + self.assertIn('double_choice.input_type = "FLOAT"', source) + self.assertNotIn("double_inverse", source) + self.assertNotIn("depth_limited", source) + self.assertNotIn("side_gate", source) + self.assertIn('base_frame.label = "Base Thickness"', source) + self.assertIn('plane_frame.label = "Middle Plane"', source) + self.assertIn('mix_frame.label = "Depth Mix"', source) + self.assertIn('output_frame.label = "Side Output"', source) + self.assertIn('side.operation = "SIGN"', source) + self.assertIn("BALLOON_DEPTH_ATTRIBUTE", source) + self.assertEqual( + self.balloon_depth.BALLOON_DEPTH_ATTRIBUTE, + "o_balloon_depth", + ) + + def test_balloon_redo_properties_use_mode_and_meter_resolution(self): import inspect @@ -949,8 +1207,7 @@ class AddonTest(unittest.TestCase): import json matrix = tuple( - tuple(row * 4 + column for column in range(4)) - for row in range(4) + tuple(row * 4 + column for column in range(4)) for row in range(4) ) values = json.loads(self.balloon.serialize_matrix(matrix)) @@ -962,34 +1219,57 @@ class AddonTest(unittest.TestCase): 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( + "selected_object = balloon if select_result else source_object", source + ) + self.assertIn( + "balloon.matrix_world = orient_balloon_world_matrix(matrix_world)", + source, + ) self.assertIn("create_material(", source) self.assertIn("material_name = material_data_name(source_object)", source) self.assertIn( - "reusable_balloon_material(material_name, color_name)", + "normal_required=uses_normal", source, ) + self.assertIn("mark_balloon_material(material, image, normal_image)", source) + self.assertIn("merge_balloon_normal_image(normal_image, update_normal)", source) self.assertIn("if reuse_material", source) self.assertIn("color_image = image.copy()", source) self.assertIn("color_image.name = color_name", source) self.assertIn( - "material_name,\n color_image,\n 0.0,", + "material_name,\n color_image,\n 1.0,", source, ) self.assertIn( - "surface_group_name=TEXTURE_MATERIAL_NODE_GROUP_NAME", + "else TEXTURE_MATERIAL_NODE_GROUP_NAME", source, ) self.assertIn("owns_material", source) self.assertIn("owns_color_image", source) + self.assertIn("balloon_depth_fields(", source) + self.assertIn("assign_balloon_shape_attributes(", source) + self.assertIn("inflation=INFLATION", source) + self.assertIn("create_balloon_base_plane(", source) + self.assertIn( + "add_balloon_shape_modifier(balloon, thickness, base_plane)", source + ) + self.assertIn("BALLOON_MATERIAL_NODE_GROUP_NAME", source) + self.assertIn("normal_path=", source) self.assertEqual( self.image_depth.TEXTURE_MATERIAL_NODE_GROUP_NAME, "O Texture Layer", ) def test_balloon_material_reuse_requires_matching_image_and_node_group(self): + class MarkedMaterial(dict): + def __init__(self, name, **properties): + super().__init__(properties) + self.name = name + self.node_tree = None + color_image = SimpleNamespace(name="cow.color") + other_image = SimpleNamespace(name="other") texture_node = SimpleNamespace(type="TEX_IMAGE", image=color_image) texture_layer = SimpleNamespace( type="GROUP", @@ -999,6 +1279,23 @@ class AddonTest(unittest.TestCase): name="cow.001", node_tree=SimpleNamespace(nodes=[texture_node, texture_layer]), ) + marked = MarkedMaterial( + "unrelated-name", + snapform_balloon_material=True, + snapform_source_image=color_image, + ) + marked_other = MarkedMaterial( + "cow", + snapform_balloon_material=True, + snapform_source_image=other_image, + ) + normal_image = SimpleNamespace(name="cow.normal") + marked_normal = MarkedMaterial( + "depth-material", + snapform_balloon_material=True, + snapform_source_image=color_image, + snapform_balloon_normal_image=normal_image, + ) wrong_group = SimpleNamespace( name="cow", node_tree=SimpleNamespace( @@ -1013,14 +1310,37 @@ class AddonTest(unittest.TestCase): ) self.fake_bpy.data = SimpleNamespace( objects=SimpleNamespace(get=lambda _name: None), - materials=[wrong_group, compatible], + materials=[ + marked_other, + wrong_group, + compatible, + marked_normal, + marked, + ], images=[color_image], ) + self.assertIs( + self.balloon.reusable_balloon_material( + "cow", + "cow.color", + color_image, + ), + marked, + ) self.assertIs( self.balloon.reusable_balloon_material("cow", "cow.color"), compatible, ) + self.assertIs( + self.balloon.reusable_balloon_material( + "cow", + "cow.color", + color_image, + normal_required=True, + ), + marked_normal, + ) self.assertIs( self.balloon.reusable_color_image("cow.color"), color_image, @@ -1056,31 +1376,40 @@ class AddonTest(unittest.TestCase): (True, True, True), ) for environment_ready, models_ready, expected in states: - with self.subTest( - environment_ready=environment_ready, - models_ready=models_ready, - expected=expected, - ), patch.object( - self.addon.runtime, - "environment_ready", - return_value=environment_ready, - ), patch.object( - self.addon.runtime, - "server_busy", - return_value=False, - ), patch( - "snapform.properties.required_models_ready", - return_value=models_ready, + with ( + self.subTest( + environment_ready=environment_ready, + models_ready=models_ready, + expected=expected, + ), + patch.object( + self.addon.runtime, + "environment_ready", + return_value=environment_ready, + ), + patch.object( + self.addon.runtime, + "server_busy", + return_value=False, + ), + patch( + "snapform.properties.required_models_ready", + return_value=models_ready, + ), ): for module, operator in operators: - with self.subTest(operator=operator.__name__), patch.object( - module, - "is_image_empty", - return_value=True, - ), patch.object( - module, - "video_total_frames", - return_value=1, + with ( + self.subTest(operator=operator.__name__), + patch.object( + module, + "is_image_empty", + return_value=True, + ), + patch.object( + module, + "video_total_frames", + return_value=1, + ), ): self.assertEqual(operator.poll(context), expected) @@ -1095,18 +1424,20 @@ class AddonTest(unittest.TestCase): (moge / "model.onnx").touch() (ben2 / "model_fp16.onnx").touch() - with patch.object( - self.addon.runtime, - "storage_root", - return_value=root, - ), patch.object( - self.addon.runtime, - "resource", - side_effect=AssertionError("UI must not start the Server"), + with ( + patch.object( + self.addon.runtime, + "storage_root", + return_value=root, + ), + patch.object( + self.addon.runtime, + "resource", + side_effect=AssertionError("UI must not start the Server"), + ), ): self.assertTrue(properties.required_models_ready()) - def test_runtime_request_uses_request_response_shape(self): result = object() responses = [] @@ -1224,13 +1555,12 @@ class AddonTest(unittest.TestCase): operator.resolution_level = 6 operator.remove_background = False settings = SimpleNamespace(video_frames=12) - context = SimpleNamespace( - scene=SimpleNamespace(snapform_settings=settings) - ) + context = SimpleNamespace(scene=SimpleNamespace(snapform_settings=settings)) - with patch.object( - self.depth_operator, "require_environment" - ), patch.object(self.depth_operator, "require_model"): + with ( + patch.object(self.depth_operator, "require_environment"), + patch.object(self.depth_operator, "require_model"), + ): parameters = operator.request(context) self.assertEqual(parameters["input"], __file__) @@ -1247,9 +1577,7 @@ class AddonTest(unittest.TestCase): result = object() context = object() - with patch.object( - self.depth_operator, "build_moge2_depth_response" - ) as build: + with patch.object(self.depth_operator, "build_moge2_depth_response") as build: operator.response(context, result) build.assert_called_once_with( @@ -1280,8 +1608,9 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - generate_moge2_depth=lambda **options: calls.append(options) - or {"RUNNING_MODAL"} + generate_moge2_depth=lambda **options: ( + calls.append(options) or {"RUNNING_MODAL"} + ) ) ) @@ -1296,5 +1625,6 @@ class AddonTest(unittest.TestCase): self.assertEqual(calls[0]["source_object_name"], "Reference") self.assertFalse(any(key.startswith("snapform_") for key in calls[0])) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_server.py b/tests/test_server.py index aab8858..e882625 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -15,7 +15,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_server_module(): server_directory = PROJECT_ROOT / "src" / "snapform" / "server" specification = importlib.util.spec_from_file_location( - "server", server_directory / "__init__.py", + "server", + server_directory / "__init__.py", submodule_search_locations=[str(server_directory)], ) module = importlib.util.module_from_spec(specification) @@ -51,14 +52,17 @@ class ServerTest(unittest.TestCase): manager = self.server_module.ModelManager(Path(".")) downloads = [] context = SimpleNamespace(progress=lambda *_args: None) - with patch.object( - manager, - "ready", - side_effect=lambda key: key == "BEN2_BASE", - ), patch.object( - manager, - "_download", - side_effect=lambda _context, key, _progress: downloads.append(key), + with ( + patch.object( + manager, + "ready", + side_effect=lambda key: key == "BEN2_BASE", + ), + patch.object( + manager, + "_download", + side_effect=lambda _context, key, _progress: downloads.append(key), + ), ): result = manager.download_missing( context, @@ -215,7 +219,7 @@ class ServerTest(unittest.TestCase): { "download-model", "download-required-models", - "refine-balloon", + "analyze-balloon", "remove-background", "extract-frames", "generate-depth", @@ -254,15 +258,11 @@ class ServerTest(unittest.TestCase): def test_cancelling_job_keeps_server_and_model_cache_alive(self): fake_models = SimpleNamespace( - snapshot=lambda: { - "loaded": {"ben2": "BEN2_BASE", "da3": ""} - }, + snapshot=lambda: {"loaded": {"ben2": "BEN2_BASE", "da3": ""}}, clear=lambda: None, close=lambda: None, ) - with patch.object( - self.server_module, "ModelManager", return_value=fake_models - ): + with patch.object(self.server_module, "ModelManager", return_value=fake_models): server = self.create_server() started = threading.Event() @@ -287,9 +287,7 @@ class ServerTest(unittest.TestCase): ) def test_job_progress_is_kept_in_server_memory(self): - server = self.server_module.JobServer( - "Test", storage_root=self.storage_root - ) + server = self.server_module.JobServer("Test", storage_root=self.storage_root) @server.job("example") def example(context, _parameters): @@ -326,9 +324,7 @@ class ServerTest(unittest.TestCase): clear=lambda: events.append("clear"), close=lambda: events.append("close"), ) - server = self.server_module.JobServer( - "Test", storage_root=self.storage_root - ) + server = self.server_module.JobServer("Test", storage_root=self.storage_root) server.add_resource("model_manager", resource) @server.job("example") @@ -356,9 +352,7 @@ class ServerTest(unittest.TestCase): clear=lambda: None, close=lambda: None, ) - with patch.object( - self.server_module, "ModelManager", return_value=fake_models - ): + with patch.object(self.server_module, "ModelManager", return_value=fake_models): server = self.create_server() with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -376,9 +370,7 @@ class ServerTest(unittest.TestCase): job_id="job", ) status = self.wait_for_job(context) - self.assertTrue( - (context.directory / "color.0001.jpg").is_file() - ) + self.assertTrue((context.directory / "color.0001.jpg").is_file()) self.assertEqual(status["state"], "succeeded") @@ -395,7 +387,7 @@ class ServerTest(unittest.TestCase): pixels[1, 1, 3] = 64 Image.fromarray(pixels, mode="RGBA").save(source_path) - _path, _selection, origin = balloon.create_lasso_input( + _path, _selection, origin, source_size = balloon.create_lasso_input( source_path, ((1, 1), (5, 1), (1, 5)), output_path, @@ -403,6 +395,7 @@ class ServerTest(unittest.TestCase): result = np.asarray(Image.open(output_path).convert("RGBA")) self.assertEqual(origin, (1, 1)) + self.assertEqual(source_size, (6, 6)) self.assertEqual(result.shape, (4, 4, 4)) self.assertEqual(tuple(result[0, 0]), (20, 40, 60, 64)) self.assertEqual(tuple(result[3, 3]), (0, 0, 0, 0)) @@ -437,7 +430,9 @@ class ServerTest(unittest.TestCase): { "input": str(source_path), "polygon": ((1, 1), (5, 1), (1, 5)), - "model": "BEN2_BASE", + "contour_analysis": True, + "depth_analysis": False, + "contour_model": "BEN2_BASE", "device": "cpu", }, ) @@ -448,6 +443,148 @@ class ServerTest(unittest.TestCase): self.assertTrue((job_directory / "foreground.png").is_file()) self.assertTrue((job_directory / "contour.json").is_file()) + def test_balloon_depth_analysis_normalizes_moge_distance(self): + import numpy as np + + balloon = sys.modules["server.jobs.balloon"] + y, x = np.indices((7, 9), dtype=np.float32) + depth = 2.0 + x * 0.1 + y * 0.2 + depth[3, 4] -= 0.5 + depth[1, 1] += 3.0 + model_mask = np.ones_like(depth, dtype=bool) + selection_mask = np.ones_like(depth, dtype=bool) + selection_mask[0, 0] = False + + depth_signal, valid, reference, baseline_vector = balloon.normalized_moge_depth( + depth, + model_mask, + selection_mask, + ) + normalized_x = x / (depth.shape[1] - 1) * 2.0 - 1.0 + normalized_y = y / (depth.shape[0] - 1) * 2.0 - 1.0 + fitted_plane = ( + baseline_vector[0] * normalized_x + + baseline_vector[1] * normalized_y + + baseline_vector[2] + ) + displacement = depth_signal - fitted_plane + + self.assertFalse(valid[0, 0]) + self.assertEqual(depth_signal[0, 0], 0.0) + self.assertGreater(float(displacement[3, 4]), 1.0) + self.assertLess(float(displacement[3, 4]), 1.01) + self.assertLess(displacement[1, 1], -1.0) + self.assertGreater(reference, 2.0) + self.assertEqual(baseline_vector.shape, (3,)) + self.assertLess(float(displacement.min()), -1.0) + self.assertLess(float(displacement.max()), 1.01) + self.assertLess(float(np.percentile(displacement, 90.0)), 0.01) + + positive = displacement[valid & (displacement > 0.0)] + self.assertAlmostEqual( + float(np.sum(np.square(positive)) / np.sum(positive)), + 1.0, + places=5, + ) + + def test_balloon_relative_depth_ignores_absolute_distance_and_scale(self): + import numpy as np + + balloon = sys.modules["server.jobs.balloon"] + y, x = np.indices((8, 10), dtype=np.float32) + depth = 4.0 + x * 0.04 + y * 0.03 + depth[2:6, 3:7] -= np.asarray( + ( + (0.1, 0.2, 0.3, 0.1), + (0.2, 0.4, 0.8, 0.2), + (0.15, 0.3, 0.5, 0.15), + (0.05, 0.1, 0.2, 0.05), + ), + dtype=np.float32, + ) + mask = np.ones_like(depth, dtype=bool) + + original = balloon.normalized_moge_depth(depth, mask, mask) + farther = balloon.normalized_moge_depth(depth + 100.0, mask, mask) + rescaled = balloon.normalized_moge_depth(depth * 7.0, mask, mask) + + np.testing.assert_allclose(farther[0], original[0], atol=2e-4) + np.testing.assert_allclose(farther[3], original[3], atol=2e-4) + np.testing.assert_allclose(rescaled[0], original[0], atol=2e-4) + np.testing.assert_allclose(rescaled[3], original[3], atol=2e-4) + def test_balloon_depth_job_sends_masked_crop_only_to_moge(self): + import numpy as np + from PIL import Image + + balloon = sys.modules["server.jobs.balloon"] + source_path = self.storage_root / "depth-source.png" + pixels = np.full((6, 6, 4), (80, 100, 120, 255), dtype=np.uint8) + Image.fromarray(pixels, mode="RGBA").save(source_path) + job_directory = self.storage_root / "depth-job" + context = SimpleNamespace( + directory=job_directory, + progress=lambda *_args: None, + check_cancelled=lambda: None, + resource=lambda _name: SimpleNamespace(), + ) + received = [] + + def infer(_manager, _parameters, images, **_kwargs): + image = Image.open(images[0]).convert("RGBA") + received.append(np.asarray(image)) + height, width = image.size[1], image.size[0] + depth = np.linspace( + 1.0, + 2.0, + height * width, + dtype=np.float32, + ).reshape(height, width) + return SimpleNamespace( + depth=depth[None], + mask=np.ones((1, height, width), dtype=bool), + normal=np.broadcast_to( + np.asarray((0.0, 0.0, -1.0), dtype=np.float32), + (1, height, width, 3), + ).copy(), + ) + + with ( + patch.object(balloon.moge2, "infer", side_effect=infer), + patch.object( + balloon, + "infer_foreground", + side_effect=AssertionError("BEN2 should not run"), + ), + ): + result = balloon.run( + context, + { + "input": str(source_path), + "polygon": ((1, 1), (5, 1), (1, 5)), + "contour_analysis": False, + "depth_analysis": True, + "depth_model": "MOGE2_VITS_NORMAL", + "resolution_level": 5, + "device": "cpu", + }, + ) + + self.assertEqual( + result, + {"depth": "depth.npz", "normal": "normal.png"}, + ) + self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0)) + self.assertFalse((job_directory / "foreground.png").exists()) + with np.load(job_directory / "depth.npz", allow_pickle=False) as data: + self.assertEqual(data["depth_signal"].shape, (4, 4)) + self.assertEqual(data["baseline_vector"].shape, (3,)) + self.assertEqual(tuple(data["origin"]), (1, 1)) + self.assertEqual(tuple(data["source_size"]), (6, 6)) + self.assertTrue(data["valid"].any()) + normal = np.asarray(Image.open(job_directory / "normal.png").convert("RGB")) + self.assertEqual(normal.shape, (6, 6, 3)) + self.assertEqual(tuple(normal[0, 0]), (128, 128, 255)) + def test_balloon_contour_uses_alpha_above_half_and_largest_component(self): import numpy as np @@ -469,5 +606,6 @@ class ServerTest(unittest.TestCase): {(12, 21), (15, 21), (15, 24), (12, 24)}, ) + if __name__ == "__main__": unittest.main() diff --git a/tools/validate_balloon_asset.py b/tools/validate_balloon_asset.py new file mode 100644 index 0000000..875213f --- /dev/null +++ b/tools/validate_balloon_asset.py @@ -0,0 +1,136 @@ +from pathlib import Path +import sys + +import bpy + + +PROJECT_ROOT = Path(__file__).parents[1] +sys.path.insert(0, str(PROJECT_ROOT / "src")) +sys.path.insert( + 0, + str(next((PROJECT_ROOT / "src" / "snapform" / "wheels").glob("blendjob-*.whl"))), +) + +from snapform.operators.balloon.depth import ( + BALLOON_DEPTH_ATTRIBUTE, + add_balloon_shape_modifier, + balloon_shape_node_group, +) +from snapform.operators.balloon.operator import ( + mark_balloon_material, + merge_balloon_normal_image, + reusable_balloon_material, +) +from snapform.operators.material import ( + BALLOON_MATERIAL_NODE_GROUP_NAME, + material_node_group, + set_modifier_input, +) + + +def input_socket(group, name): + return next( + item + for item in group.interface.items_tree + if item.item_type == "SOCKET" + and item.in_out == "INPUT" + and item.name == name + ) + + +def main(): + bpy.ops.wm.read_factory_settings(use_empty=True) + group = balloon_shape_node_group() + assert group.is_modifier + assert input_socket(group, "Double Sided Depth").default_value is False + smoother = next( + node + for node in group.nodes + if node.bl_idname == "ShaderNodeMapRange" + and node.interpolation_type == "SMOOTHERSTEP" + ) + assert smoother.clamp + assert not any( + node.bl_idname == "ShaderNodeMath" + and node.operation == "MAXIMUM" + and node.label == "Clamp Depth" + for node in group.nodes + ) + + mesh = bpy.data.meshes.new("Balloon Validation Mesh") + mesh.from_pydata( + ( + (-1.0, 0.0, 0.0), + (-0.5, 0.0, 0.0), + (0.0, 0.0, 0.0), + (1.0, 0.0, 0.0), + ), + (), + (), + ) + depth = mesh.attributes.new(BALLOON_DEPTH_ATTRIBUTE, "FLOAT", "POINT") + depth.data.foreach_set("value", (0.0, 2.0, 2.0, 2.0)) + obj = bpy.data.objects.new("Balloon Validation", mesh) + bpy.context.scene.collection.objects.link(obj) + modifier = add_balloon_shape_modifier(obj, 1.0) + evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + result = evaluated.to_mesh() + try: + x = [vertex.co.x for vertex in result.vertices] + assert abs(x[0] + 1.0) < 1e-5 + assert abs(x[1] + 1.25) < 1e-5 + assert abs(x[2]) < 1e-5 + assert abs(x[3] - 1.0) < 1e-5 + finally: + evaluated.to_mesh_clear() + + set_modifier_input(modifier, input_socket(group, "Thickness").identifier, 2.0) + obj.update_tag(refresh={"DATA"}) + bpy.context.view_layer.update() + evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + result = evaluated.to_mesh() + try: + x = [vertex.co.x for vertex in result.vertices] + assert abs(x[0] + 1.0) < 1e-5 + assert abs(x[3] - 2.0) < 1e-5 + finally: + evaluated.to_mesh_clear() + + material_group = material_node_group( + surface_group_name=BALLOON_MATERIAL_NODE_GROUP_NAME, + ) + normal_map = next( + node + for node in material_group.nodes + if node.bl_idname == "ShaderNodeNormalMap" + ) + assert normal_map.space == "TANGENT" + assert abs(input_socket(material_group, "Normal Strength").default_value - 0.5) < 1e-6 + + source_image = bpy.data.images.new("Reuse Source", 2, 1) + shared_normal = bpy.data.images.new("Shared Normal", 2, 1) + shared_normal.pixels.foreach_set( + (128 / 255, 128 / 255, 1, 1, 128 / 255, 128 / 255, 1, 1) + ) + material = bpy.data.materials.new("Unrelated Material Name") + mark_balloon_material(material, source_image, shared_normal) + assert reusable_balloon_material( + "unused", + "unused", + source_image, + normal_required=True, + ) is material + + update_normal = bpy.data.images.new("Normal Update", 2, 1) + update_normal.pixels.foreach_set( + (1, 0.25, 0.75, 1, 128 / 255, 128 / 255, 1, 1) + ) + merge_balloon_normal_image(shared_normal, update_normal) + pixels = list(shared_normal.pixels) + assert abs(pixels[0] - 1.0) < 1e-5 + assert abs(pixels[4] - 128 / 255) < 1e-5 + print(f"Validated Blender {bpy.app.version_string}: {group.name}") + + +if __name__ == "__main__": + main()