AnyImage/tests/tools/nodes/test_cutout_symmetry.py

"""Check the canonical Y mirror and final X orientation for Depth Symmetry."""

from collections import Counter

import bpy
import numpy as np
import pytest
from scipy.spatial import cKDTree

from anyimage.common.object import modifier_input_identifier, set_modifier_input
from anyimage.operators.cutout_tool.object import load_cutout_node_group
from tests.support.depth_surface import assert_closed, evaluated
from tools.nodes.common.normal_map import AXIS_ATTRIBUTE, FACE_ATTRIBUTE, ROTATION_ATTRIBUTE
from tools.nodes.groups.image_depth_cutout import build_image_depth_cutout_group
from tools.nodes.groups.image_cutout_symmetry import build_image_cutout_symmetry_group


def _set_value(obj, modifier, group):
    def setter(name, value, subtype=None):
        set_modifier_input(
            modifier,
            modifier_input_identifier(group, name, subtype=subtype),
            value,
        )
        obj.update_tag(refresh={"DATA"})
        bpy.context.view_layer.update()

    return setter


@pytest.fixture(params=["definition", "asset"])
def symmetry(request):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    count = 5
    z, x = np.meshgrid(
        np.linspace(-1, 1, count),
        np.linspace(-1, 1, count),
        indexing="ij",
    )
    faces = []
    for row in range(count - 1):
        for column in range(count - 1):
            start = row * count + column
            faces.extend(((start, start + 1, start + count + 1), (start, start + count + 1, start + count)))
    boundary = (np.maximum(abs(x), abs(z)) == 1).ravel()
    faces = [face for face in faces if not all(boundary[index] for index in face)]

    mesh = bpy.data.meshes.new("Curved Depth")
    mesh.from_pydata(np.column_stack((x.ravel(), np.zeros(x.size), z.ravel())), [], faces)
    uv = mesh.uv_layers.new(name="UVMap")
    for loop in mesh.loops:
        row, column = divmod(loop.vertex_index, count)
        uv.data[loop.index].uv = ((column + 0.5) / count, (row + 0.5) / count)
    profile = 0.3 * np.maximum(0, 1 - np.maximum(abs(x), abs(z)))
    mesh.attributes.new("o_balloon", "FLOAT", "POINT").data.foreach_set(
        "value", profile.ravel()
    )
    mesh.attributes.new("user_probe", "FLOAT", "POINT").data.foreach_set(
        "value", np.ones(x.size)
    )

    image = bpy.data.images.new("Depth", width=count, height=count, float_buffer=True)
    image.colorspace_settings.name = "Non-Color"
    image.alpha_mode = "CHANNEL_PACKED"
    depth = 1 + 0.6 * x**2 + 0.3 * z**2
    image.pixels.foreach_set(
        np.dstack((x, -z, depth, np.ones_like(x))).astype(np.float32).ravel()
    )

    obj = bpy.data.objects.new("Symmetry", mesh)
    bpy.context.collection.objects.link(obj)

    cutout_group = (
        build_image_depth_cutout_group()
        if request.param == "definition"
        else load_cutout_node_group("DEPTH_SOLID")
    )
    mirror_group = (
        build_image_cutout_symmetry_group()
        if request.param == "definition"
        else load_cutout_node_group("DEPTH_SYMMETRY")
    )
    cutout = obj.modifiers.new(cutout_group.name, "NODES")
    cutout.node_group = cutout_group
    mirror = obj.modifiers.new(mirror_group.name, "NODES")
    mirror.node_group = mirror_group

    cutout_setter = _set_value(obj, cutout, cutout_group)
    mirror_setter = _set_value(obj, mirror, mirror_group)
    mirror_setter("Fill Smooth", 0)
    for name, value in (
        ("Depth Image", image),
        ("Reference Depth", 1.45),
        ("Depth Split", 0),
        ("Boundary Smooth", 0),
    ):
        cutout_setter(name, value)
    return obj, cutout, mirror, cutout_setter, mirror_setter


def test_symmetry_interface_and_cutout_defaults(symmetry):
    _obj, cutout, mirror, _cutout_setter, _mirror_setter = symmetry
    mirror_inputs = [
        item
        for item in mirror.node_group.interface.items_tree
        if item.item_type == "SOCKET" and item.in_out == "INPUT"
    ]
    assert [item.name for item in mirror_inputs] == [
        "Geometry",
        "Direction",
        "Offset",
        "Scale",
        "Fill Sides",
        "Fill Smooth",
    ]
    assert tuple(mirror_inputs[1].default_value) == (0, 0, 1)
    assert mirror_inputs[3].default_value == 1
    assert mirror_inputs[2].min_value == -10
    assert mirror_inputs[2].max_value == 10
    assert mirror_inputs[3].min_value == 0
    assert mirror_inputs[3].max_value == 2
    assert mirror_inputs[4].default_value is True
    assert mirror_inputs[5].default_value == 2

    cutout_names = {
        item.name
        for item in cutout.node_group.interface.items_tree
        if item.item_type == "SOCKET" and item.in_out == "INPUT"
    }
    assert "Depth Axis" not in cutout_names
    thickness = [
        item
        for item in cutout.node_group.interface.items_tree
        if item.item_type == "SOCKET" and item.name == "Thickness"
    ]
    assert len(thickness) == 2 and all(item.default_value == 0 for item in thickness)


def test_final_symmetry_axis_is_x_and_z_is_up(symmetry):
    obj, _cutout, _mirror, cutout_setter, _mirror_setter = symmetry
    cutout_setter("Mode", 1)
    cutout_setter("Thickness", 0.2, subtype="DISTANCE")
    points, faces = evaluated(obj)

    assert len(points) and np.isfinite(points).all()
    assert cKDTree(points).query(points * (-1, 1, 1))[0].max() < 1e-5
    assert np.ptp(points[:, 2]) > 1.0
    assert_closed(faces)


def test_object_normal_attributes_follow_final_orientation(symmetry):
    obj, _cutout, _mirror, cutout_setter, _mirror_setter = symmetry
    cutout_setter("Mode", 1)
    cutout_setter("Thickness", 0.2, subtype="DISTANCE")
    obj.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()

    evaluated_obj = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = evaluated_obj.to_mesh()
    try:
        assert AXIS_ATTRIBUTE in mesh.attributes
        assert ROTATION_ATTRIBUTE in mesh.attributes
        assert FACE_ATTRIBUTE in mesh.attributes
        axis = np.array([item.value for item in mesh.attributes[AXIS_ATTRIBUTE].data])
        face = np.array([item.value for item in mesh.attributes[FACE_ATTRIBUTE].data])
        rotation = np.array(
            [item.vector[:] for item in mesh.attributes[ROTATION_ATTRIBUTE].data]
        )
        assert np.allclose(axis, 1)
        assert set(np.unique(face)) <= {1, 2, 3}
        assert np.isfinite(rotation).all()
    finally:
        evaluated_obj.to_mesh_clear()