AnyImage/tools/validate_depth_plane_asset.py

from pathlib import Path
import sys

import bpy


PROJECT_ROOT = Path(__file__).parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(
    0,
    str(next((PROJECT_ROOT / "src" / "anyimage" / "wheels").glob("blendjob-*.whl"))),
)

from anyimage.operators.material import (
    DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME,
    SHADELESS_NODE_GROUP_NAME,
    SURFACE_MATERIAL_NODE_GROUP_NAME,
    create_material,
    get_modifier_input,
    load_depth_plane_node_group,
    material_color_space_name,
    modifier_input_driver_path,
    modifier_input_identifier,
    set_modifier_input,
)
from anyimage.operators.depth_plane.build import add_depth_plane_processing
from anyimage.operators.surface import load_geometry_node_group


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    group = load_depth_plane_node_group()
    geometry_textures = [
        node
        for node in group.nodes
        if node.bl_idname == "GeometryNodeImageTexture"
    ]
    assert geometry_textures
    assert all(node.extension == "EXTEND" for node in geometry_textures)
    background_depth = next(
        item
        for item in group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Background Depth"
    )
    assert background_depth.subtype == "DISTANCE"
    assert background_depth.structure_type == "SINGLE"
    assert background_depth.default_value == 1.0
    subdivision = next(
        item
        for item in group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Subdivision"
    )
    assert subdivision.default_value == 6
    frame_input = next(
        item
        for item in group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Frame"
    )
    assert frame_input.structure_type == "SINGLE"
    geometry_input = next(
        item
        for item in group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Geometry"
    )
    assert geometry_input.structure_type == "SINGLE"
    all_inputs = [
        (item.name, getattr(item.parent, "name", None), item.hide_in_modifier)
        for item in group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
    ]
    visible_inputs = [name for name, parent, hidden in all_inputs if not hidden]
    assert visible_inputs[:6] == [
        "Geometry",
        "Displacement",
        "Subdivision",
        "Depth Scale",
        "Background Depth",
    ]
    assert not getattr(background_depth.parent, "name", "")
    assert not any(
        item.name == "Remove Background"
        for item in group.interface.items_tree
    )
    depth_scale = next(
        item
        for item in group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Depth Scale"
    )
    assert depth_scale.default_value == 1.0
    assert depth_scale.min_value < 0.0
    assert depth_scale.structure_type == "SINGLE"
    stores = {
        node.inputs["Name"].default_value: node
        for node in group.nodes
        if node.bl_idname == "GeometryNodeStoreNamedAttribute"
    }
    attribute_names = (
        "o_material_displacement",
        "o_background_depth",
        "o_depth_scale",
    )
    for name in attribute_names:
        assert stores[name].domain == "POINT"
    assert not any(
        node.bl_idname == "GeometryNodeGeometryToInstance"
        for node in group.nodes
    )

    mesh = bpy.data.meshes.new("Depth Plane")
    mesh.from_pydata(
        [(-1, -1, 0), (1, -1, 0), (-1, 1, 0), (1, 1, 0)],
        [],
        [(0, 1, 3, 2)],
    )
    obj = bpy.data.objects.new("Depth Plane", mesh)
    bpy.context.scene.collection.objects.link(obj)
    depth = bpy.data.images.new("Depth", width=2, height=2, float_buffer=True)
    depth.pixels[:] = [0.25, 0.25, 0.25, 1.0] * 4
    foreground = bpy.data.images.new("Foreground", width=2, height=2)
    foreground.pixels[:] = [0.8, 0.2, 0.1, 1.0] * 4
    foreground.alpha_mode = "CHANNEL_PACKED"
    material = create_material(
        "Depth Plane",
        foreground,
        1.0,
        displacement_image=depth,
    )
    mesh.materials.append(material)
    material_group = bpy.data.node_groups[DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME]
    interface_inputs = [
        item.name
        for item in material_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
    ]
    interface_outputs = [
        item.name
        for item in material_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "OUTPUT"
    ]
    assert interface_inputs == [
        "Color",
        "Alpha",
        "Depth",
        "Bump Scale",
    ]
    assert interface_outputs == [
        "Color",
        "Roughness",
        "Alpha",
        "Normal",
        "Displacement",
    ]
    depth_input = next(
        item
        for item in material_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Depth"
    )
    assert depth_input.socket_type == "NodeSocketFloat"
    bump_scale = next(
        item
        for item in material_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Bump Scale"
    )
    assert bump_scale.default_value == 1.0
    assert bump_scale.min_value < 0.0
    material_group_nodes = [
        node
        for node in material.node_tree.nodes
        if node.bl_idname == "ShaderNodeGroup"
        and node.node_tree == material_group
    ]
    assert len(material_group_nodes) == 1
    assert material_group_nodes[0].name == "O DA3 Depth Plane Layer"
    bump = next(
        node
        for node in material_group.nodes
        if node.bl_idname == "ShaderNodeBump"
    )
    bump_height_link = next(
        link
        for link in material_group.links
        if link.to_node == bump and link.to_socket == bump.inputs["Height"]
    )
    assert bump_height_link.from_node.bl_idname == "NodeGroupInput"
    assert bump_height_link.from_socket.name == "Color"
    attributes = [
        node
        for node in material_group.nodes
        if node.bl_idname == "ShaderNodeAttribute"
    ]
    assert len(attributes) == 3
    assert all(node.attribute_type == "GEOMETRY" for node in attributes)
    assert {node.attribute_name for node in attributes} == {
        "o_material_displacement",
        "o_background_depth",
        "o_depth_scale",
    }
    assert not any(
        node.bl_idname == "ShaderNodeAttribute"
        for node in material.node_tree.nodes
    )
    forbidden_outer_nodes = {
        "ShaderNodeAttribute",
        "ShaderNodeBump",
        "ShaderNodeDisplacement",
        "ShaderNodeFloatCurve",
        "ShaderNodeMath",
        "ShaderNodeMix",
    }
    assert not any(
        node.bl_idname in forbidden_outer_nodes
        for node in material.node_tree.nodes
    )
    assert all(
        node.extension == "EXTEND"
        for node in material.node_tree.nodes
        if node.bl_idname == "ShaderNodeTexImage"
    )
    shadeless_material = create_material(
        "Shadeless",
        foreground,
        1.0,
        shadeless=True,
    )
    shadeless_group = bpy.data.node_groups[SHADELESS_NODE_GROUP_NAME]
    shadeless_nodes = [
        node
        for node in shadeless_material.node_tree.nodes
        if node.bl_idname == "ShaderNodeGroup"
        and node.node_tree == shadeless_group
    ]
    assert len(shadeless_nodes) == 1
    assert not any(
        node.bl_idname == "ShaderNodeEmission"
        for node in shadeless_material.node_tree.nodes
    )
    assert [
        item.name
        for item in shadeless_group.interface.items_tree
        if item.item_type == "SOCKET" and item.in_out == "INPUT"
    ] == ["Color"]
    assert [
        item.name
        for item in shadeless_group.interface.items_tree
        if item.item_type == "SOCKET" and item.in_out == "OUTPUT"
    ] == ["Shader"]
    assert {
        node.label: node.operation
        for node in shadeless_group.nodes
        if node.bl_idname == "ShaderNodeMath"
    } == {
        "Bounce Count": "SUBTRACT",
        "Camera or Refracted": "SUBTRACT",
        "Limit Reflections": "SUBTRACT",
        "Camera Ray to Glossy": "MULTIPLY",
        "Shadow or Reflection?": "MAXIMUM",
        "Shadow, Reflect or Refract?": "MAXIMUM",
    }
    assert foreground.colorspace_settings.name == material_color_space_name(
        bpy.context.scene
    )
    surface_material = create_material("Surface", foreground, 1.0)
    surface_group = bpy.data.node_groups[SURFACE_MATERIAL_NODE_GROUP_NAME]
    surface_group_nodes = [
        node
        for node in surface_material.node_tree.nodes
        if node.bl_idname == "ShaderNodeGroup"
        and node.node_tree == surface_group
    ]
    assert len(surface_group_nodes) == 1
    assert surface_group_nodes[0].name == "O Image Layer"
    surface_inputs = [
        item.name
        for item in surface_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
    ]
    assert surface_inputs == [
        "Color",
        "Alpha",
        "Normal",
        "Normal Scale",
        "Object Space",
        "Bump Scale",
    ]
    normal_input = next(
        item
        for item in surface_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Normal"
    )
    assert tuple(normal_input.default_value) == (0.5, 0.5, 1.0, 1.0)
    normal_scale = next(
        item
        for item in surface_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Normal Scale"
    )
    assert normal_scale.default_value == 1.0
    surface_outputs = [
        item.name
        for item in surface_group.interface.items_tree
        if item.item_type == "SOCKET"
        and item.in_out == "OUTPUT"
    ]
    assert surface_outputs == ["Color", "Roughness", "Alpha", "Normal"]
    assert not any(
        node.bl_idname == "ShaderNodeDisplacement"
        for node in surface_group.nodes
    )
    refine_group = load_geometry_node_group()
    assert not any(
        item.item_type == "SOCKET"
        and item.in_out == "INPUT"
        and item.name == "Background"
        for item in refine_group.interface.items_tree
    )
    assert not any(
        node.bl_idname in forbidden_outer_nodes
        for node in surface_material.node_tree.nodes
    )
    assert all(
        node.extension == "EXTEND"
        for node in surface_material.node_tree.nodes
        if node.bl_idname == "ShaderNodeTexImage"
    )
    modifier = add_depth_plane_processing(
        obj,
        depth,
        foreground,
        material,
        animate_frame=True,
    )
    set_modifier_input(
        modifier,
        modifier_input_identifier(group, "Depth Scale"),
        -2.0,
    )
    for name, expected in (
        ("Frame Start", 1),
        ("Frames", 1),
        ("Cyclic", False),
    ):
        identifier = modifier_input_identifier(group, name)
        assert get_modifier_input(modifier, identifier) == expected
    assert not any(
        node.bl_idname == "GeometryNodeInputSceneTime"
        for node in group.nodes
    )
    frame_identifier = modifier_input_identifier(group, "Frame")
    frame_path = modifier_input_driver_path(modifier, frame_identifier)
    frame_driver = next(
        fcurve
        for fcurve in obj.animation_data.drivers
        if fcurve.data_path == frame_path
    )
    assert frame_driver.driver.type == "SCRIPTED"
    assert frame_driver.driver.expression == "frame"
    image_textures = [
        node
        for node in group.nodes
        if node.bl_idname == "GeometryNodeImageTexture"
    ]
    assert len(image_textures) == 2
    assert all(node.inputs["Frame"].is_linked for node in image_textures)
    assert image_textures[0].inputs["Frame"].links[0].from_socket == (
        image_textures[1].inputs["Frame"].links[0].from_socket
    )
    floored_modulo = [
        node
        for node in group.nodes
        if node.bl_idname == "ShaderNodeMath"
        and node.operation == "FLOORED_MODULO"
    ]
    assert len(floored_modulo) == 1
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    result = bpy.ops.object.modifier_apply(modifier=modifier.name)
    assert result == {"FINISHED"}
    assert not obj.modifiers
    assert obj.type == "MESH"
    assert len(obj.data.vertices) > 4
    for name in attribute_names:
        assert obj.data.attributes[name].domain == "POINT"
    assert all(polygon.use_smooth for polygon in obj.data.polygons)
    assert material in obj.data.materials[:]
    print("Depth plane asset validation passed")


if __name__ == "__main__":
    main()