From a0dd6e4a88b6a2c75c170fcb8578306033128566 Mon Sep 17 00:00:00 2001 From: nan Date: Sun, 13 Sep 2026 20:36:52 +0800 Subject: [PATCH] feat: Add restore cutout depth extension with edge depth correction - Introduced a new extension for restoring cutout depth with edge correction. - Implemented edge depth extension algorithm in `edge_depth.py` to enhance depth accuracy near cutout edges. - Updated `generate_moge_artifacts` to include an option for extending depth edges during depth generation. - Modified cutout job processing to enable edge depth extension based on depth request. - Added comprehensive tests for edge depth functionality, ensuring correct behavior under various scenarios. - Updated existing node loading functions to remove asset revision dependencies, simplifying the loading process. - Created new design and proposal documentation outlining the goals and decisions for the extension. --- .../.openspec.yaml | 2 + .../restore-cutout-depth-extension/design.md | 36 ++++ .../proposal.md | 23 +++ .../specs/cutout-depth-extension/spec.md | 49 ++++++ .../restore-cutout-depth-extension/tasks.md | 13 ++ src/anyimage/common/node.py | 8 +- .../operators/convert_to_panorama/object.py | 2 +- .../operators/convert_to_plane/object.py | 2 +- src/anyimage/operators/cutout_tool/object.py | 2 +- src/anyimage/server/geometry/edge_depth.py | 79 +++++++++ .../server/geometry/prediction_artifacts.py | 10 +- src/anyimage/server/jobs/cutout.py | 1 + .../operators/test_cutout_geometry.py | 3 +- .../operators/test_panorama_conversion.py | 9 +- .../operators/test_plane_conversion.py | 4 +- .../anyimage/operators/test_plane_objects.py | 6 +- tests/anyimage/server/test_edge_depth.py | 164 ++++++++++++++++++ .../server/test_geometry_artifacts.py | 1 + .../server/test_prediction_artifacts.py | 16 +- tools/nodes/arrangement/dimensions.json | 2 + tools/nodes/common/cutout_center.py | 78 +++++++++ tools/nodes/groups/image_cutout.py | 1 - tools/nodes/groups/image_depth_balloon.py | 1 - tools/nodes/groups/image_depth_panorama.py | 1 - tools/nodes/groups/image_depth_plane.py | 1 - 25 files changed, 482 insertions(+), 32 deletions(-) create mode 100644 openspec/changes/restore-cutout-depth-extension/.openspec.yaml create mode 100644 openspec/changes/restore-cutout-depth-extension/design.md create mode 100644 openspec/changes/restore-cutout-depth-extension/proposal.md create mode 100644 openspec/changes/restore-cutout-depth-extension/specs/cutout-depth-extension/spec.md create mode 100644 openspec/changes/restore-cutout-depth-extension/tasks.md create mode 100644 src/anyimage/server/geometry/edge_depth.py create mode 100644 tests/anyimage/server/test_edge_depth.py create mode 100644 tools/nodes/common/cutout_center.py diff --git a/openspec/changes/restore-cutout-depth-extension/.openspec.yaml b/openspec/changes/restore-cutout-depth-extension/.openspec.yaml new file mode 100644 index 0000000..c238415 --- /dev/null +++ b/openspec/changes/restore-cutout-depth-extension/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-09-13 diff --git a/openspec/changes/restore-cutout-depth-extension/design.md b/openspec/changes/restore-cutout-depth-extension/design.md new file mode 100644 index 0000000..5f147e8 --- /dev/null +++ b/openspec/changes/restore-cutout-depth-extension/design.md @@ -0,0 +1,36 @@ +## Context + +共享 MoGe 产物入口当前直接输出预测 XYZ。本变更在 Cutout 深度产物写入前加入边缘延伸，生成的贴图由现有几何节点直接消费。 + +独立实验位于 `C:/Users/icrdr/.codex/visualizations/2026/09/13/cutout-depth-4px/`，`report.md`、脚本、预测、对照图及 Blender 文件记录完整过程。使用 v26 节点组、BiRefNet Lite、MoGe-2 ViT-S Level 5，同一 Alpha 和预测，剪裁阈值 0.9，世界边长 0.015，Depth Split=0.1。蜘蛛/头骨补边约 35/26 ms，拖尾明显减少，细腿仍有断裂与折点。该节点版本仅是已有实验条件；实施验收使用现有节点和固定参数，仅替换深度产物。 + +## Goals + +以一次深度补边稳定 Cutout 轮廓采样，保持可靠内部、相机投影和产物有效性一致。 + +## Decisions + +### 沿用实测的固定来源规则 + +恢复 `31074bd^:src/anyimage/server/geometry/edge_depth.py` 的算法语义。剪裁阈值用于可见域及其 8 连通分量；可靠候选在该分量内满足 Alpha ≥ 0.95、Validity > 0.5、深度有限且为正。0.95 是固定来源条件，不新增随剪裁阈值变化的参数逻辑。低剪裁阈值可以保留毛发，但低 Alpha 像素仍从可靠来源取得深度。 + +不透明候选域按图片边缘延续计算距离，距离域外大于 2 个深度图像素的可靠候选作为供体。没有核心时回退到本分量未内缩的可靠候选；仍无候选则保留原预测。保持旧算法对内部区域的保护：模型 validity 孔洞只影响来源资格。 + +每个分量以最近供体 Z 延伸到边缘及可见域外 4 px。使用 sigma=0.8、半径3 px的 Gaussian，仅写回旧算法定义的补齐区；分量局部 padding 为7 px，重叠透明区域由最近可修复分量提供。完整不透明输入原样返回。 + +### 在 Cutout 产物边界启用 + +共享 `generate_moge_artifacts()` 接收默认关闭的内部补边选项，由 `server/jobs/cutout.py` 在请求深度时开启。模型只推理一次，保留原始 Frame，另取修复 Frame 写入 Cutout EXR 和 Metadata。Normal Writer 消费原始 Frame；Depth Plane 等调用保持默认行为。 + +修改 Z 后使用当前像素中心射线重建 XYZ，复用服务端已有反投影函数并处理像素/归一化内参。仅写回变化像素，保留原 validity；EXR A 继续为原图 Alpha × 原 validity，Metadata 参考深度按修复 Frame 和原参考掩码计算。复用已声明的 NumPy/SciPy，实施时核对独立服务端依赖。 + +## Risks / Trade-offs + +- 固定宽度对低分辨率、细腿覆盖比例较大 → 验证细分量回退及两张实图；维持已实测参数。 +- 高 Alpha 仍可能包含错误深度 → 保留内侧保护带；实图报告残余断裂和内部深度错误。 +- 补边有一次生成成本 → 分别记录补边耗时与后续节点求值耗时。 +- Normal 来自原始预测 → 对修复边缘检查材质着色，明确几何与法线差异。 + +## Migration Plan + +实施调整服务端深度处理与 Cutout 产物调用。新生成 Cutout 获得修复深度，已有对象可通过重新生成取得新产物。几何节点保持当前逻辑、参数和资产；相关测试与工作区外实图验收通过后交付。 diff --git a/openspec/changes/restore-cutout-depth-extension/proposal.md b/openspec/changes/restore-cutout-depth-extension/proposal.md new file mode 100644 index 0000000..fa0ca6b --- /dev/null +++ b/openspec/changes/restore-cutout-depth-extension/proposal.md @@ -0,0 +1,23 @@ +## Why + +Cutout 轮廓附近的背景深度会形成长条拖尾。蜘蛛和恐龙头骨实测表明，生成时延伸深度 4 像素可明显减少拖尾，适合在深度产物生成阶段恢复。 + +## What Changes + +- 仅为 Cutout 深度产物恢复已实测的旧版边缘延伸：可靠来源 Alpha ≥ 0.95、内侧保护 2 px、外侧延伸 4 px、局部平滑。 +- 保留旧算法的来源筛选、分量隔离与回退规则；剪裁 Alpha 继续定义可见轮廓，不新增来源阈值联动或设置。 +- 普通 Depth Plane、全景深度与 Normal 使用原始预测；Cutout 深度及 Metadata 使用同一修复结果。 + +## Capabilities + +### New Capabilities + +- `cutout-depth-extension`: Cutout 专用的可靠边缘深度延伸、产物隔离与实图验收。 + +### Modified Capabilities + +无；相关规范尚未归档至 `openspec/specs`。 + +## Impact + +涉及 Cutout Job、共享 MoGe 产物入口、服务端深度处理与相关测试。复用现有 NumPy/SciPy，核对服务端依赖声明。实现范围为生成的 Cutout 深度图及配套 Metadata；几何节点源码、参数、revision 和 `.blend` 资产保持现状。 diff --git a/openspec/changes/restore-cutout-depth-extension/specs/cutout-depth-extension/spec.md b/openspec/changes/restore-cutout-depth-extension/specs/cutout-depth-extension/spec.md new file mode 100644 index 0000000..1dfed62 --- /dev/null +++ b/openspec/changes/restore-cutout-depth-extension/specs/cutout-depth-extension/spec.md @@ -0,0 +1,49 @@ +## ADDED Requirements + +### Requirement: Reliable depth sources retain the tested alpha rule + +系统 SHALL 按剪裁 Alpha 定义可见域和8连通分量；供体候选 SHALL 位于同一可见分量内，并满足固定 Alpha ≥ 0.95、Validity > 0.5、有限正深度。系统 SHALL 保留内侧2像素保护及旧版回退规则。 + +#### Scenario: Low clipping threshold retains translucent edges +- **WHEN** 剪裁 Alpha 为0.1，边缘 Alpha 为0.1且同分量存在可靠核心 +- **THEN** 边缘属于可见域但不能充当供体，其深度使用同分量可靠来源补齐 + +#### Scenario: Narrow component lacks a core +- **WHEN** 分量内缩2像素后没有可靠供体 +- **THEN** 使用本分量未内缩的可靠候选；候选也为空时保留该分量原预测 + +### Requirement: Extension preserves reliable interior and projection + +系统 SHALL 沿用旧算法补齐轮廓边缘及可见域外4个深度图像素，使用sigma=0.8、半径3像素的分量内局部平滑。系统 SHALL 保留内部受保护区域及未修补像素，按目标像素射线重建变化像素的XYZ。 + +#### Scenario: Filtered sampling crosses an outline +- **WHEN** 采样邻域跨越主体轮廓 +- **THEN** 邻近透明侧4像素范围提供延伸深度，各分量采用自身供体和平滑输入，其他可见分量保持各自处理结果 + +#### Scenario: Corrected points retain their pixel location +- **WHEN** 像素继承供体Z +- **THEN** XYZ仍投影至目标像素，归一化投影误差小于1e-6；原Alpha和validity保持不变 + +#### Scenario: Opaque input and interior validity holes +- **WHEN** 输入完全不透明，或透明图内部存在模型validity孔洞 +- **THEN** 完全不透明输入原样返回；内部孔洞只限制供体资格，不扩大轮廓修补区域 + +### Requirement: Only cutout depth artifacts use corrected frames + +系统 SHALL 仅由Cutout深度请求开启延伸，使用同一修复Frame写入深度和Metadata。共享入口默认 SHALL 关闭修复，Normal SHALL 使用原始Frame，EXR A SHALL 保持原Alpha乘原validity。 + +#### Scenario: Cutout requests depth and normals +- **WHEN** 同一Cutout请求同时生成深度和Normal +- **THEN** 只执行一次模型预测，深度及Metadata来自修复Frame，Normal来自原始Frame + +#### Scenario: Other artifact requests +- **WHEN** 普通Depth Plane、全景深度或仅Normal请求运行 +- **THEN** 沿用各自原始预测产物行为 + +### Requirement: Real image verification uses a fixed prediction + +验收 SHALL 在工作区外使用用户蜘蛛和头骨原图、固定同一Alpha与预测，在现有几何节点及相同参数下仅替换原始/延伸深度产物，记录截图、耗时与残余问题。 + +#### Scenario: Full modifier evaluation +- **WHEN** 两张图完成深度产物对照 +- **THEN** 全部结果非空、坐标有限、无节点警告，报告一次补边与节点求值的独立耗时，并检查Normal着色 diff --git a/openspec/changes/restore-cutout-depth-extension/tasks.md b/openspec/changes/restore-cutout-depth-extension/tasks.md new file mode 100644 index 0000000..167800c --- /dev/null +++ b/openspec/changes/restore-cutout-depth-extension/tasks.md @@ -0,0 +1,13 @@ +## 1. Cutout 深度延伸 + +- [x] 1.1 恢复实测旧算法：固定0.95可靠来源条件、内侧2 px保护、外侧4 px延伸、分量内平滑和细分量回退；复用现有服务端反投影函数，核对NumPy/SciPy服务端声明。 +- [x] 1.2 为共享产物入口增加默认关闭的内部选项，仅由Cutout深度请求开启；深度和Metadata使用修复Frame，Normal与有效性通道保留原始语义。 +- [x] 1.3 补充相关行为测试：低剪裁Alpha不降低供体要求、分量隔离、无核心回退、贴图边界、完全不透明输入、内部validity孔洞、XYZ投影和未修改像素。 +- [x] 1.4 验证Cutout深度/Normal联合请求只推理一次，以及普通Depth Plane、全景和仅Normal路径的产物隔离。 + +## 2. 验收 + +- [x] 2.1 运行受影响的服务端与Cutout产物测试，确认进程正常退出。 +- [x] 2.2 在工作区外复用实测蜘蛛/头骨Alpha和预测，固定现有节点及参数，仅替换原始/延伸深度产物，记录同视角图、补边耗时、节点求值耗时与残余问题。 +- [x] 2.3 检查使用原始Normal的边缘着色，以及生产Cutout产物与独立实验的数值一致性。 +- [x] 2.4 检查最终差异与OpenSpec校验，确认改动限于深度产物生成及相关测试，汇总结果。 diff --git a/src/anyimage/common/node.py b/src/anyimage/common/node.py index fabe14f..d1ac329 100644 --- a/src/anyimage/common/node.py +++ b/src/anyimage/common/node.py @@ -10,12 +10,10 @@ def node_asset_path(): return Path(__file__).resolve().parents[1] / "assets" / NODE_ASSET_NAME -def load_node_group(group_name, *, asset_revision=None): +def load_node_group(group_name): existing = bpy.data.node_groups.get(group_name) if existing is not None: - if asset_revision is None or existing.get("anyimage_asset_revision") == asset_revision: - return existing - existing.name = group_name + " (Saved)" + return existing asset_path = node_asset_path() if not asset_path.is_file(): @@ -27,7 +25,7 @@ def load_node_group(group_name, *, asset_revision=None): directory=directory, filename=group_name, link=False, - do_reuse_local_id=True, + do_reuse_local_id=False, ) if "FINISHED" not in result: raise RuntimeError("Unable to append the node asset") diff --git a/src/anyimage/operators/convert_to_panorama/object.py b/src/anyimage/operators/convert_to_panorama/object.py index ee1ff01..6eba5ec 100644 --- a/src/anyimage/operators/convert_to_panorama/object.py +++ b/src/anyimage/operators/convert_to_panorama/object.py @@ -32,7 +32,7 @@ def create_panorama_from_result(context, result, operator): try: with material_color_image(input_path=operator.color_path) as color: metadata = load_panorama_metadata(result.file("depth_metadata")) - group = load_node_group("O Image Depth Panorama", asset_revision=15) + group = load_node_group("O Image Depth Panorama") depth = load_depth_result_image(result.file("depth"), source, metadata) material = create_image_material(source.data, color, shadeless=True, scene=context.scene, texture_extension="REPEAT") diff --git a/src/anyimage/operators/convert_to_plane/object.py b/src/anyimage/operators/convert_to_plane/object.py index 0a69116..05c56e0 100644 --- a/src/anyimage/operators/convert_to_plane/object.py +++ b/src/anyimage/operators/convert_to_plane/object.py @@ -62,7 +62,7 @@ def create_depth_plane_object( node_group = load_node_group({ "DEPTH": "O Image Depth Plane", "RELIEF": "O Image Relief Plane", - }[plane_type], asset_revision=16 if plane_type == "DEPTH" else None) + }[plane_type]) source_name = source_object.name bounds = image_empty_bounds(source_object) mesh = create_image_plane_mesh(source_name, bounds) diff --git a/src/anyimage/operators/cutout_tool/object.py b/src/anyimage/operators/cutout_tool/object.py index d42e1e9..a6f8487 100644 --- a/src/anyimage/operators/cutout_tool/object.py +++ b/src/anyimage/operators/cutout_tool/object.py @@ -72,7 +72,7 @@ def _assign_float_point_attribute(mesh, name, values): def load_cutout_node_group(shape): - return load_node_group(CUTOUT_NODE_GROUP_NAMES[shape], asset_revision=26 if shape == "DEPTH_SOLID" else 14 if shape == "DEPTH_BALLOON" else 13) + return load_node_group(CUTOUT_NODE_GROUP_NAMES[shape]) def _connect_cutout_normal_scale(material): diff --git a/src/anyimage/server/geometry/edge_depth.py b/src/anyimage/server/geometry/edge_depth.py new file mode 100644 index 0000000..c158973 --- /dev/null +++ b/src/anyimage/server/geometry/edge_depth.py @@ -0,0 +1,79 @@ +"""Extend reliable depth along each visible component's silhouette.""" + +from dataclasses import replace + + +def extend_edge_depth(frame, alpha, alpha_threshold=0.9): + """Return edge-corrected geometry while preserving reliable camera points.""" + import numpy as np + from scipy import ndimage + + from ..models.onnx_moge2 import depth_to_points + + depth = frame.depth + alpha = np.asarray(alpha, dtype=np.float32) + if alpha.shape != depth.shape: + raise ValueError("Edge Alpha must match the depth dimensions") + if np.all(alpha == 1.0): + return frame + + visible = (alpha > 0.0) & (alpha >= alpha_threshold) + opaque = visible & (alpha >= 0.95) + candidates = ( + opaque + & (frame.validity > 0.5) + & np.isfinite(depth) + & (depth > 0.0) + ) + if not candidates.any(): + return frame + interior_distance = ndimage.distance_transform_edt( + np.pad(opaque, 3, mode="edge") + )[3:-3, 3:-3] + interior = opaque & (interior_distance > 2.0) + core = candidates & interior + del interior_distance + labels, _ = ndimage.label(visible, structure=np.ones((3, 3), dtype=bool)) + corrected = depth.copy() + background_distance = np.full(depth.shape, np.inf, dtype=np.float32) + + for component_id, bounds in enumerate(ndimage.find_objects(labels), start=1): + region = tuple( + slice(max(0, bound.start - 7), min(size, bound.stop + 7)) + for bound, size in zip(bounds, depth.shape) + ) + component = labels[region] == component_id + donors = core[region] & component + if not donors.any(): + donors = candidates[region] & component + if not donors.any(): + continue + + indices = ndimage.distance_transform_edt( + ~donors, return_distances=False, return_indices=True + ) + extended = depth[region][tuple(indices)] + smoothed = ndimage.gaussian_filter(extended, sigma=0.8, radius=3) + distance = ndimage.distance_transform_edt(~component) + background = ( + ~visible[region] + & (distance <= 4.0) + & (distance < background_distance[region]) + ) + target = (component & ~interior[region] & ~donors) | background + corrected[region][target] = smoothed[target] + background_distance[region][background] = distance[background] + + changed = corrected != depth + if not changed.any(): + return frame + points = frame.points + if points is not None: + height, width = depth.shape + intrinsics = frame.intrinsics.copy() + intrinsics[0] /= width + intrinsics[1] /= height + projected = depth_to_points(corrected, intrinsics) + points = points.copy() + points[changed] = projected[changed] + return replace(frame, depth=corrected, points=points) diff --git a/src/anyimage/server/geometry/prediction_artifacts.py b/src/anyimage/server/geometry/prediction_artifacts.py index 6e1ff78..4272034 100644 --- a/src/anyimage/server/geometry/prediction_artifacts.py +++ b/src/anyimage/server/geometry/prediction_artifacts.py @@ -8,6 +8,7 @@ from .depth_texture import ( write_depth_texture, ) from .normal_texture import write_camera_normal_texture +from .edge_depth import extend_edge_depth def _limited_model_input(source_path, output_path, max_input_size): @@ -53,6 +54,7 @@ def generate_moge_artifacts( generate_depth=False, normal_mode="NONE", alpha_threshold=0.9, + extend_depth_edges=False, ): """Run MoGe once and write only the requested file artifacts.""" normal_mode = str(normal_mode).upper() @@ -75,11 +77,15 @@ def generate_moge_artifacts( if generate_depth: context.progress(0.45, "Generating depth texture") alpha = _image_alpha(input_path, frame.depth.shape) + depth_frame = ( + extend_edge_depth(frame, alpha, alpha_threshold) + if extend_depth_edges else frame + ) result["depth"] = write_depth_texture( - frame, context.directory / "depth.exr", alpha=alpha, + depth_frame, context.directory / "depth.exr", alpha=alpha, ) result["depth_metadata"] = write_depth_metadata( - frame, + depth_frame, (alpha > 0.0) & (alpha >= alpha_threshold), context.directory / "depth.json", ) diff --git a/src/anyimage/server/jobs/cutout.py b/src/anyimage/server/jobs/cutout.py index 607370a..52d4f9d 100644 --- a/src/anyimage/server/jobs/cutout.py +++ b/src/anyimage/server/jobs/cutout.py @@ -18,6 +18,7 @@ def run(context, parameters): parameters, input_path, generate_depth=generate_depth, + extend_depth_edges=generate_depth, normal_mode=normal_mode, alpha_threshold=float(parameters.get("alpha_threshold", 0.9)), ) diff --git a/tests/anyimage/operators/test_cutout_geometry.py b/tests/anyimage/operators/test_cutout_geometry.py index 0dc9a70..83e4377 100644 --- a/tests/anyimage/operators/test_cutout_geometry.py +++ b/tests/anyimage/operators/test_cutout_geometry.py @@ -40,8 +40,7 @@ def test_cutout_node_groups_are_loaded_from_the_asset(): ) as load: for shape in cutout_shape.CUTOUT_NODE_GROUP_NAMES: assert cutout_object.load_cutout_node_group(shape) is loaded - load.assert_called_with(cutout_shape.CUTOUT_NODE_GROUP_NAMES[shape], - asset_revision=26 if shape == "DEPTH_SOLID" else 14 if shape == "DEPTH_BALLOON" else 13) + load.assert_called_with(cutout_shape.CUTOUT_NODE_GROUP_NAMES[shape]) assert load.call_count == 4 diff --git a/tests/anyimage/operators/test_panorama_conversion.py b/tests/anyimage/operators/test_panorama_conversion.py index 5aa7f97..2925d67 100644 --- a/tests/anyimage/operators/test_panorama_conversion.py +++ b/tests/anyimage/operators/test_panorama_conversion.py @@ -81,17 +81,16 @@ def artifacts(tmp_path, source): def test_result_replaces_source_with_zero_rotation_and_shadeless_material(source, artifacts, monkeypatch): result, operator = artifacts - saved = bpy.data.node_groups.new("O Image Depth Panorama", "GeometryNodeTree") - saved["anyimage_asset_revision"] = 11 + from anyimage.common.node import load_node_group + saved = load_node_group("O Image Depth Panorama") matrix = Matrix.Translation((2, -3, 4)) @ Matrix.Rotation(0.6, 4, "Y") @ Matrix.Diagonal((2, 3, 4, 1)) source.matrix_world = matrix source.empty_image_offset = (0.3, -0.8) undo = Mock(return_value={"FINISHED"}) monkeypatch.setattr(conversion, "bpy", SimpleNamespace(data=bpy.data, ops=SimpleNamespace(ed=SimpleNamespace(undo_push=undo)))) obj = conversion.create_panorama_from_result(bpy.context, result, operator) - assert obj.modifiers[0].node_group != saved - assert obj.modifiers[0].node_group["anyimage_asset_revision"] == 15 - assert saved in bpy.data.node_groups.values() and saved["anyimage_asset_revision"] == 11 + assert obj.modifiers[0].node_group == saved + assert saved.name == "O Image Depth Panorama" assert obj.name == operator.source_object_name and obj.type == "MESH" np.testing.assert_allclose(obj.location, matrix.translation, atol=1e-6) np.testing.assert_allclose(obj.scale, matrix.to_scale(), atol=1e-6) diff --git a/tests/anyimage/operators/test_plane_conversion.py b/tests/anyimage/operators/test_plane_conversion.py index c89c69c..5aae35f 100644 --- a/tests/anyimage/operators/test_plane_conversion.py +++ b/tests/anyimage/operators/test_plane_conversion.py @@ -11,7 +11,7 @@ def test_conversion_preserves_offset_world_placement_and_initializes_inputs(surf obj, _, _, inputs, image = surface group = obj.modifiers[0].node_group - monkeypatch.setattr(conversion, "load_node_group", lambda name, *, asset_revision: group) + monkeypatch.setattr(conversion, "load_node_group", lambda name: group) source = bpy.data.objects.new("Offset Image", None) source.empty_display_type = "IMAGE" source.data = image @@ -48,7 +48,7 @@ def test_failed_modifier_initialization_removes_partial_object(surface, monkeypa from anyimage.operators.convert_to_plane import object as conversion obj, _, _, _, image = surface - monkeypatch.setattr(conversion, "load_node_group", lambda name, *, asset_revision: obj.modifiers[0].node_group) + monkeypatch.setattr(conversion, "load_node_group", lambda name: obj.modifiers[0].node_group) source = bpy.data.objects.new("Source", None) source.empty_display_type = "IMAGE" source.data = image diff --git a/tests/anyimage/operators/test_plane_objects.py b/tests/anyimage/operators/test_plane_objects.py index 6fd43f3..af87370 100644 --- a/tests/anyimage/operators/test_plane_objects.py +++ b/tests/anyimage/operators/test_plane_objects.py @@ -74,7 +74,6 @@ def test_clipboard_plane_uses_the_mesh_material_slot(): def test_converted_planes_center_the_origin_without_moving_world_corners(): bpy.ops.wm.read_factory_settings(use_empty=True) saved = build_image_depth_plane_group(build_image_plane_group()) - saved["anyimage_asset_revision"] = 11 def create_source(name): image = bpy.data.images.new(name, width=200, height=100) @@ -150,9 +149,8 @@ def test_converted_planes_center_the_origin_without_moving_world_corners(): assert_centered(depth_plane, depth_bounds, depth_matrix) assert depth_plane["anyimage_mesh_shape"] == "DEPTH_PLANE" depth_modifier = depth_plane.modifiers[0] - assert depth_modifier.node_group != saved - assert depth_modifier.node_group["anyimage_asset_revision"] == 16 - assert saved in bpy.data.node_groups.values() and saved["anyimage_asset_revision"] == 11 + assert depth_modifier.node_group == saved + assert saved.name == "O Image Depth Plane" uniform_scale = depth_modifier[ object_data.modifier_input_identifier( depth_modifier.node_group, diff --git a/tests/anyimage/server/test_edge_depth.py b/tests/anyimage/server/test_edge_depth.py new file mode 100644 index 0000000..169e7d9 --- /dev/null +++ b/tests/anyimage/server/test_edge_depth.py @@ -0,0 +1,164 @@ +import numpy as np +import pytest + + +from server.geometry.edge_depth import extend_edge_depth +from server.models.geometry import GeometryFrame +from server.models.onnx_moge2 import depth_to_points + + +def make_frame(depth): + depth = np.asarray(depth, dtype=np.float32) + height, width = depth.shape + normalized = np.array([[1.2, 0, 0.4], [0, 1.4, 0.6], [0, 0, 1]], np.float32) + intrinsics = normalized.copy() + intrinsics[0] *= width + intrinsics[1] *= height + return GeometryFrame( + depth=depth, + validity=np.ones_like(depth), + intrinsics=intrinsics, + points=depth_to_points(depth, normalized), + normal=np.zeros((*depth.shape, 3), np.float32), + ) + + +def silhouette(): + alpha = np.zeros((25, 35), np.float32) + alpha[4:21, 4:22] = 1 + alpha[4:21, 22:25] = 0.2 + depth = np.full(alpha.shape, 20, np.float32) + depth[alpha == 1] = 2 + return make_frame(depth), alpha + + +def test_hair_extension_preserves_core_alpha_and_projection(): + frame, alpha = silhouette() + original_alpha = alpha.copy() + result = extend_edge_depth(frame, alpha, alpha_threshold=0.1) + assert np.all(result.depth[alpha >= 0.1] == 2) + np.testing.assert_array_equal(result.points[9:16, 9:16], frame.points[9:16, 9:16]) + np.testing.assert_array_equal(alpha, original_alpha) + np.testing.assert_array_equal(frame.depth[alpha == 0.2], 20) + np.testing.assert_allclose( + result.points[..., :2] / result.points[..., 2:], + frame.points[..., :2] / frame.points[..., 2:], + atol=1e-6, + ) + assert result.validity is frame.validity + assert result.normal is frame.normal + assert np.isfinite(result.points).all() + + +def test_sampling_padding_is_bounded(): + frame, alpha = silhouette() + result = extend_edge_depth(frame, alpha, alpha_threshold=0.1) + assert np.all(result.depth[12, 24:29] == 2) + assert result.depth[12, 29] == 20 + assert (result.depth[12, 24] + result.depth[12, 25]) / 2 == 2 + + +def test_edge_padding_tracks_selected_threshold(): + frame, alpha = silhouette() + loose = extend_edge_depth(frame, alpha, alpha_threshold=0.1) + strict = extend_edge_depth(frame, alpha) + assert loose.depth[12, 28] == 2 + assert strict.depth[12, 28] == 20 + + +def test_strict_threshold_excludes_lower_alpha_donors(): + alpha = np.zeros((10, 10), np.float32) + alpha[3:7, 3:7] = 0.96 + alpha[5, 5] = 1 + depth = np.full(alpha.shape, 10, np.float32) + depth[5, 5] = 2 + result = extend_edge_depth(make_frame(depth), alpha, alpha_threshold=1) + assert result.depth[5, 6] == 2 + + +def test_image_rectangle_does_not_erode_subject(): + alpha = np.ones((15, 20), np.float32) + alpha[:, 15:] = 0 + depth = np.full(alpha.shape, 2, np.float32) + depth[0, 0] = 3 + frame = make_frame(depth) + result = extend_edge_depth(frame, alpha) + assert result.depth[0, 0] == 3 + np.testing.assert_array_equal(result.points[0, 0], frame.points[0, 0]) + + +def test_small_components_fall_back_without_borrowing_other_depths(): + alpha = np.zeros((15, 30), np.float32) + alpha[5, 5:7] = [0.95, 0.1] + alpha[5, 10:12] = [1, 0.2] + alpha[5, 14] = 0.94 + depth = np.full(alpha.shape, 30, np.float32) + depth[5, 5] = 2 + depth[5, 10] = 8 + result = extend_edge_depth(make_frame(depth), alpha, alpha_threshold=0.1) + assert result.depth[5, 6] == 2 + assert result.depth[5, 11] == 8 + assert result.depth[5, 14] == 30 + assert result.depth[5, 7] == 2 + assert result.depth[5, 9] == 8 + + +def test_diagonal_pixels_share_a_component(): + alpha = np.zeros((10, 10), np.float32) + alpha[4, 4] = 1 + alpha[5, 5] = 0.2 + depth = np.full(alpha.shape, 10, np.float32) + depth[4, 4] = 2 + result = extend_edge_depth(make_frame(depth), alpha, alpha_threshold=0.1) + assert result.depth[5, 5] == 2 + + +@pytest.mark.parametrize("validity", [0.0, 0.2, 0.5]) +def test_invalid_donors_are_excluded(validity): + frame, alpha = silhouette() + frame.depth[10, 6] = 100 + frame.validity[10, 6] = validity + result = extend_edge_depth(frame, alpha) + assert result.depth[10, 4] == 2 + assert result.depth[10, 6] == 100 + assert result.validity[10, 6] == np.float32(validity) + + +def test_internal_mask_hole_does_not_expand_alpha_repair(): + frame, alpha = silhouette() + frame.depth[9:14, 10:15] = 100 + frame.validity[9:14, 10:15] = 0 + original = frame.depth.copy() + result = extend_edge_depth(frame, alpha) + np.testing.assert_array_equal(result.depth[6:19, 6:20], original[6:19, 6:20]) + assert result.depth[12, 24] == 2 + + +def test_opaque_image_keeps_low_mask_depth(): + frame = make_frame(np.ones((10, 10))) + frame.depth[5, 5] = 100 + frame.validity[5, 5] = 0 + assert extend_edge_depth(frame, np.ones((10, 10))) is frame + + +def test_component_without_valid_donors_keeps_its_depth(): + frame, alpha = silhouette() + frame.validity[alpha > 0] = 0 + assert extend_edge_depth(frame, alpha) is frame + + +def test_local_smoothing_reduces_donor_step_and_preserves_core(): + frame, alpha = silhouette() + frame.depth[12:, alpha.shape[1] // 4:22] = 4 + result = extend_edge_depth(frame, alpha) + assert 2 < result.depth[11, 23] < 4 + assert 2 < result.depth[12, 23] < 4 + np.testing.assert_array_equal(result.depth[8:17, 8:17], frame.depth[8:17, 8:17]) + + +@pytest.mark.parametrize("alpha_value", [0, 0.3, 1]) +def test_images_without_repair_return_original_frame(alpha_value): + frame = make_frame(np.ones((5, 8))) + assert extend_edge_depth(frame, np.full((5, 8), alpha_value)) is frame + + diff --git a/tests/anyimage/server/test_geometry_artifacts.py b/tests/anyimage/server/test_geometry_artifacts.py index a98c62a..eef27ec 100644 --- a/tests/anyimage/server/test_geometry_artifacts.py +++ b/tests/anyimage/server/test_geometry_artifacts.py @@ -252,4 +252,5 @@ class GeometryArtifactsTest(ServerTestCase): generate_depth=True, normal_mode="OBJECT", alpha_threshold=0.9, + extend_depth_edges=True, ) diff --git a/tests/anyimage/server/test_prediction_artifacts.py b/tests/anyimage/server/test_prediction_artifacts.py index 26d7a79..accabc3 100644 --- a/tests/anyimage/server/test_prediction_artifacts.py +++ b/tests/anyimage/server/test_prediction_artifacts.py @@ -58,8 +58,9 @@ def test_alpha_resize_retains_threshold_values(tmp_path): (False, "OBJECT"), ], ) -def test_artifacts_preserve_original_geometry_with_alpha( - tmp_path, generate_depth, normal_mode +@pytest.mark.parametrize("extend_depth_edges", [False, True]) +def test_artifacts_extend_only_requested_depth_and_preserve_normals( + tmp_path, generate_depth, normal_mode, extend_depth_edges ): frame, alpha = silhouette() rgba = np.full((*alpha.shape, 4), 255, np.uint8) @@ -73,7 +74,7 @@ def test_artifacts_preserve_original_geometry_with_alpha( check_cancelled=lambda: None, ) with ( - patch.object(artifacts.moge, "infer", return_value=frame), + patch.object(artifacts.moge, "infer", return_value=frame) as infer, patch.object( artifacts, "write_depth_texture", return_value=tmp_path / "depth.exr", @@ -84,10 +85,14 @@ def test_artifacts_preserve_original_geometry_with_alpha( patch.object( artifacts, "_write_normal", return_value=tmp_path / "normal.png" ) as normal, + patch.object(artifacts, "extend_edge_depth", wraps=artifacts.extend_edge_depth) as extend, ): artifacts.generate_moge_artifacts( context, {}, path, generate_depth=generate_depth, normal_mode=normal_mode, + **({"extend_depth_edges": True} if extend_depth_edges else {}), ) + infer.assert_called_once() + assert extend.call_count == int(generate_depth and extend_depth_edges) if normal_mode != "NONE": assert normal.call_args.args[0] is frame else: @@ -96,8 +101,9 @@ def test_artifacts_preserve_original_geometry_with_alpha( written_frame = depth.call_args.args[0] np.testing.assert_array_equal(depth.call_args.kwargs["alpha"], alpha) assert metadata.call_args.args[0] is written_frame - assert written_frame is frame - assert written_frame.depth[12, 24] == 20 + assert (written_frame is frame) == (not extend_depth_edges) + assert written_frame.depth[12, 24] == (2 if extend_depth_edges else 20) + assert frame.depth[12, 24] == 20 assert written_frame.validity is frame.validity else: metadata.assert_not_called() diff --git a/tools/nodes/arrangement/dimensions.json b/tools/nodes/arrangement/dimensions.json index 6300571..87eb184 100644 --- a/tools/nodes/arrangement/dimensions.json +++ b/tools/nodes/arrangement/dimensions.json @@ -53,6 +53,7 @@ "GeometryNodeRepeatOutput": 3, "GeometryNodeSampleIndex": 3, "GeometryNodeSampleNearest": 1, + "GeometryNodeSampleNearestSurface": 1, "GeometryNodeSeparateGeometry": 1, "GeometryNodeSetMaterialIndex": 0, "GeometryNodeSetPosition": 0, @@ -72,6 +73,7 @@ "ShaderNodeMath": 2, "ShaderNodeMix": 3, "ShaderNodeSeparateXYZ": 0, + "ShaderNodeValue": 1, "ShaderNodeVectorMath": 1 } } diff --git a/tools/nodes/common/cutout_center.py b/tools/nodes/common/cutout_center.py new file mode 100644 index 0000000..4d09753 --- /dev/null +++ b/tools/nodes/common/cutout_center.py @@ -0,0 +1,78 @@ +"""Build a shared center on the original Cutout mesh.""" + + +from .nodes import evaluate_field, group_input, sample_field + + +def vector(group, operation, first, second=None): + node = group.nodes.new("ShaderNodeVectorMath") + node.operation = operation + for value, socket in ((first, node.inputs[0]), (second, node.inputs[3 if operation == "SCALE" else 1])): + if value is None: + continue + if hasattr(value, "node"): + group.links.new(value, socket) + else: + socket.default_value = value + return node.outputs[0] + + +def axes(group, value): + node = group.nodes.new("ShaderNodeSeparateXYZ") + group.links.new(value, node.inputs[0]) + return node.outputs + + +def combine(group, x, y, z=0): + node = group.nodes.new("ShaderNodeCombineXYZ") + for value, socket in zip((x, y, z), node.inputs): + if hasattr(value, "node"): + group.links.new(value, socket) + else: + socket.default_value = value + return node.outputs[0] + + +def set_position(group, geometry, position): + node = group.nodes.new("GeometryNodeSetPosition") + group.links.new(geometry, node.inputs["Geometry"]) + group.links.new(position, node.inputs["Position"]) + return node.outputs[0] + + +def sample_surface(group, geometry, value, data_type, domain_position, query): + """Keep point indices on points before the face interpolation adapts domains.""" + field = evaluate_field(group, sample_field(group, geometry, value, data_type), data_type, "POINT") + flat = set_position(group, geometry, domain_position) + sample = group.nodes.new("GeometryNodeSampleNearestSurface") + sample.data_type = data_type + group.links.new(flat, sample.inputs["Mesh"]) + group.links.new(field, sample.inputs["Value"]) + group.links.new(query, sample.inputs["Sample Position"]) + return sample.outputs[0] + + +def build_center_surface(group, unsplit): + """Smooth the unsplit displaced surface without moving its XY coordinates.""" + nodes, links = group.nodes, group.links + controls = group_input(nodes, {"Center Smooth"}) + position = nodes.new("GeometryNodeInputPosition").outputs[0] + components = axes(group, position) + uv = nodes.new("GeometryNodeInputNamedAttribute") + uv.data_type = "FLOAT_VECTOR" + uv.inputs["Name"].default_value = "UVMap" + coordinates = uv.outputs["Attribute"] + blur = nodes.new("GeometryNodeBlurAttribute") + blur.data_type = "FLOAT" + links.new(components["Z"], blur.inputs["Value"]) + links.new(controls.outputs["Center Smooth"], blur.inputs["Iterations"]) + support = set_position(group, unsplit, combine(group, components["X"], components["Y"], blur.outputs[0])) + height = sample_surface(group, support, components["Z"], "FLOAT", coordinates, coordinates) + normal = nodes.new("GeometryNodeInputNormal").outputs[0] + normal_blur = nodes.new("GeometryNodeBlurAttribute") + normal_blur.data_type = "FLOAT_VECTOR" + links.new(evaluate_field(group, normal, "FLOAT_VECTOR", "POINT"), normal_blur.inputs["Value"]) + links.new(controls.outputs["Center Smooth"], normal_blur.inputs["Iterations"]) + normal = normal_blur.outputs[0] + normal = vector(group, "NORMALIZE", sample_surface(group, support, normal, "FLOAT_VECTOR", coordinates, coordinates)) + return {"geometry": support, "preview": support, "height": height, "normal": normal} diff --git a/tools/nodes/groups/image_cutout.py b/tools/nodes/groups/image_cutout.py index a9762b3..f626ea0 100644 --- a/tools/nodes/groups/image_cutout.py +++ b/tools/nodes/groups/image_cutout.py @@ -155,7 +155,6 @@ def build_image_cutout_group(): group = bpy.data.node_groups.new(name, 'GeometryNodeTree') group.is_modifier, group.use_fake_user = (True, True) group.color_tag = 'GEOMETRY' - group['anyimage_asset_revision'] = 13 interface_socket(group, 'Geometry', 'INPUT', 'NodeSocketGeometry') interface_socket(group, 'Mode', 'INPUT', 'NodeSocketMenu') float_input(group, 'Thickness', 1.0).description = 'Control how much the surface bulges in Balloon mode. Zero removes the bulge.' diff --git a/tools/nodes/groups/image_depth_balloon.py b/tools/nodes/groups/image_depth_balloon.py index c4f22ed..45df602 100644 --- a/tools/nodes/groups/image_depth_balloon.py +++ b/tools/nodes/groups/image_depth_balloon.py @@ -107,7 +107,6 @@ def build_image_depth_balloon_group(): group = bpy.data.node_groups.new(name, 'GeometryNodeTree') group.is_modifier, group.use_fake_user = (True, True) group.color_tag = 'GEOMETRY' - group['anyimage_asset_revision'] = 14 interface_socket(group, 'Geometry', 'INPUT', 'NodeSocketGeometry') float_input(group, 'Thickness', 1.0).description = 'Control how much the surface bulges. Zero removes the bulge.' float_input(group, 'Depth Scale', 1.0) diff --git a/tools/nodes/groups/image_depth_panorama.py b/tools/nodes/groups/image_depth_panorama.py index 2dd21cd..53edc79 100644 --- a/tools/nodes/groups/image_depth_panorama.py +++ b/tools/nodes/groups/image_depth_panorama.py @@ -39,7 +39,6 @@ def _quad_sphere(group, segments): def build_image_depth_panorama_group(): group = bpy.data.node_groups.new("O Image Depth Panorama", "GeometryNodeTree") group.is_modifier = group.use_fake_user = True - group["anyimage_asset_revision"] = 15 interface_socket(group, "Geometry", "INPUT", "NodeSocketGeometry") subdivide = interface_socket(group, "Subdivide", "INPUT", "NodeSocketInt") subdivide.description = "Increase mesh density to represent finer shapes. Higher levels create more vertices and require more processing time." diff --git a/tools/nodes/groups/image_depth_plane.py b/tools/nodes/groups/image_depth_plane.py index 386e8e8..30f16b9 100644 --- a/tools/nodes/groups/image_depth_plane.py +++ b/tools/nodes/groups/image_depth_plane.py @@ -25,7 +25,6 @@ def build_image_depth_plane_group(image_plane_group): "GeometryNodeTree", ) group.is_modifier = True - group["anyimage_asset_revision"] = 16 geometry = _interface_socket(group, "Geometry", "INPUT", "NodeSocketGeometry") geometry.structure_type = "SINGLE" subdivide = _interface_socket(group, "Subdivide", "INPUT", "NodeSocketInt")