AnyImage/tools/nodes/groups/image_relief_plane.py

"""Build the Relief Plane Geometry Nodes group."""

import bpy

from ..common.nodes import group_input as _group_input, interface_socket as _interface_socket
from ..common.nodes import prepare_node_group, float_input, vector_input
from ..common.reference_plane import reference_plane_depth
from ..common.smoothing import expand_smooth


RELIEF_PLANE_NODE_GROUP_NAME = "O Image Relief Plane"


def _math(nodes, operation):
    node = nodes.new("ShaderNodeMath")
    node.operation = operation
    return node


def build_image_relief_plane_group(image_plane_group):
    group = bpy.data.node_groups.get(RELIEF_PLANE_NODE_GROUP_NAME)
    if group is not None:
        return group
    group = bpy.data.node_groups.new(
        RELIEF_PLANE_NODE_GROUP_NAME,
        "GeometryNodeTree",
    )
    group.is_modifier = True
    geometry = _interface_socket(group, "Geometry", "INPUT", "NodeSocketGeometry")
    geometry.structure_type = "SINGLE"
    subdivide = _interface_socket(group, "Subdivide", "INPUT", "NodeSocketInt")
    subdivide.default_value = 6
    subdivide.min_value = 0
    subdivide.max_value = 10
    thickness = _interface_socket(group, "Thickness", "INPUT", "NodeSocketFloat")
    thickness.default_value = 0.5
    thickness.min_value = 0.0
    thickness.subtype = "DISTANCE"
    thickness.description = "Fixed relief base thickness"
    depth_scale = _interface_socket(group, "Depth Scale", "INPUT", "NodeSocketFloat")
    depth_scale.default_value = 0.5
    depth_scale.min_value = 0.0
    smooth = _interface_socket(group, "Smooth", "INPUT", "NodeSocketInt")
    smooth.default_value, smooth.min_value, smooth.max_value = 0, 0, 20
    options = group.interface.new_panel(name="Options", default_closed=True)
    base_plane_depth = _interface_socket(
        group,
        "Reference Depth",
        "INPUT",
        "NodeSocketFloat",
        parent=options,
    )
    base_plane_depth.default_value = 1.0
    base_plane_depth.min_value = 0.0
    base_plane_depth.subtype = "DISTANCE"
    float_input(group, "Smooth Weight", 1.0, maximum=1, subtype="FACTOR", parent=options)
    vector_input(group, "Depth Direction", (0, 0, 1), subtype="DIRECTION", parent=options)
    data = group.interface.new_panel(name="Data", default_closed=True)
    uniform_scale = _interface_socket(
        group,
        "Uniform Scale",
        "INPUT",
        "NodeSocketFloat",
        parent=data,
    )
    uniform_scale.default_value = 1.0
    uniform_scale.min_value = 0.0
    uniform_scale.hide_in_modifier = True
    depth_image = _interface_socket(
        group,
        "Depth Image",
        "INPUT",
        "NodeSocketImage",
        parent=data,
    )
    depth_image.hide_in_modifier = True
    _interface_socket(group, "Geometry", "OUTPUT", "NodeSocketGeometry")

    nodes = group.nodes
    links = group.links
    geometry_input = _group_input(
        nodes,
        {"Geometry", "Subdivide"},
    )
    shape_input = _group_input(
        nodes,
        {"Thickness", "Depth Scale"},
    )
    image_input = _group_input(
        nodes,
        {"Depth Image", "Uniform Scale"},
    )
    group_output = nodes.new("NodeGroupOutput")
    image_plane = nodes.new("GeometryNodeGroup")
    image_plane.node_tree = image_plane_group
    links.new(geometry_input.outputs["Geometry"], image_plane.inputs["Geometry"])
    links.new(geometry_input.outputs["Subdivide"], image_plane.inputs["Subdivide"])
    links.new(shape_input.outputs["Thickness"], image_plane.inputs["Thickness"])

    positive_half = _math(nodes, "MULTIPLY")
    positive_half.inputs[1].default_value = 0.5
    links.new(shape_input.outputs["Thickness"], positive_half.inputs[0])
    translation = nodes.new("ShaderNodeCombineXYZ")
    links.new(positive_half.outputs[0], translation.inputs["Z"])
    place_block = nodes.new("GeometryNodeTransform")
    links.new(image_plane.outputs["Geometry"], place_block.inputs["Geometry"])
    links.new(translation.outputs["Vector"], place_block.inputs["Translation"])

    uv = nodes.new("GeometryNodeInputNamedAttribute")
    uv.data_type = "FLOAT_VECTOR"
    uv.inputs["Name"].default_value = "UVMap"
    texture = nodes.new("GeometryNodeImageTexture")
    texture.interpolation = "Linear"
    texture.extension = "EXTEND"
    links.new(image_input.outputs["Depth Image"], texture.inputs["Image"])
    links.new(uv.outputs["Attribute"], texture.inputs["Vector"])
    depth_components = nodes.new("ShaderNodeSeparateXYZ")
    links.new(texture.outputs["Color"], depth_components.inputs["Vector"])
    metric_depth = _math(nodes, "MULTIPLY")
    links.new(depth_components.outputs["Z"], metric_depth.inputs[0])
    links.new(image_input.outputs["Uniform Scale"], metric_depth.inputs[1])
    relative_depth = _math(nodes, "SUBTRACT")
    reference_input = _group_input(nodes, {"Depth Direction", "Reference Depth"})
    plane_depth = reference_plane_depth(
        group, reference_input.outputs["Depth Direction"], reference_input.outputs["Reference Depth"],
    )
    links.new(plane_depth, relative_depth.inputs[0])
    links.new(metric_depth.outputs[0], relative_depth.inputs[1])
    full_offset = _math(nodes, "MULTIPLY")
    links.new(relative_depth.outputs[0], full_offset.inputs[0])
    links.new(shape_input.outputs["Depth Scale"], full_offset.inputs[1])

    inward_limit = _math(nodes, "MULTIPLY")
    inward_limit.inputs[1].default_value = -0.999
    links.new(shape_input.outputs["Thickness"], inward_limit.inputs[0])
    limited_offset = _math(nodes, "MAXIMUM")
    links.new(full_offset.outputs[0], limited_offset.inputs[0])
    links.new(inward_limit.outputs[0], limited_offset.inputs[1])
    has_thickness = _math(nodes, "GREATER_THAN")
    has_thickness.inputs[1].default_value = 1e-6
    links.new(shape_input.outputs["Thickness"], has_thickness.inputs[0])
    flat_offset = _math(nodes, "MAXIMUM")
    flat_offset.inputs[1].default_value = 0.0
    links.new(full_offset.outputs[0], flat_offset.inputs[0])
    safe_offset = nodes.new("GeometryNodeSwitch")
    safe_offset.input_type = "FLOAT"
    links.new(has_thickness.outputs[0], safe_offset.inputs["Switch"])
    links.new(flat_offset.outputs[0], safe_offset.inputs["False"])
    links.new(limited_offset.outputs[0], safe_offset.inputs["True"])

    position = nodes.new("GeometryNodeInputPosition")
    position_components = nodes.new("ShaderNodeSeparateXYZ")
    links.new(position.outputs["Position"], position_components.inputs["Vector"])
    layer_ratio = _math(nodes, "DIVIDE")
    links.new(position_components.outputs["Z"], layer_ratio.inputs[0])
    links.new(shape_input.outputs["Thickness"], layer_ratio.inputs[1])
    nonnegative_ratio = _math(nodes, "MAXIMUM")
    nonnegative_ratio.inputs[1].default_value = 0.0
    links.new(layer_ratio.outputs[0], nonnegative_ratio.inputs[0])
    clamped_ratio = _math(nodes, "MINIMUM")
    clamped_ratio.inputs[1].default_value = 1.0
    links.new(nonnegative_ratio.outputs[0], clamped_ratio.inputs[0])
    layer_weight = nodes.new("GeometryNodeSwitch")
    layer_weight.input_type = "FLOAT"
    layer_weight.inputs["False"].default_value = 1.0
    links.new(has_thickness.outputs[0], layer_weight.inputs["Switch"])
    links.new(clamped_ratio.outputs[0], layer_weight.inputs["True"])

    weighted_offset = _math(nodes, "MULTIPLY")
    links.new(safe_offset.outputs["Output"], weighted_offset.inputs[0])
    links.new(layer_weight.outputs["Output"], weighted_offset.inputs[1])
    offset = nodes.new("ShaderNodeCombineXYZ")
    links.new(weighted_offset.outputs[0], offset.inputs["Z"])
    set_position = nodes.new("GeometryNodeSetPosition")
    links.new(place_block.outputs["Geometry"], set_position.inputs["Geometry"])
    links.new(offset.outputs["Vector"], set_position.inputs["Offset"])
    shade_smooth = nodes.new("GeometryNodeSetShadeSmooth")
    shade_smooth.domain = "FACE"
    shade_smooth.inputs["Shade Smooth"].default_value = True
    geometry = expand_smooth(
        group, set_position.outputs["Geometry"], iterations_name="Smooth", weight_name="Smooth Weight",
    )
    links.new(geometry, shade_smooth.inputs["Geometry"])
    links.new(shade_smooth.outputs["Geometry"], group_output.inputs["Geometry"])
    group.color_tag = "GEOMETRY"
    group.description = (
        "Build a fixed-base image relief."
    )
    group.use_fake_user = True
    prepare_node_group(group)
    return group