AnyImage/tests/tools/nodes/test_cutout_symmetry.py

"""Check composition against the original Projected geometry baseline."""
from collections import Counter
import json
from pathlib import Path

import bpy
import numpy as np
import pytest
from scipy.spatial import cKDTree
from mathutils import Euler, Vector

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 evaluated, assert_closed
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
from tools.nodes.common.normal_map import ROTATION_ATTRIBUTE, FACE_ATTRIBUTE, AXIS_ATTRIBUTE


@pytest.fixture(params=['definition', 'asset'])
def symmetry(request):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    count = 5
    y, x = np.meshgrid(np.linspace(-1, 1, count), np.linspace(-1, 1, count), indexing='ij')
    faces = []
    for row in range(count - 1):
        for col in range(count - 1):
            a = row * count + col
            faces.extend(((a, a + 1, a + count + 1), (a, a + count + 1, a + count)))
    # Match Cutout's outline cleanup before recording or comparing projection.
    boundary = (np.maximum(abs(x), abs(y)) == 1).ravel()
    faces = [face for face in faces if not all(boundary[i] for i in face)]
    mesh = bpy.data.meshes.new('Curved Depth')
    mesh.from_pydata(np.column_stack((x.ravel(), y.ravel(), np.zeros(x.size))), [], faces)
    uv = mesh.uv_layers.new(name='UVMap')
    for loop in mesh.loops:
        row, col = divmod(loop.vertex_index, count)
        uv.data[loop.index].uv = ((col + .5) / count, (row + .5) / count)
    profile = .3 * np.maximum(0, 1 - np.maximum(abs(x), abs(y)))
    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 + .6 * x ** 2 + .3 * y ** 2
    image.pixels.foreach_set(np.dstack((x, -y, depth, np.ones_like(x))).astype(np.float32).ravel())
    obj = bpy.data.objects.new('Symmetry', mesh)
    bpy.context.collection.objects.link(obj)
    setters = []
    for name, builder in [('DEPTH_SOLID', build_image_depth_cutout_group),
                          ('DEPTH_SYMMETRY', build_image_cutout_symmetry_group)]:
        group = builder() if request.param == 'definition' else load_cutout_node_group(name)
        modifier = obj.modifiers.new(group.name, 'NODES')
        modifier.node_group = group

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

        setters.append(set_value)
    cutout, mirror = setters
    mirror('Fill Smooth', 0)
    for name, value in [('Depth Image', image), ('Reference Depth', 1.45),
                        ('Depth Split', 0), ('Boundary Smooth', 0), ('Smooth', 0), ('Depth Axis', 0)]:
        cutout(name, value)
    yield obj, cutout, mirror
    bpy.ops.wm.read_factory_settings(use_empty=True)


def oriented_faces(points, faces):
    """Ignore vertex numbering while retaining polygon winding."""
    keys = [tuple(np.round(point, 6)) for point in points]
    result = []
    for face in faces:
        loop = tuple(keys[i] for i in face)
        result.append(min(loop[i:] + loop[:i] for i in range(len(loop))))
    return Counter(result)


@pytest.mark.parametrize('mode', [0, 1])
def test_zero_thickness_matches_original_projected(symmetry, mode):
    obj, cutout, mirror = symmetry
    cutout('Mode', mode)
    baseline = json.loads((Path(__file__).parent / 'fixtures/projected_symmetry.json').read_text())
    for case in baseline['cases']:
        mirror('Symmetry Direction', case['direction'])
        mirror('Fill Sides', case['filled'])
        points, faces = evaluated(obj)
        expected = np.array(case['points'])
        assert len(points) == len(expected)
        assert cKDTree(points).query(expected)[0].max() < 1e-6
        assert oriented_faces(points, faces) == oriented_faces(expected, case['faces'])


@pytest.mark.parametrize('mode', [0, 1])
@pytest.mark.parametrize('smooth', [0, 3])
def test_completed_geometry_is_mirrored_and_closed(symmetry, mode, smooth):
    obj, cutout, mirror = symmetry
    cutout('Mode', mode)
    cutout('Thickness', .2 if mode else .6, 'DISTANCE' if mode else 'NONE')
    mirror('Symmetry Direction', (.3, -.2, 1))
    mirror('Fill Smooth', smooth)
    points, faces = evaluated(obj)
    assert len(points) and np.isfinite(points).all()
    assert cKDTree(points).query(points * (1, 1, -1))[0].max() < 1e-6
    assert_closed(faces)
    orientation = Counter()
    for face in faces:
        for a, b in zip(face, (*face[1:], face[0])):
            orientation[tuple(sorted((a, b)))] += 1 if a < b else -1
    assert not any(orientation.values())
    result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    mesh = result.to_mesh()
    try:
        assert 'UVMap' in mesh.uv_layers and 'user_probe' in mesh.attributes
        direction = Vector((.3, -.2, 1)).normalized()
        wall_count = 0
        for polygon in mesh.polygons:
            i = polygon.index
            rotation = Euler(mesh.attributes[ROTATION_ATTRIBUTE].data[i].vector)
            np.testing.assert_allclose(rotation.to_matrix() @ Vector((0, 0, 1)), direction, atol=1e-6)
            assert mesh.attributes[AXIS_ATTRIBUTE].data[i].value == 0
            if mesh.attributes[FACE_ATTRIBUTE].data[i].value == 3:
                wall_count += 1
                if smooth == 0:
                    assert abs(polygon.normal.z) < 1e-5
        assert wall_count
    finally:
        result.to_mesh_clear()


def test_fill_smoothing_preserves_surfaces_and_only_moves_wall_interiors(symmetry):
    obj, _, mirror = symmetry
    mirror('Fill Sides', False)
    original, original_faces = evaluated(obj)
    mirror('Fill Smooth', 3)
    points, faces = evaluated(obj)
    np.testing.assert_array_equal(points, original)
    assert faces == original_faces
    mirror('Fill Sides', True)
    points, faces = evaluated(obj)
    assert_closed(faces)
    assert cKDTree(points).query(original)[0].max() < 1e-6
    assert cKDTree(points).query(points * (1, 1, -1))[0].max() < 1e-6
    assert len(points) > len(original)
    assert cKDTree(original[:, :2]).query(points[:, :2])[0].max() > .001
    assert oriented_faces(original, original_faces).items() <= oriented_faces(points, faces).items()


def test_direction_offset_and_empty_geometry(symmetry):
    obj, _, mirror = symmetry
    mirror('Depth Offset', 2.)
    mirror('Symmetry Direction', (.3, -.2, 1))
    original, original_faces = evaluated(obj)
    mirror('Depth Offset', 2.2)
    points, faces = evaluated(obj)
    assert faces == original_faces
    np.testing.assert_allclose(points[:, :2], original[:, :2], atol=1e-6)
    np.testing.assert_allclose(abs(points[:, 2]) - abs(original[:, 2]), .2, atol=1e-6)
    mirror('Depth Offset', -10.)
    points, faces = evaluated(obj)
    assert not len(points) and not faces


def test_symmetry_interface_and_cutout_defaults(symmetry):
    obj, _, _ = symmetry
    inputs = [item for item in obj.modifiers[1].node_group.interface.items_tree
              if item.item_type == 'SOCKET' and item.in_out == 'INPUT']
    assert [item.name for item in inputs] == ['Geometry', 'Symmetry Direction', 'Depth Offset', 'Fill Sides', 'Fill Smooth']
    assert inputs[-1].default_value == 2 and inputs[-2].default_value is True
    thickness = [item for item in obj.modifiers[0].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)