AnyImage/tests/support/image_objects.py

import bpy
import numpy as np


from anyimage.common import image as image_data
from anyimage.common import object as object_data
from anyimage.common.selection import SelectionMask
from anyimage.operators.cutout_tool import geometry as cutout_geometry
from anyimage.operators.cutout_tool import object as cutout_object
from tests.support.color_image import color_image as prepared_color_image


def _inputs(group):
    return [
        item.name
        for item in group.interface.items_tree
        if item.item_type == "SOCKET" and item.in_out == "INPUT"
    ]


def _group_node(group, node_tree):
    return next(
        node
        for node in node_tree.nodes
        if node.bl_idname == "ShaderNodeGroup" and node.node_tree == group
    )


def _create_cutout_shape(
    source,
    shape,
    edge_length,
    *,
    bounds=None,
    depth_image=None,
    depth_metadata=None,
    selection_values=None,
    gesture="LASSO",
):
    if bounds is None:
        bounds = (0, 0, *tuple(int(value) for value in source.data.size))
    left, top, right, bottom = bounds
    if selection_values is None:
        selection_values = np.ones((bottom - top, right - left), dtype=np.float32)
    selection_mask = SelectionMask(selection_values, bounds)
    source_rgba = image_data.image_rgba(source.data)
    color_image = prepared_color_image(
        source.data,
        selection_mask.bounds,
        source_rgba,
    )
    base_shape = cutout_geometry.build_base_shape(
        bpy.context,
        source,
        edge_length,
        selection_mask.values,
        selection_mask.bounds,
    )
    return cutout_object.create_shape_object(
        bpy.context,
        source,
        shape,
        edge_length,
        selection_mask.values,
        selection_mask.bounds,
        base_shape,
        depth_image=depth_image,
        depth_metadata=depth_metadata,
        gesture=gesture,
        color_image=color_image,
    )


def _depth_image(points, valid, name="Depth"):
    points = np.asarray(points, dtype=np.float32)
    valid = np.asarray(valid, dtype=bool)
    height, width = valid.shape
    pixels = np.zeros((height, width, 4), dtype=np.float32)
    pixels[..., :3] = points
    pixels[..., 3] = valid
    image = bpy.data.images.new(
        name,
        width=width,
        height=height,
        alpha=True,
        float_buffer=True,
    )
    image.colorspace_settings.name = "Non-Color"
    image.alpha_mode = "CHANNEL_PACKED"
    image.pixels.foreach_set(np.flipud(pixels).ravel())
    image.pack()
    return image


def _depth_metadata(intrinsics, image_size, reference_depth):
    return {
        "intrinsics": tuple(tuple(float(value) for value in row) for row in intrinsics),
        "image_size": tuple(int(value) for value in image_size),
        "reference_depth": float(reference_depth),
    }


def _evaluated_positions(obj):
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    result = evaluated.to_mesh()
    try:
        return np.asarray([tuple(vertex.co) for vertex in result.vertices])
    finally:
        evaluated.to_mesh_clear()


def _add_cutout_modifier(obj, shape):
    group = cutout_object.load_cutout_node_group(shape)
    modifier = obj.modifiers.new(group.name, "NODES")
    modifier.node_group = group
    if shape in {"FLAT", "DEPTH_SOLID"}:
        object_data.set_modifier_input(modifier, object_data.modifier_input_identifier(group, "Mode"), 1)
    object_data.set_modifier_input(
        modifier,
        object_data.modifier_input_identifier(group, "Cleanup Threshold"),
        0.0,
    )
    if shape == "DEPTH_SOLID":
        object_data.set_modifier_input(
            modifier,
            object_data.modifier_input_identifier(group, "Smooth Weight"),
            0.0,
        )
    return modifier


def _depth_surface_step_object():
    vertices = [(float(x), float(y), 0.0) for y in range(3) for x in range(3)]
    faces = []
    for y in range(2):
        for x in range(2):
            a = y * 3 + x
            faces.extend(((a, a + 1, a + 4), (a, a + 4, a + 3)))
    data = bpy.data.meshes.new("Depth Surface Step")
    data.from_pydata(vertices, (), faces)
    uv_layer = data.uv_layers.new(name="UVMap")
    for loop in data.loops:
        coordinate = data.vertices[loop.vertex_index].co
        uv_layer.data[loop.index].uv = (coordinate.x * 0.5, coordinate.y * 0.5)
    obj = bpy.data.objects.new("Depth Surface Step", data)
    bpy.context.collection.objects.link(obj)
    modifier = _add_cutout_modifier(obj, "DEPTH_SOLID")

    points = np.empty((3, 3, 3), dtype=np.float32)
    for y in range(3):
        for x in range(3):
            points[y, x] = (float(x), -float(y), 1.0)
    points[0, 0, 2] = 6.0
    depth_image = _depth_image(points, np.ones((3, 3), dtype=bool))
    for name, value in (
        ("Depth Image", depth_image),
        ("Depth Scale", 1.0),
        ("Uniform Scale", 1.0),
        ("Reference Depth", 1.0),
        ("Thickness", 0.0),
    ):
        object_data.set_modifier_input(
            modifier,
            object_data.modifier_input_identifier(
                modifier.node_group, name,
                subtype="DISTANCE" if name == "Thickness" else None,
            ),
            value,
        )
    return obj, modifier


def _evaluated_topology(obj):
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    result = evaluated.to_mesh()
    try:
        face_edges = {
            tuple(sorted((start, end)))
            for polygon in result.polygons
            for start, end in zip(
                polygon.vertices,
                (*polygon.vertices[1:], polygon.vertices[0]),
            )
        }
        mesh_edges = {
            tuple(sorted(edge.vertices))
            for edge in result.edges
        }
        face_vertices = {
            vertex
            for polygon in result.polygons
            for vertex in polygon.vertices
        }
        return (
            len(result.vertices),
            len(result.polygons),
            mesh_edges - face_edges,
            set(range(len(result.vertices))) - face_vertices,
        )
    finally:
        evaluated.to_mesh_clear()