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.

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
 
 

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)

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

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,

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
+
+

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,
         )

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()