AnyImage/tools/nodes/groups/image_depth_plane.py

"""Build the camera-projected Depth 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
from ..common.smoothing import expand_smooth
from ..common.depth_surface import (
    project_depth_surface, thicken_depth_surface, remove_surface_attributes,
    split_depth_surface, store_surface_camera,
)

DEPTH_PLANE_NODE_GROUP_NAME = "O Image Depth Plane"


def build_image_depth_plane_group(image_plane_group):
    group = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME)
    if group is not None:
        return group
    group = bpy.data.node_groups.new(
        DEPTH_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.0
    thickness.min_value = 0.0
    thickness.subtype = "DISTANCE"
    thickness.description = "Inward shell distance"
    depth_scale = _interface_socket(group, "Depth Scale", "INPUT", "NodeSocketFloat")
    depth_scale.default_value = 1.0
    depth_scale.min_value = 0.0
    iterations = _interface_socket(group, "Smooth", "INPUT", "NodeSocketInt")
    iterations.default_value = 0
    iterations.min_value, iterations.max_value = 0, 20
    valid_only = _interface_socket(group, "Valid Only", "INPUT", "NodeSocketBool")
    valid_only.default_value = True
    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"
    normal_smooth = _interface_socket(
        group, "Normal Smooth", "INPUT", "NodeSocketInt", parent=options,
    )
    normal_smooth.default_value, normal_smooth.min_value, normal_smooth.max_value = 50, 0, 50
    normal_smooth.description = "Smooth the shell thickness direction"
    smooth_weight = _interface_socket(group, "Smooth Weight", "INPUT", "NodeSocketFloat", parent=options)
    smooth_weight.default_value = 1.0
    smooth_weight.min_value, smooth_weight.max_value = 0.0, 1.0
    smooth_weight.subtype = "FACTOR"
    validity_threshold = _interface_socket(group, "Validity Threshold", "INPUT", "NodeSocketFloat", parent=options)
    validity_threshold.default_value = 0.9
    validity_threshold.min_value, validity_threshold.max_value = 0.0, 1.0
    validity_threshold.subtype = "FACTOR"
    float_input(group, "Split Threshold", 0.0, maximum=1, subtype="FACTOR", 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, links = group.nodes, group.links
    controls = _group_input(nodes, {"Geometry", "Subdivide"})
    plane = nodes.new("GeometryNodeGroup")
    plane.node_tree = image_plane_group
    plane.inputs["Thickness"].default_value = 0
    links.new(controls.outputs["Geometry"], plane.inputs["Geometry"])
    links.new(controls.outputs["Subdivide"], plane.inputs["Subdivide"])
    texture_input = _group_input(nodes, {"Depth Image"})
    scale_input = _group_input(nodes, {"Uniform Scale"})
    uv = nodes.new("GeometryNodeInputNamedAttribute")
    uv.data_type = "FLOAT_VECTOR"
    uv.inputs["Name"].default_value = "UVMap"
    texture = nodes.new("GeometryNodeImageTexture")
    texture.interpolation, texture.extension = "Linear", "EXTEND"
    links.new(uv.outputs["Attribute"], texture.inputs["Vector"])
    links.new(texture_input.outputs["Depth Image"], texture.inputs["Image"])
    threshold_input = _group_input(nodes, {"Validity Threshold"})
    channels = nodes.new("FunctionNodeSeparateColor")
    channels.mode = "RGB"
    links.new(texture.outputs["Color"], channels.inputs["Color"])
    invalid = nodes.new("FunctionNodeCompare")
    invalid.data_type, invalid.operation = "FLOAT", "LESS_THAN"
    links.new(threshold_input.outputs["Validity Threshold"], invalid.inputs["B"])
    links.new(channels.outputs["Alpha"], invalid.inputs["A"])
    points = nodes.new("GeometryNodeFieldOnDomain")
    points.data_type, points.domain = "BOOLEAN", "POINT"
    links.new(invalid.outputs["Result"], points.inputs["Value"])
    delete = nodes.new("GeometryNodeDeleteGeometry")
    delete.domain, delete.mode = "FACE", "ALL"
    links.new(plane.outputs["Geometry"], delete.inputs["Geometry"])
    links.new(points.outputs["Value"], delete.inputs["Selection"])
    cull_controls = _group_input(nodes, {"Valid Only"})
    cull = nodes.new("GeometryNodeSwitch")
    cull.input_type = "GEOMETRY"
    links.new(cull_controls.outputs["Valid Only"], cull.inputs["Switch"])
    links.new(delete.outputs["Geometry"], cull.inputs["True"])
    links.new(plane.outputs["Geometry"], cull.inputs["False"])
    split_input = _group_input(nodes, {"Split Threshold"})
    geometry = split_depth_surface(
        group, cull.outputs["Output"], texture_input.outputs["Depth Image"],
        split_input.outputs["Split Threshold"], scale_input.outputs["Uniform Scale"],
    )
    geometry, camera = store_surface_camera(
        group, geometry, texture.outputs["Color"], split_input.outputs["Split Threshold"],
    )
    depth_input = _group_input(nodes, {"Reference Depth", "Depth Scale"})
    surface = project_depth_surface(
        group, geometry, camera,
        scale_input.outputs["Uniform Scale"], depth_input.outputs["Reference Depth"],
        depth_input.outputs["Depth Scale"],
    )
    shell_input = _group_input(nodes, {"Thickness", "Normal Smooth"})
    shell = thicken_depth_surface(
        group, surface, shell_input.outputs["Thickness"],
        normal_smooth=shell_input.outputs["Normal Smooth"],
    )
    positive = nodes.new("ShaderNodeMath")
    positive.operation = "GREATER_THAN"
    links.new(shell_input.outputs["Thickness"], positive.inputs[0])
    positive.inputs[1].default_value = 0
    choose = nodes.new("GeometryNodeSwitch")
    choose.input_type = "GEOMETRY"
    links.new(positive.outputs[0], choose.inputs["Switch"])
    links.new(surface, choose.inputs["False"])
    links.new(shell, choose.inputs["True"])
    clean = remove_surface_attributes(group, choose.outputs["Output"])
    geometry = expand_smooth(
        group, clean, iterations_name="Smooth", weight_name="Smooth Weight",
    )
    smooth = nodes.new("GeometryNodeSetShadeSmooth")
    smooth.domain = "FACE"
    links.new(geometry, smooth.inputs["Geometry"])
    output = nodes.new("NodeGroupOutput")
    links.new(smooth.outputs["Geometry"], output.inputs["Geometry"])
    group.color_tag = "GEOMETRY"
    group.description = "Build a camera-projected image shell."
    group.use_fake_user = True
    prepare_node_group(group)
    return group