AnyImage/tests/tools/nodes/test_depth_strip_cleanup.py

"""Evaluate opposite-boundary cleanup on explicit mesh topology."""

import bmesh
import bpy
import numpy as np
import pytest

from tools.nodes.common.depth_surface import _remove_strip_faces


def cleaned_faces(vertices, faces, *, triangles=False):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    mesh = bpy.data.meshes.new("Topology")
    mesh.from_pydata(vertices, [], faces)
    ids = mesh.attributes.new("source_face", "INT", "FACE")
    ids.data.foreach_set("value", np.arange(len(faces), dtype=np.int32))
    uv = mesh.uv_layers.new(name="UVMap")
    for loop in mesh.loops:
        uv.data[loop.index].uv = mesh.vertices[loop.vertex_index].co[:2]
    obj = bpy.data.objects.new("Topology", mesh)
    bpy.context.collection.objects.link(obj)
    group = bpy.data.node_groups.new("Strip cleanup test", "GeometryNodeTree")
    group.interface.new_socket(name="Geometry", in_out="INPUT", socket_type="NodeSocketGeometry")
    group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = group.nodes.new("NodeGroupInput")
    result = _remove_strip_faces(group, source.outputs["Geometry"], triangles=triangles)
    group.links.new(result, group.nodes.new("NodeGroupOutput").inputs["Geometry"])
    obj.modifiers.new("Cleanup", "NODES").node_group = group
    bpy.context.view_layer.update()
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    result = evaluated.to_mesh()
    try:
        assert all(np.isfinite(vertex.co).all() for vertex in result.vertices)
        if not result.polygons:
            assert not result.vertices and not result.edges
            return set()
        assert "UVMap" in result.attributes
        for loop in result.loops:
            np.testing.assert_allclose(result.uv_layers["UVMap"].data[loop.index].uv,
                                       result.vertices[loop.vertex_index].co[:2])
        return {value.value for value in result.attributes["source_face"].data}
    finally:
        evaluated.to_mesh_clear()


@pytest.mark.parametrize("cells, expected", [
    ([(0, 0)], set()),
    ([(x, 0) for x in range(5)], set()),
    ([(x, y) for y in range(3) for x in range(3)], set(range(9))),
    # The center becomes isolated only after its three arms disappear.
    ([(0, 0), (1, 0), (2, 0), (1, 1)], {1}),
])
@pytest.mark.parametrize("offset, reverse", [(0, False), (1, False), (2, True), (3, True)])
def test_quad_strips_and_corners_are_classified_once(cells, expected, offset, reverse):
    vertices, indices, faces = [], {}, []
    for x, y in cells:
        face = []
        for point in [(x, y, 0), (x + 1, y, 0), (x + 1, y + 1, 0), (x, y + 1, 0)]:
            if point not in indices:
                indices[point] = len(vertices)
                vertices.append(point)
            face.append(indices[point])
        face = face[offset:] + face[:offset]
        faces.append(face[::-1] if reverse else face)
    assert cleaned_faces(vertices, faces) == expected


@pytest.mark.parametrize("width,height,expected", [(2, 2, 8), (3, 3, 18)])
def test_triangle_pairs_remove_strips_and_preserve_surface_corners(width, height, expected):
    vertices = [(x, y, 0) for y in range(height + 1) for x in range(width + 1)]
    faces = []
    for y in range(height):
        for x in range(width):
            a = y * (width + 1) + x
            b, c, d = a + 1, a + width + 2, a + width + 1
            faces.extend([(a, b, c), (a, c, d)])
    # Reverse alternate faces to test matching by vertex identity, not winding.
    faces = [face[::-1] if i % 2 else face[1:] + face[:1] for i, face in enumerate(faces)]
    result = cleaned_faces(vertices, faces, triangles=True)
    assert len(result) == expected


@pytest.mark.parametrize("kind", ["plane", "cutout"])
def test_strips_are_removed_before_projection_and_closed_thickness(kind):
    from tests.support.depth_surface import surface, evaluated, assert_closed, component_count
    from tests.support.planes import create_surface

    if kind == "plane":
        obj, set_value, _, _, image = create_surface("DEPTH")
        set_value("Subdivide", 1)
    else:
        obj, set_value = surface(step=0.0)
        image = bpy.data.images["Camera"]
    width, height = image.size
    pixels = np.array(image.pixels[:], np.float32).reshape(height, width, 4)
    pixels[..., 2] = 1
    pixels[:, width // 4:3 * width // 4, 2] = 21
    image.pixels.foreach_set(pixels.ravel())
    image.update()
    set_value("Depth Scale", 1.0)
    set_value("Reference Depth", 11.0)
    if kind == "cutout":
        bm = bmesh.new()
        bm.from_mesh(obj.data)
        bmesh.ops.triangulate(bm, faces=list(bm.faces))
        bm.to_mesh(obj.data)
        bm.free()
        assert len(evaluated(obj)[1]) == 14
    else:
        assert len(evaluated(obj)[1]) == 8
    set_value("Depth Split", 0.5)
    vertices, faces = evaluated(obj)
    if kind == "cutout":
        assert len(faces) == 6 and len(vertices) == 7
    else:
        assert len(faces) == 4 and len(vertices) == 9
    assert component_count(faces) == 1
    set_value("Thickness", 0.2)
    vertices, shell = evaluated(obj)
    assert len(shell) == (18 if kind == "cutout" else 16)
    assert_closed(shell)
    assert np.isfinite(vertices).all()
    assert component_count(shell) == 1
    set_value("Thickness", 0.0)
    assert len(evaluated(obj)[1]) == (6 if kind == "cutout" else 4)


@pytest.mark.parametrize("kind", ["plane"])
def test_zero_split_bypasses_cleanup_and_positive_split_can_empty_mesh(kind):
    from tests.support.depth_surface import surface, evaluated
    from tests.support.planes import create_surface

    obj, set_value, _, _, _ = create_surface("DEPTH")
    set_value("Subdivide", 0)
    original, faces = evaluated(obj)
    assert len(faces) == 2
    set_value("Depth Split", 0.001)
    assert not evaluated(obj)[1]
    set_value("Thickness", 0.2)
    assert not evaluated(obj)[1]
    set_value("Thickness", 0.0)
    set_value("Depth Split", 0.0)
    restored, restored_faces = evaluated(obj)
    assert restored_faces == faces
    np.testing.assert_array_equal(restored, original)


def test_cutout_removes_an_all_boundary_triangle_strip():
    import bmesh
    from tests.support.depth_surface import surface, evaluated

    obj, set_value = surface(step=0.0, resolution=1)
    mesh = bmesh.new()
    mesh.from_mesh(obj.data)
    bmesh.ops.triangulate(mesh, faces=list(mesh.faces))
    mesh.to_mesh(obj.data)
    mesh.free()
    # Every vertex lies on the boundary, even before any depth separation.
    for split in (0.0, 0.001, 0.0):
        set_value("Depth Split", split)
        for thickness in (0.0, 0.2):
            set_value("Thickness", thickness)
            points, faces = evaluated(obj)
            assert not faces and not len(points)