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

BIN
src/snapform/assets/O_SnapForm.blend LFS
View File

Binary file not shown.

619
src/snapform/image_depth.py
View File

@@ -13,6 +13,14 @@ REFERENCE_DEPTH_KEY = "snapform_reference_depth"
ORIGINAL_DISPLACEMENT_STRENGTH = 1.0
ORIGINAL_DISPLACEMENT_MIDLEVEL = 0.0
BUMP_BASE_DISTANCE = 0.001
DEPTH_PLANE_NODE_GROUP_NAME = "O DA3 Depth Plane"
NODE_ASSET_NAME = "O_SnapForm.blend"
REMOVE_BACKGROUND_ATTRIBUTE = "o_remove_background"
MATERIAL_DISPLACEMENT_ATTRIBUTE = "o_material_displacement"
BACKGROUND_DEPTH_ATTRIBUTE = "o_background_depth"
DEPTH_SCALE_ATTRIBUTE = "o_depth_scale"
SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface"
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Material"
IMAGE_NAME_SUFFIXES = (
".jpeg",
".jpg",
@@ -430,6 +438,7 @@ def create_removed_background_image(
str(foreground_paths[0]),
check_existing=False,
)
result_image.alpha_mode = "CHANNEL_PACKED"
force_sequence_source(result_image)
except (OSError, RuntimeError) as error:
if "result_image" in locals():
@@ -471,6 +480,7 @@ def create_removed_background_image(
str(output_path),
check_existing=False,
)
result_image.alpha_mode = "CHANNEL_PACKED"
except RuntimeError as error:
raise RuntimeError("Unable to load the background result") from error
source_image = source_object.data
@@ -794,10 +804,13 @@ def create_plane_mesh(source_object, bounds=None, subdivision_level=5):
x_min, x_max, y_min, y_max = (
bounds if bounds is not None else plane_bounds(source_object)
)
x_segments, y_segments = grid_segments(
(x_min, x_max, y_min, y_max),
subdivision_level,
)
if subdivision_level is None:
x_segments, y_segments = 1, 1
else:
x_segments, y_segments = grid_segments(
(x_min, x_max, y_min, y_max),
subdivision_level,
)
vertices = []
uv_coordinates = []
for y_index in range(y_segments + 1):
@@ -836,11 +849,171 @@ def create_plane_mesh(source_object, bounds=None, subdivision_level=5):
return mesh
def material_node_group(depth_plane=False):
group_name = (
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME
if depth_plane
else SURFACE_MATERIAL_NODE_GROUP_NAME
)
existing = bpy.data.node_groups.get(group_name)
if existing is not None:
return existing
group = bpy.data.node_groups.new(group_name, "ShaderNodeTree")
interface = group.interface
for name, socket_type in (
("Color", "NodeSocketColor"),
("Alpha", "NodeSocketFloat"),
("Depth", "NodeSocketFloat"),
("Bump Scale", "NodeSocketFloat"),
):
socket = interface.new_socket(
name=name,
in_out="INPUT",
socket_type=socket_type,
)
if name == "Alpha":
socket.default_value = 1.0
socket.min_value = 0.0
socket.max_value = 1.0
socket.subtype = "FACTOR"
elif name == "Bump Scale":
socket.default_value = 1.0
socket.min_value = -10000.0
socket.max_value = 10000.0
for name, socket_type in (
("Color", "NodeSocketColor"),
("Roughness", "NodeSocketFloat"),
("Alpha", "NodeSocketFloat"),
("Normal", "NodeSocketVector"),
("Displacement", "NodeSocketVector"),
):
interface.new_socket(
name=name,
in_out="OUTPUT",
socket_type=socket_type,
)
nodes = group.nodes
links = group.links
group_input = nodes.new("NodeGroupInput")
group_input.location = (-500.0, 100.0)
group_output = nodes.new("NodeGroupOutput")
group_output.location = (300.0, 100.0)
roughness_curve = nodes.new("ShaderNodeFloatCurve")
roughness_curve.mapping.initialize()
curve_points = roughness_curve.mapping.curves[0].points
curve_points[0].location = (0.0, 0.5)
curve_points[-1].location = (1.0, 0.1)
roughness_curve.mapping.update()
roughness_curve.location = (-250.0, 350.0)
bump = nodes.new("ShaderNodeBump")
bump.location = (-250.0, 0.0)
displacement = nodes.new("ShaderNodeDisplacement")
displacement.inputs["Midlevel"].default_value = 0.0
displacement.inputs["Scale"].default_value = 1.0
displacement.location = (-250.0, -300.0)
processed_alpha = group_input.outputs["Alpha"]
processed_depth = group_input.outputs["Depth"]
if depth_plane:
group_input.location = (-1200.0, 100.0)
group_output.location = (1050.0, 100.0)
roughness_curve.location = (-300.0, 450.0)
bump.location = (750.0, 50.0)
displacement.location = (750.0, -300.0)
remove_attribute = nodes.new("ShaderNodeAttribute")
remove_attribute.attribute_type = "GEOMETRY"
remove_attribute.attribute_name = REMOVE_BACKGROUND_ATTRIBUTE
remove_attribute.location = (-950.0, 0.0)
background_attribute = nodes.new("ShaderNodeAttribute")
background_attribute.attribute_type = "GEOMETRY"
background_attribute.attribute_name = BACKGROUND_DEPTH_ATTRIBUTE
background_attribute.location = (-950.0, -500.0)
material_attribute = nodes.new("ShaderNodeAttribute")
material_attribute.attribute_type = "GEOMETRY"
material_attribute.attribute_name = MATERIAL_DISPLACEMENT_ATTRIBUTE
material_attribute.location = (-50.0, -650.0)
strength_attribute = nodes.new("ShaderNodeAttribute")
strength_attribute.attribute_type = "GEOMETRY"
strength_attribute.attribute_name = DEPTH_SCALE_ATTRIBUTE
strength_attribute.location = (200.0, -650.0)
alpha_mix = nodes.new("ShaderNodeMix")
alpha_mix.data_type = "FLOAT"
alpha_mix.inputs[2].default_value = 1.0
alpha_mix.location = (-650.0, 50.0)
foreground_mask = nodes.new("ShaderNodeMath")
foreground_mask.operation = "GREATER_THAN"
foreground_mask.inputs[1].default_value = 0.5
foreground_mask.location = (-950.0, -250.0)
clamp_foreground = nodes.new("ShaderNodeMath")
clamp_foreground.operation = "MAXIMUM"
clamp_foreground.location = (-650.0, -400.0)
masked_depth = nodes.new("ShaderNodeMix")
masked_depth.data_type = "FLOAT"
masked_depth.location = (-350.0, -350.0)
background_mix = nodes.new("ShaderNodeMix")
background_mix.data_type = "FLOAT"
background_mix.location = (-50.0, -350.0)
enabled_depth = nodes.new("ShaderNodeMath")
enabled_depth.operation = "MULTIPLY"
enabled_depth.location = (250.0, -350.0)
scaled_depth = nodes.new("ShaderNodeMath")
scaled_depth.operation = "MULTIPLY"
scaled_depth.location = (450.0, -350.0)
links.new(remove_attribute.outputs["Fac"], alpha_mix.inputs[0])
links.new(group_input.outputs["Alpha"], alpha_mix.inputs[3])
links.new(group_input.outputs["Alpha"], foreground_mask.inputs[0])
links.new(group_input.outputs["Depth"], clamp_foreground.inputs[0])
links.new(background_attribute.outputs["Fac"], clamp_foreground.inputs[1])
links.new(foreground_mask.outputs[0], masked_depth.inputs[0])
links.new(background_attribute.outputs["Fac"], masked_depth.inputs[2])
links.new(clamp_foreground.outputs[0], masked_depth.inputs[3])
links.new(remove_attribute.outputs["Fac"], background_mix.inputs[0])
links.new(group_input.outputs["Depth"], background_mix.inputs[2])
links.new(masked_depth.outputs["Result"], background_mix.inputs[3])
links.new(background_mix.outputs["Result"], enabled_depth.inputs[0])
links.new(material_attribute.outputs["Fac"], enabled_depth.inputs[1])
links.new(enabled_depth.outputs[0], scaled_depth.inputs[0])
links.new(strength_attribute.outputs["Fac"], scaled_depth.inputs[1])
processed_alpha = alpha_mix.outputs["Result"]
processed_depth = scaled_depth.outputs[0]
links.new(group_input.outputs["Color"], group_output.inputs["Color"])
links.new(group_input.outputs["Color"], roughness_curve.inputs["Value"])
links.new(
roughness_curve.outputs["Value"],
group_output.inputs["Roughness"],
)
links.new(processed_alpha, group_output.inputs["Alpha"])
links.new(processed_depth, bump.inputs["Height"])
bump_distance = nodes.new("ShaderNodeMath")
bump_distance.operation = "MULTIPLY"
bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE
bump_distance.location = (500.0, 100.0) if depth_plane else (-250.0, 150.0)
links.new(group_input.outputs["Bump Scale"], bump_distance.inputs[0])
links.new(bump_distance.outputs[0], bump.inputs["Distance"])
links.new(bump.outputs["Normal"], group_output.inputs["Normal"])
links.new(processed_depth, displacement.inputs["Height"])
links.new(
displacement.outputs["Displacement"],
group_output.inputs["Displacement"],
)
group.color_tag = "SHADER"
group.description = (
"Prepare DA3 depth-plane material outputs."
if depth_plane
else "Prepare DA3 surface material outputs."
)
group.use_fake_user = True
return group
def create_material(
name,
color_image,
depth_scale,
displacement_image=None,
foreground_image=None,
displacement_midlevel=0.0,
displacement_scale=1.0,
unshaded=False,
@@ -848,6 +1021,8 @@ def create_material(
frame_start=1,
color_frame_offset=0,
color_frame_duration=None,
foreground_frame_offset=0,
foreground_frame_duration=None,
depth_frame_offset=0,
depth_frame_duration=None,
use_cyclic=None,
@@ -860,13 +1035,17 @@ def create_material(
nodes.clear()
color_texture = nodes.new("ShaderNodeTexImage")
uv_map = nodes.new("ShaderNodeUVMap")
output = nodes.new("ShaderNodeOutputMaterial")
color_texture.image = color_image
uv_map.uv_map = "UVMap"
uv_map.location = (-1250.0, 350.0)
material_image = foreground_image or color_image
color_texture.image = material_image
color_texture.interpolation = "Linear"
color_texture.extension = "EXTEND"
color_texture.location = (-640.0, 180.0)
output.location = (460.0, 180.0)
color_texture.location = (-1000.0, 350.0)
output.location = (850.0, 350.0)
extracted_frames = frame_count if frame_count > 1 else None
color_frames = (
color_frame_duration
@@ -878,29 +1057,51 @@ def create_material(
if depth_frame_duration is not None
else extracted_frames
)
if frame_count > 1 or is_animated_image(color_image):
material_frames = color_frames
material_frame_offset = color_frame_offset
if foreground_image is not None:
material_frames = (
foreground_frame_duration
if foreground_frame_duration is not None
else extracted_frames
)
material_frame_offset = foreground_frame_offset
if frame_count > 1 or is_animated_image(material_image):
configure_image_user(
color_texture.image_user,
color_frames,
material_frames,
frame_start=frame_start,
frame_offset=color_frame_offset,
frame_offset=material_frame_offset,
use_cyclic=use_cyclic,
)
links = material.node_tree.links
links.new(uv_map.outputs["UV"], color_texture.inputs["Vector"])
alpha = color_texture.outputs["Alpha"]
material_inputs = nodes.new("ShaderNodeGroup")
is_depth_plane = displacement_image is not None
material_inputs.node_tree = material_node_group(is_depth_plane)
material_inputs.name = (
"O DA3 Depth Plane" if is_depth_plane else "O DA3 Surface"
)
material_inputs.label = material_inputs.name
material_inputs.location = (-300.0, 350.0)
material_inputs.inputs["Bump Scale"].default_value = depth_scale
links.new(color_texture.outputs["Color"], material_inputs.inputs["Color"])
links.new(alpha, material_inputs.inputs["Alpha"])
if unshaded:
emission = nodes.new("ShaderNodeEmission")
transparent = nodes.new("ShaderNodeBsdfTransparent")
mix_shader = nodes.new("ShaderNodeMixShader")
emission.location = (-100.0, 260.0)
transparent.location = (-100.0, 20.0)
mix_shader.location = (200.0, 180.0)
emission.location = (0.0, 550.0)
transparent.location = (0.0, 150.0)
mix_shader.location = (300.0, 350.0)
links.new(
color_texture.outputs["Color"],
material_inputs.outputs["Color"],
emission.inputs["Color"],
)
links.new(
color_texture.outputs["Alpha"],
material_inputs.outputs["Alpha"],
mix_shader.inputs["Fac"],
)
links.new(
@@ -917,39 +1118,22 @@ def create_material(
)
else:
shader = nodes.new("ShaderNodeBsdfPrincipled")
bump = nodes.new("ShaderNodeBump")
roughness_curve = nodes.new("ShaderNodeFloatCurve")
bump.inputs["Distance"].default_value = BUMP_BASE_DISTANCE * depth_scale
roughness_curve.mapping.initialize()
curve_points = roughness_curve.mapping.curves[0].points
curve_points[0].location = (0.0, 0.5)
curve_points[-1].location = (1.0, 0.1)
roughness_curve.mapping.update()
roughness_curve.location = (-300.0, 320.0)
bump.location = (-300.0, -80.0)
shader.location = (80.0, 180.0)
shader.location = (250.0, 400.0)
links.new(
color_texture.outputs["Color"],
material_inputs.outputs["Color"],
shader.inputs["Base Color"],
)
links.new(color_texture.outputs["Alpha"], shader.inputs["Alpha"])
links.new(color_texture.outputs["Color"], bump.inputs["Height"])
links.new(bump.outputs["Normal"], shader.inputs["Normal"])
links.new(material_inputs.outputs["Alpha"], shader.inputs["Alpha"])
links.new(material_inputs.outputs["Normal"], shader.inputs["Normal"])
links.new(
color_texture.outputs["Color"],
roughness_curve.inputs["Value"],
)
links.new(
roughness_curve.outputs["Value"],
material_inputs.outputs["Roughness"],
shader.inputs["Roughness"],
)
links.new(shader.outputs["BSDF"], output.inputs["Surface"])
if displacement_image is not None:
depth_texture = nodes.new("ShaderNodeTexImage")
displacement = nodes.new("ShaderNodeDisplacement")
depth_texture.image = displacement_image
depth_texture.interpolation = "Linear"
depth_texture.extension = "EXTEND"
@@ -961,21 +1145,30 @@ def create_material(
frame_offset=depth_frame_offset,
use_cyclic=use_cyclic,
)
displacement.inputs["Midlevel"].default_value = (
displacement_midlevel
)
displacement.inputs["Scale"].default_value = displacement_scale
depth_texture.location = (-300.0, -360.0)
displacement.location = (80.0, -260.0)
depth_texture.location = (-1000.0, -400.0)
links.new(uv_map.outputs["UV"], depth_texture.inputs["Vector"])
displacement_height = depth_texture.outputs["Color"]
if displacement_midlevel != 0.0 or displacement_scale != 1.0:
depth_adjust = nodes.new("ShaderNodeMath")
depth_adjust.operation = "SUBTRACT"
depth_adjust.inputs[1].default_value = displacement_midlevel
depth_adjust.location = (500.0, -550.0)
depth_scale_node = nodes.new("ShaderNodeMath")
depth_scale_node.operation = "MULTIPLY"
depth_scale_node.inputs[1].default_value = displacement_scale
depth_scale_node.location = (700.0, -550.0)
links.new(displacement_height, depth_adjust.inputs[0])
links.new(depth_adjust.outputs[0], depth_scale_node.inputs[0])
links.new(depth_scale_node.outputs[0], material_inputs.inputs["Depth"])
else:
links.new(displacement_height, material_inputs.inputs["Depth"])
links.new(
depth_texture.outputs["Color"],
displacement.inputs["Height"],
)
links.new(
displacement.outputs["Displacement"],
material_inputs.outputs["Displacement"],
output.inputs["Displacement"],
)
set_displacement_only(material)
else:
links.new(color_texture.outputs["Color"], material_inputs.inputs["Depth"])
return material
@@ -990,6 +1183,84 @@ def set_displacement_only(material):
return
def node_asset_path():
return Path(__file__).with_name("assets") / NODE_ASSET_NAME
def load_depth_plane_node_group():
existing = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME)
if existing is not None:
return existing
asset_path = node_asset_path()
if not asset_path.is_file():
raise RuntimeError(f"Missing Geometry Nodes asset: {asset_path}")
use_pack = bpy.app.version >= (5, 0, 0)
if not use_pack:
directory = str(asset_path) + "/NodeTree/"
result = bpy.ops.wm.append(
filepath=directory + DEPTH_PLANE_NODE_GROUP_NAME,
directory=directory,
filename=DEPTH_PLANE_NODE_GROUP_NAME,
link=False,
do_reuse_local_id=True,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to append the Geometry Nodes asset")
node_group = bpy.data.node_groups.get(DEPTH_PLANE_NODE_GROUP_NAME)
if node_group is None:
raise RuntimeError("Appended Geometry Nodes asset is unavailable")
node_group.use_fake_user = True
return node_group
with bpy.data.libraries.load(str(asset_path), link=True) as (source, target):
if DEPTH_PLANE_NODE_GROUP_NAME not in source.node_groups:
raise RuntimeError(
"Geometry Nodes asset is missing "
f"{DEPTH_PLANE_NODE_GROUP_NAME}"
)
target.node_groups = [DEPTH_PLANE_NODE_GROUP_NAME]
linked_group = target.node_groups[0]
node_group = bpy.data.pack_linked_ids_hierarchy(root_id=linked_group)
if linked_group != node_group and linked_group.users == 0:
bpy.data.node_groups.remove(linked_group)
node_group.use_fake_user = True
return node_group
def modifier_input_identifier(node_group, name):
for item in node_group.interface.items_tree:
if item.item_type != "SOCKET" or item.in_out != "INPUT":
continue
if item.name == name:
return item.identifier
raise RuntimeError(f"Geometry Nodes input is missing: {name}")
def add_depth_plane_processing(
plane_object,
depth_image,
foreground_image,
material,
frame_start=1,
frame_duration=1,
use_cyclic=False,
):
node_group = load_depth_plane_node_group()
modifier = plane_object.modifiers.new(name=node_group.name, type="NODES")
modifier.node_group = node_group
modifier.show_group_selector = False
for name, value in (
("Depth Image", depth_image),
("Foreground Image", foreground_image),
("Material", material),
("Frame Start", frame_start),
("Frames", max(int(frame_duration or 1), 1)),
("Cyclic", use_cyclic),
):
identifier = modifier_input_identifier(node_group, name)
modifier[identifier] = value
return modifier
def finish_depth_plane(
context,
source_object,
@@ -1010,11 +1281,10 @@ def finish_depth_plane(
def build_image_depth_plane(
context,
source_object,
subdivision_level,
delete_source,
displacement_via="MODIFIER",
unshaded=False,
color_image=None,
foreground_image=None,
depth_frames_directory=None,
depth_image=None,
):
@@ -1046,48 +1316,21 @@ def build_image_depth_plane(
depth_source_image = frames_image
elif not is_animated_image(source_object.data):
depth_source_image = copy_depth_image(source_object)
mesh = create_plane_mesh(
source_object,
subdivision_level=subdivision_level,
foreground_image = (
foreground_image or color_image or source_object.data
)
mesh = create_plane_mesh(source_object, subdivision_level=None)
plane_object = bpy.data.objects.new(source_name, mesh)
context.collection.objects.link(plane_object)
plane_object.matrix_world = source_object.matrix_world.copy()
subdivision = plane_object.modifiers.new("Subdivision", "SUBSURF")
subdivision.subdivision_type = "SIMPLE"
subdivision.levels = 1
subdivision.render_levels = 2
if displacement_via == "MODIFIER":
texture = bpy.data.textures.new(
depth_data_name(source_object),
type="IMAGE",
)
texture.image = depth_source_image
texture.extension = "EXTEND"
if uses_extracted or is_animated_image(source_object.data):
configure_image_user(
texture.image_user,
texture_frame_count,
frame_start=source_frame_start,
frame_offset=texture_frame_offset,
use_cyclic=source_use_cyclic,
)
displace = plane_object.modifiers.new("Displace", "DISPLACE")
displace.texture = texture
displace.texture_coords = "UV"
displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL
depth_scale = abs(float(plane_object.scale.z))
displacement_image = (
depth_source_image if displacement_via == "MATERIAL" else None
)
material = create_material(
source_name,
color_image if color_image is not None else source_object.data,
depth_scale,
displacement_image=displacement_image,
displacement_image=depth_source_image,
foreground_image=foreground_image,
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
unshaded=unshaded,
@@ -1095,11 +1338,22 @@ def build_image_depth_plane(
frame_start=source_frame_start,
color_frame_offset=texture_frame_offset,
color_frame_duration=texture_frame_count,
foreground_frame_offset=0,
foreground_frame_duration=frame_count,
depth_frame_offset=texture_frame_offset,
depth_frame_duration=texture_frame_count,
use_cyclic=source_use_cyclic,
)
mesh.materials.append(material)
add_depth_plane_processing(
plane_object,
depth_source_image,
foreground_image,
material,
frame_start=source_frame_start,
frame_duration=frame_count,
use_cyclic=source_use_cyclic,
)
plane_object["snapform_frame_count"] = frame_count
return finish_depth_plane(
context,
@@ -1113,12 +1367,11 @@ def build_depth_plane(
context,
source_object,
npz_path,
subdivision_level,
delete_source,
displacement_via="MODIFIER",
unshaded=False,
depth_name="",
color_image=None,
foreground_image=None,
depth_frames_directory=None,
):
source_name = source_object.name
@@ -1136,13 +1389,16 @@ def build_depth_plane(
)
depth_name = depth_name or depth_data_name(source_object)
frame_count = len(depth_arrays)
foreground_image = (
foreground_image or color_image or source_object.data
)
(
source_frame_start,
source_frame_offset,
source_frame_duration,
source_use_cyclic,
) = image_user_settings(source_object)
mesh = create_plane_mesh(source_object, bounds, subdivision_level)
mesh = create_plane_mesh(source_object, bounds, None)
plane_object = bpy.data.objects.new(
source_name,
mesh,
@@ -1170,39 +1426,12 @@ def build_depth_plane(
depth_frames_directory,
)
subdivision = plane_object.modifiers.new("Subdivision", "SUBSURF")
subdivision.subdivision_type = "SIMPLE"
subdivision.levels = 1
subdivision.render_levels = 2
if displacement_via == "MODIFIER":
texture = bpy.data.textures.new(
depth_name,
type="IMAGE",
)
texture.image = depth_image
texture.extension = "EXTEND"
if frame_count > 1:
configure_image_user(
texture.image_user,
frame_count,
frame_start=source_frame_start,
use_cyclic=source_use_cyclic,
)
displace = plane_object.modifiers.new("Displace", "DISPLACE")
displace.texture = texture
displace.texture_coords = "UV"
displace.direction = "NORMAL"
displace.strength = DISPLACEMENT_STRENGTH
displace.mid_level = DISPLACEMENT_MIDLEVEL
material = create_material(
source_name,
color_image if color_image is not None else source_object.data,
depth_scale,
displacement_image=(
depth_image if displacement_via == "MATERIAL" else None
),
displacement_image=depth_image,
foreground_image=foreground_image,
displacement_midlevel=DISPLACEMENT_MIDLEVEL,
displacement_scale=DISPLACEMENT_STRENGTH,
unshaded=unshaded,
@@ -1212,11 +1441,22 @@ def build_depth_plane(
color_frame_duration=(
frame_count if color_image is not None else source_frame_duration
),
foreground_frame_offset=0,
foreground_frame_duration=frame_count,
depth_frame_offset=0,
depth_frame_duration=frame_count,
use_cyclic=source_use_cyclic,
)
mesh.materials.append(material)
add_depth_plane_processing(
plane_object,
depth_image,
foreground_image,
material,
frame_start=source_frame_start,
frame_duration=frame_count,
use_cyclic=source_use_cyclic,
)
plane_object["snapform_reference_depth"] = reference_depth
plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0]
plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2]
@@ -1254,24 +1494,10 @@ class BuildDepthPlane(bpy.types.Operator):
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty(
default=5,
min=0,
max=10,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_displacement_via: bpy.props.EnumProperty(
items=(
("MODIFIER", "Modifier", "Use a Displace modifier"),
("MATERIAL", "Material", "Use material displacement"),
),
default="MODIFIER",
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
@@ -1295,14 +1521,15 @@ class BuildDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
try:
color_image = None
foreground_image = None
if self.snapform_background_path:
color_image = create_removed_background_image(
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
elif self.snapform_color_frames_dir:
color_image = None
if self.snapform_color_frames_dir:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
@@ -1311,12 +1538,11 @@ class BuildDepthPlane(bpy.types.Operator):
context,
source_object,
npz_path,
self.snapform_subdivision_level,
self.snapform_delete_source,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
depth_name=self.snapform_depth_data_name,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
@@ -1347,24 +1573,10 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
snapform_depth_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty(
default=5,
min=0,
max=10,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_displacement_via: bpy.props.EnumProperty(
items=(
("MODIFIER", "Modifier", "Use a Displace modifier"),
("MATERIAL", "Material", "Use material displacement"),
),
default="MODIFIER",
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
@@ -1379,14 +1591,15 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
)
return {"CANCELLED"}
try:
color_image = None
foreground_image = None
if self.snapform_background_path:
color_image = create_removed_background_image(
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
elif self.snapform_color_frames_dir:
color_image = None
if self.snapform_color_frames_dir:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
@@ -1411,11 +1624,10 @@ class BuildOriginalDepthPlane(bpy.types.Operator):
build_image_depth_plane(
context,
source_object,
self.snapform_subdivision_level,
self.snapform_delete_source,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
@@ -1527,27 +1739,11 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
),
default="ORIGINAL",
)
snapform_displacement_via: bpy.props.EnumProperty(
name="Displace via",
description="How the depth displacement is presented",
items=(
("MODIFIER", "Modifier", "Use a Displace modifier"),
("MATERIAL", "Material", "Use material displacement"),
),
default="MODIFIER",
)
snapform_unshaded: bpy.props.BoolProperty(
name="Unshaded",
description="Display the source color through an Emission shader",
default=False,
)
snapform_background_depth_distance: bpy.props.FloatProperty(
name="Background Depth Distance",
description="Set excluded background to this fixed depth distance; zero disables background exclusion",
default=0.0,
min=0.0,
soft_max=10.0,
)
snapform_model: bpy.props.EnumProperty(
name="Model",
description="Model used to estimate image depth",
@@ -1572,24 +1768,19 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
description="Delete the source Image Empty after conversion",
default=True,
)
snapform_subdivision_level: bpy.props.IntProperty(
name="Subdivision",
description="Power of two for the short-edge segment count; the long edge follows the image aspect ratio",
default=5,
min=0,
max=10,
)
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy, model_ready
from .models import MODELS
from .models import BEN2_MODEL, MODELS
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and any(model_ready(model_key) for model_key in MODELS)
and model_ready(BEN2_MODEL.key)
)
def invoke(self, context, _event):
@@ -1608,9 +1799,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return context.window_manager.invoke_props_dialog(self, width=340)
def draw(self, context):
from .models import BEN2_MODEL
from .runtime import model_ready
layout = self.layout
layout.prop(self, "snapform_depth_source")
if self.snapform_depth_source == "ESTIMATE":
@@ -1637,14 +1825,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
)
settings_header.label(text="Options")
if settings_body is not None:
settings_body.prop(self, "snapform_displacement_via")
settings_body.prop(self, "snapform_subdivision_level")
remove_background_row = settings_body.row()
remove_background_row.enabled = model_ready(BEN2_MODEL.key)
remove_background_row.prop(
self,
"snapform_background_depth_distance",
)
settings_body.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_delete_original")
@@ -1652,7 +1832,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
from .operators import (
IMAGE_DEPTH_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
)
source_object = context.object
@@ -1660,57 +1839,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
context.scene.snapform_settings.video_frames = (
default_extract_frames(source_object)
)
if (
self.snapform_depth_source == "ORIGINAL"
and self.snapform_background_depth_distance <= 0.0
):
if video_total_frames(source_object.data) > 1:
try:
source_path, delete_input = prepare_image_input(
source_object.data,
context.scene,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.run_job(
job_type=ORIGINAL_FRAMES_JOB,
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_subdivision_level=self.snapform_subdivision_level,
snapform_delete_source=self.snapform_delete_original,
snapform_delete_input=delete_input,
snapform_displacement_via=self.snapform_displacement_via,
snapform_unshaded=self.snapform_unshaded,
)
if "CANCELLED" not in result:
return {"FINISHED"}
if delete_input:
remove_temporary_input(source_path)
return {"CANCELLED"}
try:
build_image_depth_plane(
context,
source_object,
self.snapform_subdivision_level,
self.snapform_delete_original,
self.snapform_displacement_via,
self.snapform_unshaded,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
undo_result = bpy.ops.ed.undo_push(
message="Convert to Depth Plane",
)
if "FINISHED" not in undo_result:
self.report(
{"ERROR"},
"Unable to create the depth plane Undo step",
)
return {"CANCELLED"}
return {"FINISHED"}
if self.snapform_depth_source == "ORIGINAL":
try:
source_path, delete_input = prepare_image_input(
@@ -1724,14 +1852,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
job_type=ORIGINAL_BACKGROUND_JOB,
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_subdivision_level=self.snapform_subdivision_level,
snapform_delete_source=self.snapform_delete_original,
snapform_delete_input=delete_input,
snapform_displacement_via=self.snapform_displacement_via,
snapform_unshaded=self.snapform_unshaded,
snapform_background_depth_distance=(
self.snapform_background_depth_distance
),
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -1765,14 +1888,10 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
snapform_model_key=self.snapform_model,
snapform_use_ray_pose=self.snapform_use_ray_pose,
snapform_process_res=process_res,
snapform_subdivision_level=self.snapform_subdivision_level,
snapform_delete_source=self.snapform_delete_original,
snapform_delete_input=delete_input,
snapform_displacement_via=self.snapform_displacement_via,
snapform_unshaded=self.snapform_unshaded,
snapform_background_depth_distance=(
self.snapform_background_depth_distance
),
snapform_remove_background=True,
)
if "CANCELLED" in result:
if delete_input:
@@ -1797,7 +1916,7 @@ class SnapFormImageMenu(bpy.types.Menu):
)
layout.operator(
"snapform.convert_image_to_mesh",
text="Convert to Mesh",
text="Convert to Surface",
icon="MESH_GRID",
)
layout.separator()

18
src/snapform/image_mesh.py
View File

@@ -33,7 +33,7 @@ FACE_ATTRIBUTES = {
"o_edge_stretch": ("FLOAT", "face_edge_stretch"),
"o_background_face": ("BOOLEAN", "face_background"),
}
NODE_GROUP_NAME = "O Refine DA3 Mesh"
NODE_GROUP_NAME = "O Refine DA3 Surface"
NODE_ASSET_NAME = "O_SnapForm.blend"
@@ -60,7 +60,7 @@ def load_image_mesh_data(npz_path):
missing = [key for key in required if key not in archive.files]
if missing:
raise ValueError(
f"Mesh data is missing: {', '.join(missing)}"
f"Surface data is missing: {', '.join(missing)}"
)
vertices = np.asarray(archive["vertices"], dtype=np.float32)
faces = np.asarray(archive["faces"], dtype=np.int32)
@@ -81,19 +81,19 @@ def load_image_mesh_data(npz_path):
}.items()
}
except OSError as error:
raise ValueError(f"Unable to read mesh data: {npz_path}") from error
raise ValueError(f"Unable to read surface data: {npz_path}") from error
if vertices.ndim != 2 or vertices.shape[1] != 3:
raise ValueError(f"Unsupported vertices shape: {vertices.shape}")
if faces.ndim != 2 or faces.shape[1] != 3:
raise ValueError(f"Unsupported faces shape: {faces.shape}")
if uv.shape != (len(vertices), 2):
raise ValueError("Mesh UV and vertex counts do not match")
raise ValueError("Surface UV and vertex counts do not match")
if intrinsics.shape != (3, 3):
raise ValueError(f"Unsupported intrinsics shape: {intrinsics.shape}")
if len(image_size) != 2 or min(image_size) <= 0:
raise ValueError(f"Unsupported image size: {image_size}")
if not reference_depth > 0.0:
raise ValueError("Mesh reference depth must be positive")
raise ValueError("Surface reference depth must be positive")
for _name, (_data_type, key) in POINT_ATTRIBUTES.items():
if len(attributes[key]) != len(vertices):
raise ValueError(f"Point attribute has the wrong length: {key}")
@@ -394,7 +394,7 @@ def build_image_mesh(
class BuildImageMesh(bpy.types.Operator):
bl_idname = "snapform.build_image_mesh"
bl_label = "Build Image Mesh"
bl_label = "Build Image Surface"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
@@ -438,8 +438,8 @@ class BuildImageMesh(bpy.types.Operator):
class ConvertImageToMesh(bpy.types.Operator):
bl_idname = "snapform.convert_image_to_mesh"
bl_label = "Convert to Mesh"
bl_description = "Convert an Image Empty into a perspective textured mesh"
bl_label = "Convert to Surface"
bl_description = "Convert an Image Empty into a perspective textured surface"
bl_options = {"REGISTER"}
snapform_model: bpy.props.EnumProperty(
@@ -529,7 +529,7 @@ class ConvertImageToMesh(bpy.types.Operator):
source_object = context.object
if video_total_frames(source_object.data) != 1:
self.report({"ERROR"}, "Convert to Mesh currently supports still images")
self.report({"ERROR"}, "Convert to Surface currently supports still images")
return {"CANCELLED"}
try:
process_res = scaled_process_resolution(

78
src/snapform/operators.py
View File

@@ -73,12 +73,9 @@ class ActiveJob:
source_object_name: str = ""
depth_data_name: str = ""
process_res: int = IMAGE_DEPTH_DEFAULT_PROCESS_RES
subdivision_level: int = 5
delete_source: bool = True
displacement_via: str = "MODIFIER"
unshaded: bool = False
remove_background: bool = False
background_depth_distance: float = 0.0
mesh_create_camera: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
@@ -141,7 +138,7 @@ class RunJob(bpy.types.Operator):
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(IMAGE_MESH_JOB, "Convert Image to Mesh", ""),
(IMAGE_MESH_JOB, "Convert Image to Surface", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
@@ -183,29 +180,10 @@ class RunJob(bpy.types.Operator):
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty(
default=5,
min=0,
max=10,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_displacement_via: bpy.props.EnumProperty(
items=(
("MODIFIER", "Modifier", ""),
("MATERIAL", "Material", ""),
),
default="MODIFIER",
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_depth_distance: bpy.props.FloatProperty(
default=0.0,
min=0.0,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_mesh_create_camera: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
@@ -214,6 +192,7 @@ class RunJob(bpy.types.Operator):
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
global _active_job
@@ -251,15 +230,11 @@ class RunJob(bpy.types.Operator):
source_object_name=self.snapform_source_object_name,
depth_data_name=self.snapform_depth_data_name,
process_res=self.snapform_process_res,
subdivision_level=self.snapform_subdivision_level,
delete_source=self.snapform_delete_source,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
remove_background=(
self.snapform_background_depth_distance > 0.0
or bool(getattr(self, "snapform_remove_background", False))
remove_background=bool(
getattr(self, "snapform_remove_background", False)
),
background_depth_distance=self.snapform_background_depth_distance,
mesh_create_camera=getattr(
self,
"snapform_mesh_create_camera",
@@ -350,10 +325,9 @@ class RunJob(bpy.types.Operator):
)
self.snapform_input_path = str(archived_input)
self.snapform_delete_input = False
background_depth_distance = float(
getattr(self, "snapform_background_depth_distance", 0.0)
remove_background = bool(
getattr(self, "snapform_remove_background", False)
)
remove_background = background_depth_distance > 0.0
if self.job_type == IMAGE_MESH_JOB:
remove_background = True
if self.job_type == "INSTALL_RUNTIME":
@@ -413,11 +387,6 @@ class RunJob(bpy.types.Operator):
BEN2_MODEL.key,
start_frame=start_frame,
device=self._task_device(settings),
background_depth_distance=(
background_depth_distance
if self.job_type == ORIGINAL_BACKGROUND_JOB
else 0.0
),
),
paths,
)
@@ -483,11 +452,6 @@ class RunJob(bpy.types.Operator):
self.job_type in {IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}
and remove_background
),
background_depth_distance=(
background_depth_distance
if self.job_type == IMAGE_DEPTH_JOB
else 0.0
),
device=self._task_device(settings),
),
paths,
@@ -520,7 +484,7 @@ class RunJob(bpy.types.Operator):
if self.job_type == IMAGE_DEPTH_JOB:
return "Starting depth estimation"
if self.job_type == IMAGE_MESH_JOB:
return "Starting perspective mesh generation"
return "Starting perspective surface generation"
if self.job_type == BACKGROUND_REMOVAL_JOB:
return "Starting BEN2 background removal"
if self.job_type == IMAGE_BACKGROUND_JOB:
@@ -696,19 +660,19 @@ class RunJob(bpy.types.Operator):
status["message"] = "Depth plane is ready"
if job.kind == IMAGE_MESH_JOB:
try:
settings.status = "Building perspective mesh"
settings.status = "Building perspective surface"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_image_mesh(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "Mesh import failed"
settings.status = "Surface import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Perspective mesh is ready"
status["message"] = "Perspective surface is ready"
if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}:
try:
settings.status = "Building depth plane"
@@ -761,10 +725,10 @@ class RunJob(bpy.types.Operator):
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:
raise RuntimeError("Unable to build the perspective mesh")
undo_result = bpy.ops.ed.undo_push(message="Convert to Mesh")
raise RuntimeError("Unable to build the perspective surface")
undo_result = bpy.ops.ed.undo_push(message="Convert to Surface")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the mesh Undo step")
raise RuntimeError("Unable to create the surface Undo step")
def _build_depth_plane(self, context, job):
result_directory = job.result_directory
@@ -784,18 +748,14 @@ class RunJob(bpy.types.Operator):
snapform_depth_data_name=job.depth_data_name,
snapform_background_path=background_path,
snapform_color_frames_dir=(
str(result_directory)
if is_sequence and not foreground_paths
else ""
str(result_directory) if is_sequence else ""
),
snapform_depth_frames_dir=(
str(result_directory)
if is_sequence
else ""
),
snapform_subdivision_level=job.subdivision_level,
snapform_delete_source=job.delete_source,
snapform_displacement_via=job.displacement_via,
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:
@@ -815,7 +775,7 @@ class RunJob(bpy.types.Operator):
if is_sequence
else result_directory / "foreground.png"
)
elif is_sequence:
if is_sequence:
color_frames_dir = str(result_directory)
result = bpy.ops.snapform.build_original_depth_plane(
snapform_source_object_name=job.source_object_name,
@@ -824,14 +784,8 @@ class RunJob(bpy.types.Operator):
snapform_depth_frames_dir=(
str(result_directory) if is_sequence else ""
),
snapform_depth_npz_path=(
str(result_directory / "original-depth.npz")
if job.kind == ORIGINAL_BACKGROUND_JOB
else ""
),
snapform_subdivision_level=job.subdivision_level,
snapform_depth_npz_path="",
snapform_delete_source=job.delete_source,
snapform_displacement_via=job.displacement_via,
snapform_unshaded=job.unshaded,
)
if "FINISHED" not in result:

2
src/snapform/panel.py
View File

@@ -216,7 +216,7 @@ class Da3OutputDebugPanel(bpy.types.Panel):
ply_box.prop(settings, "gs_views_interval")
elif settings.export_format == "mesh":
mesh_box = layout.box()
mesh_box.label(text="Mesh Settings")
mesh_box.label(text="Surface Settings")
mesh_box.prop(settings, "conf_thresh_percentile")
mesh_box.prop(settings, "mesh_filter_edges")
mesh_box.prop(settings, "mesh_max_edge_length")

6
src/snapform/properties.py
View File

@@ -67,7 +67,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
)
conf_thresh_percentile: bpy.props.FloatProperty(
name="Confidence Threshold",
description="Remove the lowest-confidence percentile of points or mesh faces",
description="Remove the lowest-confidence percentile of points or surface faces",
default=40.0,
min=0.0,
max=100.0,
@@ -135,8 +135,8 @@ class SnapFormSettings(bpy.types.PropertyGroup):
),
(
"mesh",
"Textured Mesh (GLB)",
"Generate one filtered textured depth mesh per input image",
"Textured Surface (GLB)",
"Generate one filtered textured depth surface per input image",
),
),
default="glb",

8
src/snapform/runtime.py
View File

@@ -748,7 +748,6 @@ def remove_background_request(
video_frames=None,
start_frame=1,
device=None,
background_depth_distance=0.0,
):
return HostJobRequest(
{
@@ -766,9 +765,6 @@ def remove_background_request(
else video_frames
),
"start_frame": start_frame,
"background_depth_distance": float(
background_depth_distance
),
},
}
)
@@ -812,7 +808,6 @@ def reconstruct_request(
process_res_method=None,
use_ray_pose=None,
remove_background=False,
background_depth_distance=0.0,
mesh_filter_edges=None,
start_frame=1,
device=None,
@@ -851,7 +846,6 @@ def reconstruct_request(
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
"ref_view_strategy": settings.ref_view_strategy,
"remove_background_model_dir": "",
"background_depth_distance": float(background_depth_distance),
"mesh_filter_edges": bool(
getattr(settings, "mesh_filter_edges", True)
if mesh_filter_edges is None
@@ -861,7 +855,7 @@ def reconstruct_request(
getattr(settings, "mesh_max_edge_length", 0.1)
),
}
if remove_background or background_depth_distance > 0.0:
if remove_background:
parameters["remove_background_model_dir"] = str(
model_directory(BEN2_MODEL.key)
)

152
src/snapform/worker.py
View File

@@ -48,64 +48,6 @@ def check_cancelled(cancel_check):
raise TaskCancelled("Task cancelled")
def apply_foreground_depth_distance(
prediction,
foreground_paths,
background_depth_distance,
):
"""Set fully transparent pixels to a fixed negative displacement."""
import numpy as np
from PIL import Image
if len(foreground_paths) != len(prediction.depth):
raise RuntimeError("Foreground mask count does not match depth frames")
masks = []
height, width = prediction.depth.shape[-2:]
for foreground_path in foreground_paths:
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize((width, height), resample=Image.Resampling.NEAREST)
masks.append(np.asarray(alpha) > 0)
apply_depth_background_array(
prediction,
np.stack(masks),
background_depth_distance,
)
def apply_depth_background_array(
prediction,
foreground_mask,
background_depth_distance,
):
import numpy as np
depth = np.asarray(prediction.depth)
masked_depth = depth.copy()
background_depth = -background_depth_distance
for frame_index in range(len(masked_depth)):
np.maximum(
masked_depth[frame_index],
background_depth,
out=masked_depth[frame_index],
)
masked_depth[frame_index][~foreground_mask[frame_index]] = (
background_depth
)
set_depth_border(masked_depth[frame_index], background_depth)
prediction.depth = masked_depth
def set_depth_border(depth, background_depth):
"""Set the outermost pixel ring to the fixed background depth."""
if depth.ndim != 2 or not depth.size:
raise ValueError("Depth frame must be a non-empty 2D array")
depth[0, :] = background_depth
depth[-1, :] = background_depth
depth[:, 0] = background_depth
depth[:, -1] = background_depth
def convert_ai_depth_for_displacement(prediction):
"""Convert DA3 distance into Blender's positive displacement direction."""
import numpy as np
@@ -137,55 +79,6 @@ def write_npz_reference_depth(npz_path, reference_depth):
temporary_path.replace(npz_path)
def write_original_background_depth(
image_paths,
foreground_paths,
background_depth_distance,
output_directory,
):
"""Write original-image luminance with a fixed negative background depth."""
import numpy as np
from PIL import Image
if len(image_paths) != len(foreground_paths):
raise RuntimeError("Foreground mask count does not match source frames")
depth_frames = []
for image_path, foreground_path in zip(image_paths, foreground_paths):
with Image.open(image_path) as source:
luminance = np.asarray(source.convert("L"), dtype=np.float32) / 255.0
with Image.open(foreground_path) as foreground:
alpha = np.asarray(
foreground.convert("RGBA").getchannel("A"),
dtype=np.uint8,
)
if luminance.shape != alpha.shape:
raise RuntimeError("Foreground mask size does not match source frame")
depth_frames.append(
compose_original_background_depth(
luminance,
alpha,
background_depth_distance,
)
)
np.savez_compressed(
Path(output_directory) / "original-depth.npz",
depth=np.stack(depth_frames),
)
def compose_original_background_depth(
luminance,
alpha,
background_depth_distance,
):
import numpy as np
background_depth = -background_depth_distance
depth = np.where(alpha > 0, luminance, background_depth)
set_depth_border(depth, background_depth)
return depth
def write_status(
path,
progress,
@@ -608,7 +501,7 @@ def export_image_mesh_data(
depth_frames = np.asarray(prediction.depth, dtype=np.float32)
confidence_frames = np.asarray(prediction.conf, dtype=np.float32)
if depth_frames.shape[0] != 1:
raise ValueError("Convert to Mesh currently supports one image")
raise ValueError("Convert to Surface currently supports one image")
raw_depth = depth_frames[0]
confidence = confidence_frames[0]
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
@@ -636,7 +529,7 @@ def export_image_mesh_data(
)
if not len(faces):
raise ValueError(
"DA3 output does not contain any valid mesh faces"
"DA3 output does not contain any valid surface faces"
)
progress(5, "Projecting UV and structured attributes")
vertices, faces = compact_mesh(
@@ -658,7 +551,7 @@ def export_image_mesh_data(
reference_depth = float(np.median(valid_depth))
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
progress(6, "Writing mesh data")
progress(6, "Writing surface data")
result_path = output_directory / "mesh.npz"
np.savez_compressed(
result_path,
@@ -876,7 +769,7 @@ def export_textured_mesh_glb(
referenced_points.append(mesh.vertices)
if not meshes:
raise ValueError(
"Mesh filters removed every face; lower Confidence Threshold or disable filters"
"Surface filters removed every face; lower Confidence Threshold or disable filters"
)
world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64)
@@ -1343,6 +1236,7 @@ def infer_foreground(
image,
refine_foreground=False,
)
foreground = preserve_source_rgb(image, foreground)
if device.type == "cuda":
torch.cuda.synchronize()
inference_ms = (time.perf_counter() - inference_started) * 1000.0
@@ -1354,6 +1248,18 @@ def infer_foreground(
return foreground, load_ms, inference_ms
def preserve_source_rgb(source, foreground):
"""Combine the untouched source RGB with BEN2's foreground alpha."""
source_rgba = source.convert("RGBA")
alpha = foreground.convert("RGBA").getchannel("A")
if alpha.size != source_rgba.size:
from PIL import Image
alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS)
source_rgba.putalpha(alpha)
return source_rgba
def infer_foreground_frames(
images,
model_dir,
@@ -1456,21 +1362,6 @@ def remove_background(args, model_cache=None, cancel_check=None):
cancel_check=cancel_check,
)
copy_frames(images, output_directory)
background_depth_distance = float(
getattr(args, "background_depth_distance", 0.0)
)
if background_depth_distance > 0.0:
foreground_paths = (
[output_directory / "foreground.png"]
if len(images) == 1
else sorted(output_directory.glob("foreground.*.png"))
)
write_original_background_depth(
images,
foreground_paths,
background_depth_distance,
output_directory,
)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -1662,9 +1553,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
"remove_background_model_dir",
"",
)
background_depth_distance = float(
getattr(args, "background_depth_distance", 0.0)
)
foreground_images = []
if remove_background_model_dir:
write_status(args.status, 0.12, "Removing image background")
@@ -1747,12 +1635,6 @@ def reconstruct(args, model_cache=None, cancel_check=None):
if args.export_format == "mini_npz":
reference_depth = depth_reference_value(prediction.depth[0])
convert_ai_depth_for_displacement(prediction)
if background_depth_distance > 0.0:
apply_foreground_depth_distance(
prediction,
foreground_images,
background_depth_distance,
)
write_status(args.status, 0.88, f"Exporting {args.export_format}")
if args.export_format == "mesh":
export_textured_mesh_glb(