icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add tests for foreground preservation and refactor geometry node asset

Browse Source 
- Implement test to ensure foreground preserves source RGB and Ben2 alpha in WorkerTest.
- Remove outdated tests related to background distance and original background depth.
- Update geometry node asset to reflect new naming conventions and structure.
- Introduce validation for depth plane asset, ensuring correct input/output configurations and attribute handling.
- Adjust confidence and edge parameters in geometry node asset to support a wider range of values.
This commit is contained in:
马南
2026-08-11 23:25:51 +08:00
parent 07d836d9ad
commit 7cd5681a8b
15 changed files with 1296 additions and 660 deletions
Download Patch File Download Diff File Expand all files Collapse all files

462
tools/build_geometry_nodes_asset.py
View File

@@ -3,7 +3,12 @@ from pathlib import Path
import bpy
GROUP_NAME = "O Refine DA3 Mesh"
GROUP_NAME = "O Refine DA3 Surface"
DEPTH_PLANE_GROUP_NAME = "O DA3 Depth Plane"
REMOVE_BACKGROUND_ATTRIBUTE = "o_remove_background"
MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement"
BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth"
DEPTH_SCALE_ATTRIBUTE = "o_depth_scale"
OUTPUT_PATH = (
Path(__file__).parents[1]
/ "src"
@@ -356,7 +361,13 @@ def build_group():
group = bpy.data.node_groups.new(GROUP_NAME, "GeometryNodeTree")
group.is_modifier = True
interface = group.interface
new_socket(interface, "Geometry", "INPUT", "NodeSocketGeometry")
geometry_input_socket = new_socket(
interface,
"Geometry",
"INPUT",
"NodeSocketGeometry",
)
geometry_input_socket.structure_type = "SINGLE"
resolution = new_socket(
interface,
@@ -400,7 +411,7 @@ def build_group():
)
confidence.default_value = 0.0
confidence.min_value = 0.0
confidence.max_value = 1.0
confidence.max_value = 100.0
confidence.subtype = "PERCENTAGE"
edge = new_socket(
interface,
@@ -409,9 +420,9 @@ def build_group():
"NodeSocketFloat",
parent=culling_panel,
)
edge.default_value = 0.1
edge.default_value = 10.0
edge.min_value = 0.0
edge.max_value = 1.0
edge.max_value = 100.0
edge.subtype = "PERCENTAGE"
background = new_socket(
interface,
@@ -459,6 +470,10 @@ def build_group():
confidence_attribute.data_type = "FLOAT"
confidence_attribute.inputs["Name"].default_value = "o_face_confidence"
confidence_attribute.location = (-1800.0, 600.0)
confidence_scale = nodes.new("ShaderNodeMath")
confidence_scale.operation = "MULTIPLY"
confidence_scale.inputs[1].default_value = 0.01
confidence_scale.location = (-1600.0, 800.0)
confidence_compare = new_compare(
nodes,
"FLOAT",
@@ -472,6 +487,10 @@ def build_group():
)
links.new(
confidence_input.outputs["Confidence"],
confidence_scale.inputs[0],
)
links.new(
confidence_scale.outputs[0],
input_socket(confidence_compare, "B", 1),
)
@@ -479,6 +498,10 @@ def build_group():
edge_attribute.data_type = "FLOAT"
edge_attribute.inputs["Name"].default_value = "o_edge_stretch"
edge_attribute.location = (-1800.0, 350.0)
edge_scale = nodes.new("ShaderNodeMath")
edge_scale.operation = "MULTIPLY"
edge_scale.inputs[1].default_value = 0.01
edge_scale.location = (-1600.0, -50.0)
edge_maximum = nodes.new("ShaderNodeMath")
edge_maximum.operation = "MAXIMUM"
edge_maximum.inputs[1].default_value = 0.0001
@@ -494,7 +517,8 @@ def build_group():
(-1200.0, 350.0),
1.0,
)
links.new(stretch_input.outputs["Stretch"], edge_maximum.inputs[0])
links.new(stretch_input.outputs["Stretch"], edge_scale.inputs[0])
links.new(edge_scale.outputs[0], edge_maximum.inputs[0])
links.new(edge_maximum.outputs[0], edge_reciprocal.inputs[1])
links.new(edge_attribute.outputs["Attribute"], input_socket(edge_compare, "A", 0))
links.new(edge_reciprocal.outputs[0], input_socket(edge_compare, "B", 1))
@@ -526,7 +550,7 @@ def build_group():
0.999999,
)
links.new(
stretch_input.outputs["Stretch"],
edge_scale.outputs[0],
input_socket(full_stretch, "A", 0),
)
culling_selection = nodes.new("FunctionNodeBooleanMath")
@@ -621,10 +645,434 @@ def build_group():
return group
def build_sequence_frame(nodes, links, scene_frame, frames_input, y):
relative_frame = nodes.new("ShaderNodeMath")
relative_frame.operation = "SUBTRACT"
relative_frame.location = (300.0, y)
links.new(scene_frame, relative_frame.inputs[0])
links.new(frames_input.outputs["Frame Start"], relative_frame.inputs[1])
safe_frames = nodes.new("ShaderNodeMath")
safe_frames.operation = "MAXIMUM"
safe_frames.inputs[1].default_value = 1.0
safe_frames.location = (300.0, y - 200.0)
links.new(frames_input.outputs["Frames"], safe_frames.inputs[0])
first_frame = nodes.new("ShaderNodeMath")
first_frame.operation = "ADD"
first_frame.inputs[1].default_value = 1.0
first_frame.location = (500.0, y)
links.new(relative_frame.outputs[0], first_frame.inputs[0])
clamp_minimum = nodes.new("ShaderNodeMath")
clamp_minimum.operation = "MAXIMUM"
clamp_minimum.inputs[1].default_value = 1.0
clamp_minimum.location = (700.0, y)
links.new(first_frame.outputs[0], clamp_minimum.inputs[0])
clamp_maximum = nodes.new("ShaderNodeMath")
clamp_maximum.operation = "MINIMUM"
clamp_maximum.location = (900.0, y)
links.new(clamp_minimum.outputs[0], clamp_maximum.inputs[0])
links.new(safe_frames.outputs[0], clamp_maximum.inputs[1])
cyclic_frame = nodes.new("ShaderNodeMath")
cyclic_frame.operation = "FLOORED_MODULO"
cyclic_frame.location = (700.0, y - 200.0)
links.new(relative_frame.outputs[0], cyclic_frame.inputs[0])
links.new(safe_frames.outputs[0], cyclic_frame.inputs[1])
cyclic_first_frame = nodes.new("ShaderNodeMath")
cyclic_first_frame.operation = "ADD"
cyclic_first_frame.inputs[1].default_value = 1.0
cyclic_first_frame.location = (900.0, y - 200.0)
links.new(cyclic_frame.outputs[0], cyclic_first_frame.inputs[0])
frame_switch = nodes.new("GeometryNodeSwitch")
frame_switch.input_type = "INT"
frame_switch.location = (1100.0, y)
links.new(frames_input.outputs["Cyclic"], frame_switch.inputs["Switch"])
links.new(clamp_maximum.outputs[0], frame_switch.inputs["False"])
links.new(cyclic_first_frame.outputs[0], frame_switch.inputs["True"])
return frame_switch.outputs["Output"]
def build_depth_plane_group():
group = bpy.data.node_groups.new(DEPTH_PLANE_GROUP_NAME, "GeometryNodeTree")
group.is_modifier = True
interface = group.interface
geometry_input_socket = new_socket(
interface,
"Geometry",
"INPUT",
"NodeSocketGeometry",
)
geometry_input_socket.structure_type = "SINGLE"
displacement = new_socket(
interface,
"Displacement",
"INPUT",
"NodeSocketMenu",
)
displacement_strength = new_socket(
interface,
"Depth Scale",
"INPUT",
"NodeSocketFloat",
)
displacement_strength.default_value = 1.0
displacement_strength.min_value = -10000.0
displacement_strength.max_value = 10000.0
displacement_strength.structure_type = "SINGLE"
subdivision = new_socket(
interface,
"Subdivision",
"INPUT",
"NodeSocketInt",
)
subdivision.default_value = 5
subdivision.min_value = 0
subdivision.max_value = 10
background_panel = interface.new_panel(
name="Remove Background",
default_closed=True,
)
remove_background = new_socket(
interface,
"Remove Background",
"INPUT",
"NodeSocketBool",
parent=background_panel,
)
remove_background.default_value = False
remove_background.is_panel_toggle = True
background_depth = new_socket(
interface,
"Background Depth",
"INPUT",
"NodeSocketFloat",
parent=background_panel,
)
background_depth.default_value = 0.0
background_depth.min_value = -10000.0
background_depth.max_value = 10000.0
background_depth.subtype = "DISTANCE"
data_panel = interface.new_panel(name="Data", default_closed=True)
depth_image = new_socket(
interface,
"Depth Image",
"INPUT",
"NodeSocketImage",
parent=data_panel,
)
foreground_image = new_socket(
interface,
"Foreground Image",
"INPUT",
"NodeSocketImage",
parent=data_panel,
)
material = new_socket(
interface,
"Material",
"INPUT",
"NodeSocketMaterial",
parent=data_panel,
)
frame_data = []
for name, socket_type, default_value in (
("Frame Start", "NodeSocketInt", 1),
("Frames", "NodeSocketInt", 1),
("Cyclic", "NodeSocketBool", False),
):
socket = new_socket(
interface,
name,
"INPUT",
socket_type,
parent=data_panel,
)
socket.default_value = default_value
if name == "Frames":
socket.min_value = 1
frame_data.append(socket)
depth_image.hide_in_modifier = True
foreground_image.hide_in_modifier = True
material.hide_in_modifier = True
for socket in frame_data:
socket.hide_in_modifier = True
new_socket(interface, "Geometry", "OUTPUT", "NodeSocketGeometry")
nodes = group.nodes
links = group.links
bounds_input = new_group_input(nodes, (-2200.0, 700.0), {"Geometry"})
grid_input = new_group_input(nodes, (-2200.0, 150.0), {"Subdivision"})
images_input = new_group_input(
nodes,
(100.0, 1250.0),
{"Depth Image", "Foreground Image"},
)
frames_input = new_group_input(
nodes,
(50.0, 1750.0),
{
"Frame Start",
"Frames",
"Cyclic",
},
)
background_input = new_group_input(
nodes,
(350.0, 250.0),
{"Remove Background", "Background Depth"},
)
displacement_input = new_group_input(
nodes,
(950.0, -150.0),
{"Displacement", "Depth Scale"},
)
material_input = new_group_input(nodes, (1450.0, 250.0), {"Material"})
attributes_input = new_group_input(
nodes,
(2150.0, 50.0),
{"Background Depth"},
)
group_output = nodes.new("NodeGroupOutput")
group_output.location = (3050.0, 650.0)
bounds = nodes.new("GeometryNodeBoundBox")
bounds.location = (-1950.0, 700.0)
links.new(bounds_input.outputs["Geometry"], bounds.inputs["Geometry"])
dimensions = nodes.new("ShaderNodeVectorMath")
dimensions.operation = "SUBTRACT"
dimensions.location = (-1750.0, 700.0)
links.new(bounds.outputs["Max"], dimensions.inputs[0])
links.new(bounds.outputs["Min"], dimensions.inputs[1])
separate_dimensions = nodes.new("ShaderNodeSeparateXYZ")
separate_dimensions.location = (-1550.0, 700.0)
links.new(dimensions.outputs["Vector"], separate_dimensions.inputs[0])
short_segments = nodes.new("ShaderNodeMath")
short_segments.operation = "POWER"
short_segments.inputs[0].default_value = 2.0
short_segments.location = (-1950.0, 150.0)
links.new(grid_input.outputs["Subdivision"], short_segments.inputs[1])
vertex_counts = []
for index, (axis, other_axis) in enumerate((("X", "Y"), ("Y", "X"))):
y = 500.0 - index * 300.0
ratio = nodes.new("ShaderNodeMath")
ratio.operation = "DIVIDE"
ratio.location = (-1350.0, y)
links.new(separate_dimensions.outputs[axis], ratio.inputs[0])
links.new(separate_dimensions.outputs[other_axis], ratio.inputs[1])
at_least_one = nodes.new("ShaderNodeMath")
at_least_one.operation = "MAXIMUM"
at_least_one.inputs[1].default_value = 1.0
at_least_one.location = (-1150.0, y)
links.new(ratio.outputs[0], at_least_one.inputs[0])
scaled_segments = nodes.new("ShaderNodeMath")
scaled_segments.operation = "MULTIPLY"
scaled_segments.location = (-950.0, y)
links.new(short_segments.outputs[0], scaled_segments.inputs[0])
links.new(at_least_one.outputs[0], scaled_segments.inputs[1])
rounded_segments = nodes.new("ShaderNodeMath")
rounded_segments.operation = "ROUND"
rounded_segments.location = (-750.0, y)
links.new(scaled_segments.outputs[0], rounded_segments.inputs[0])
vertices = nodes.new("ShaderNodeMath")
vertices.operation = "ADD"
vertices.inputs[1].default_value = 1.0
vertices.location = (-550.0, y)
links.new(rounded_segments.outputs[0], vertices.inputs[0])
vertex_counts.append(vertices.outputs[0])
grid = nodes.new("GeometryNodeMeshGrid")
grid.location = (-300.0, 650.0)
links.new(separate_dimensions.outputs["X"], grid.inputs["Size X"])
links.new(separate_dimensions.outputs["Y"], grid.inputs["Size Y"])
links.new(vertex_counts[0], grid.inputs["Vertices X"])
links.new(vertex_counts[1], grid.inputs["Vertices Y"])
store_uv = nodes.new("GeometryNodeStoreNamedAttribute")
store_uv.data_type = "FLOAT2"
store_uv.domain = "CORNER"
store_uv.inputs["Name"].default_value = "UVMap"
store_uv.location = (-50.0, 650.0)
links.new(grid.outputs["Mesh"], store_uv.inputs["Geometry"])
links.new(grid.outputs["UV Map"], store_uv.inputs["Value"])
center = nodes.new("ShaderNodeVectorMath")
center.operation = "ADD"
center.location = (-550.0, 950.0)
links.new(bounds.outputs["Min"], center.inputs[0])
links.new(bounds.outputs["Max"], center.inputs[1])
half_center = nodes.new("ShaderNodeVectorMath")
half_center.operation = "SCALE"
half_center.inputs[3].default_value = 0.5
half_center.location = (-300.0, 950.0)
links.new(center.outputs["Vector"], half_center.inputs[0])
transform_grid = nodes.new("GeometryNodeTransform")
transform_grid.location = (200.0, 650.0)
links.new(store_uv.outputs["Geometry"], transform_grid.inputs["Geometry"])
links.new(half_center.outputs["Vector"], transform_grid.inputs["Translation"])
depth_texture = nodes.new("GeometryNodeImageTexture")
depth_texture.interpolation = "Linear"
depth_texture.extension = "EXTEND"
depth_texture.location = (350.0, 1150.0)
foreground_texture = nodes.new("GeometryNodeImageTexture")
foreground_texture.interpolation = "Closest"
foreground_texture.extension = "EXTEND"
foreground_texture.location = (350.0, 900.0)
links.new(images_input.outputs["Depth Image"], depth_texture.inputs["Image"])
links.new(
images_input.outputs["Foreground Image"],
foreground_texture.inputs["Image"],
)
links.new(grid.outputs["UV Map"], depth_texture.inputs["Vector"])
links.new(grid.outputs["UV Map"], foreground_texture.inputs["Vector"])
scene_time = nodes.new("GeometryNodeInputSceneTime")
scene_time.location = (50.0, 2200.0)
sequence_frame = build_sequence_frame(
nodes,
links,
scene_time.outputs["Frame"],
frames_input,
2200.0,
)
links.new(sequence_frame, depth_texture.inputs["Frame"])
links.new(sequence_frame, foreground_texture.inputs["Frame"])
clamp_foreground = nodes.new("ShaderNodeMath")
clamp_foreground.operation = "MAXIMUM"
clamp_foreground.location = (600.0, 650.0)
links.new(depth_texture.outputs["Color"], clamp_foreground.inputs[0])
links.new(
background_input.outputs["Background Depth"],
clamp_foreground.inputs[1],
)
foreground_mask = new_compare(
nodes,
"FLOAT",
"GREATER_THAN",
(600.0, 900.0),
0.5,
)
links.new(
foreground_texture.outputs["Alpha"],
input_socket(foreground_mask, "A", 0),
)
masked_depth = nodes.new("ShaderNodeMix")
masked_depth.data_type = "FLOAT"
masked_depth.location = (800.0, 650.0)
links.new(foreground_mask.outputs["Result"], masked_depth.inputs[0])
links.new(
background_input.outputs["Background Depth"],
masked_depth.inputs[2],
)
links.new(clamp_foreground.outputs[0], masked_depth.inputs[3])
selected_depth = nodes.new("ShaderNodeMix")
selected_depth.data_type = "FLOAT"
selected_depth.location = (1000.0, 650.0)
remove_background_value = background_input.outputs["Remove Background"]
links.new(remove_background_value, selected_depth.inputs[0])
links.new(depth_texture.outputs["Color"], selected_depth.inputs[2])
links.new(masked_depth.outputs["Result"], selected_depth.inputs[3])
displacement_menu = nodes.new("GeometryNodeMenuSwitch")
displacement_menu.data_type = "FLOAT"
displacement_menu.enum_items.clear()
displacement_menu.enum_items.new("Geometry")
displacement_menu.enum_items.new("Material")
displacement_menu.inputs["Geometry"].default_value = 1.0
displacement_menu.inputs["Material"].default_value = 0.0
displacement_menu.location = (1150.0, 200.0)
links.new(
displacement_input.outputs["Displacement"],
displacement_menu.inputs["Menu"],
)
displacement.default_value = "Geometry"
geometry_mode = displacement_menu.outputs[0]
material_mode = nodes.new("ShaderNodeMath")
material_mode.operation = "SUBTRACT"
material_mode.inputs[0].default_value = 1.0
material_mode.location = (1350.0, 100.0)
links.new(geometry_mode, material_mode.inputs[1])
displacement_scale = nodes.new("ShaderNodeMath")
displacement_scale.operation = "MULTIPLY"
displacement_scale.location = (1200.0, 650.0)
links.new(selected_depth.outputs["Result"], displacement_scale.inputs[0])
geometry_strength = nodes.new("ShaderNodeMath")
geometry_strength.operation = "MULTIPLY"
geometry_strength.location = (1150.0, 450.0)
links.new(geometry_mode, geometry_strength.inputs[0])
links.new(displacement_input.outputs["Depth Scale"], geometry_strength.inputs[1])
links.new(geometry_strength.outputs[0], displacement_scale.inputs[1])
normal = nodes.new("GeometryNodeInputNormal")
normal.location = (1200.0, 950.0)
offset = nodes.new("ShaderNodeVectorMath")
offset.operation = "SCALE"
offset.location = (1400.0, 750.0)
links.new(normal.outputs["Normal"], offset.inputs[0])
links.new(displacement_scale.outputs[0], offset.inputs[3])
set_position = nodes.new("GeometryNodeSetPosition")
set_position.location = (1450.0, 650.0)
links.new(transform_grid.outputs["Geometry"], set_position.inputs["Geometry"])
links.new(offset.outputs["Vector"], set_position.inputs["Offset"])
set_material = nodes.new("GeometryNodeSetMaterial")
set_material.location = (1650.0, 650.0)
links.new(set_position.outputs["Geometry"], set_material.inputs["Geometry"])
links.new(material_input.outputs["Material"], set_material.inputs["Material"])
shade_smooth = nodes.new("GeometryNodeSetShadeSmooth")
shade_smooth.domain = "FACE"
shade_smooth.inputs["Shade Smooth"].default_value = True
shade_smooth.location = (1850.0, 650.0)
links.new(set_material.outputs["Geometry"], shade_smooth.inputs["Geometry"])
attribute_values = (
(
REMOVE_BACKGROUND_ATTRIBUTE,
"BOOLEAN",
remove_background_value,
),
(
MATERIAL_DISPLACEMENT_ATTRIBUTE,
"BOOLEAN",
material_mode.outputs[0],
),
(
BACKGROUND_DEPTH_ATTRIBUTE,
"FLOAT",
attributes_input.outputs["Background Depth"],
),
(
DEPTH_SCALE_ATTRIBUTE,
"FLOAT",
displacement_input.outputs["Depth Scale"],
),
)
geometry = shade_smooth.outputs["Geometry"]
x = 2100.0
for name, data_type, value in attribute_values:
store = nodes.new("GeometryNodeStoreNamedAttribute")
store.data_type = data_type
store.domain = "POINT"
store.inputs["Name"].default_value = name
store.location = (x, 650.0)
links.new(geometry, store.inputs["Geometry"])
links.new(value, store.inputs["Value"])
geometry = store.outputs["Geometry"]
x += 200.0
links.new(geometry, group_output.inputs["Geometry"])
group.color_tag = "GEOMETRY"
group.description = (
"Build an aspect-adaptive DA3 depth grid, clamp its foreground, "
"and pass background and displacement controls to the material."
)
group.use_fake_user = True
return group
def main():
bpy.ops.wm.read_factory_settings(use_empty=True)
bpy.context.preferences.filepaths.save_version = 0
build_group()
build_depth_plane_group()
OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
bpy.ops.wm.save_as_mainfile(filepath=str(OUTPUT_PATH), check_existing=False)