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

Refactor image mesh processing and background removal functionality

Browse Source 
- Updated NODE_GROUP_NAME from "O Refine DA3 Surface" to "O DA3 Surface".
- Introduced `remove_surface_background` function to filter out background faces and compact geometry.
- Enhanced `build_image_mesh` to conditionally remove background faces based on a new property.
- Added `snapform_remove_background` and `snapform_keep_original` properties to operators for better control over background handling.
- Updated UI to reflect new properties and improved resolution scaling logic.
- Refactored tests to cover new background removal functionality and ensure proper geometry handling.
- Adjusted geometry nodes and validation scripts to remove references to background handling.
This commit is contained in:
马南
2026-08-12 16:15:14 +08:00
parent 7cd5681a8b
commit c3a151ad2c
13 changed files with 477 additions and 271 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
src/snapform/image_depth.py
View File

@@ -15,12 +15,14 @@ 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"
SURFACE_MATERIAL_NODE_GROUP_NAME = "O DA3 Surface Layer"
DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME = "O DA3 Depth Plane Layer"
DA3_PATCH_SIZE = 14
DA3_MIN_PROCESS_RES = 252
DA3_MAX_PROCESS_RES = 1008
IMAGE_NAME_SUFFIXES = (
".jpeg",
".jpg",
@@ -180,13 +182,45 @@ def removed_background_data_name(source_object):
return f"{image_base_name(source_object.data)}.foreground"
def scaled_process_resolution(image, scale):
def scaled_process_resolution(image, scale_percent):
"""Return the scaled long edge clamped and aligned for DA3 inference."""
width, height = image.size
if width <= 0 or height <= 0:
raise ValueError("The source image has no readable dimensions")
if scale < 0.1 or scale > 1.0:
raise ValueError("Resolution Scale must be between 0.1 and 1.0")
return int(max(width, height) * scale + 0.5)
if scale_percent < 1.0 or scale_percent > 100.0:
raise ValueError("Resolution Scale must be between 1% and 100%")
process_res = max(width, height) * scale_percent / 100.0
process_res = max(DA3_MIN_PROCESS_RES, process_res)
process_res = min(DA3_MAX_PROCESS_RES, process_res)
return int(process_res / DA3_PATCH_SIZE + 0.5) * DA3_PATCH_SIZE
def process_resolution_scale_property():
return bpy.props.FloatProperty(
name="Resolution Scale",
description=(
"Scale the source image long edge for DA3 processing; the result "
"is limited to 252-1008 px and aligned to a multiple of 14"
),
default=100.0,
min=1.0,
max=100.0,
step=100,
precision=0,
subtype="PERCENTAGE",
)
def draw_process_resolution(layout, owner, image):
process_res = scaled_process_resolution(
image,
owner.snapform_process_res_factor,
)
layout.prop(
owner,
"snapform_process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
def save_packed_image(image, scene):
@@ -880,13 +914,15 @@ def material_node_group(depth_plane=False):
socket.default_value = 1.0
socket.min_value = -10000.0
socket.max_value = 10000.0
for name, socket_type in (
output_sockets = [
("Color", "NodeSocketColor"),
("Roughness", "NodeSocketFloat"),
("Alpha", "NodeSocketFloat"),
("Normal", "NodeSocketVector"),
("Displacement", "NodeSocketVector"),
):
]
if depth_plane:
output_sockets.append(("Displacement", "NodeSocketVector"))
for name, socket_type in output_sockets:
interface.new_socket(
name=name,
in_out="OUTPUT",
@@ -908,23 +944,20 @@ def material_node_group(depth_plane=False):
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)
displacement = None
processed_alpha = group_input.outputs["Alpha"]
processed_depth = group_input.outputs["Depth"]
bump_depth = group_input.outputs["Color"]
if depth_plane:
displacement = nodes.new("ShaderNodeDisplacement")
displacement.inputs["Midlevel"].default_value = 0.0
displacement.inputs["Scale"].default_value = 1.0
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
@@ -937,10 +970,6 @@ def material_node_group(depth_plane=False):
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
@@ -951,31 +980,23 @@ def material_node_group(depth_plane=False):
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(masked_depth.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_alpha = foreground_mask.outputs[0]
processed_depth = scaled_depth.outputs[0]
links.new(group_input.outputs["Color"], group_output.inputs["Color"])
@@ -985,7 +1006,7 @@ def material_node_group(depth_plane=False):
group_output.inputs["Roughness"],
)
links.new(processed_alpha, group_output.inputs["Alpha"])
links.new(processed_depth, bump.inputs["Height"])
links.new(bump_depth, bump.inputs["Height"])
bump_distance = nodes.new("ShaderNodeMath")
bump_distance.operation = "MULTIPLY"
bump_distance.inputs[1].default_value = BUMP_BASE_DISTANCE
@@ -993,11 +1014,12 @@ def material_node_group(depth_plane=False):
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"],
)
if displacement is not None:
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."
@@ -1081,9 +1103,7 @@ def create_material(
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.name = material_inputs.node_tree.name
material_inputs.label = material_inputs.name
material_inputs.location = (-300.0, 350.0)
material_inputs.inputs["Bump Scale"].default_value = depth_scale
@@ -1243,21 +1263,32 @@ def add_depth_plane_processing(
frame_start=1,
frame_duration=1,
use_cyclic=False,
animate_frame=False,
):
node_group = load_depth_plane_node_group()
modifier = plane_object.modifiers.new(name=node_group.name, type="NODES")
modifier = plane_object.modifiers.new(
name=DEPTH_PLANE_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", 1),
("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
if animate_frame:
frame_identifier = modifier_input_identifier(node_group, "Frame")
driver = modifier.driver_add(f'["{frame_identifier}"]')
driver.driver.type = "SCRIPTED"
driver.driver.expression = "frame"
plane_object.update_tag(refresh={"DATA"})
return modifier
@@ -1353,6 +1384,7 @@ def build_image_depth_plane(
frame_start=source_frame_start,
frame_duration=frame_count,
use_cyclic=source_use_cyclic,
animate_frame=is_animated_image(source_object.data),
)
plane_object["snapform_frame_count"] = frame_count
return finish_depth_plane(
@@ -1456,6 +1488,7 @@ def build_depth_plane(
frame_start=source_frame_start,
frame_duration=frame_count,
use_cyclic=source_use_cyclic,
animate_frame=is_animated_image(source_object.data),
)
plane_object["snapform_reference_depth"] = reference_depth
plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0]
@@ -1744,6 +1777,11 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
description="Display the source color through an Emission shader",
default=False,
)
snapform_remove_background: bpy.props.BoolProperty(
name="Remove Background",
description="Remove the background during conversion with BEN2",
default=False,
)
snapform_model: bpy.props.EnumProperty(
name="Model",
description="Model used to estimate image depth",
@@ -1754,33 +1792,24 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
description="Use ray-based camera pose estimation",
default=True,
)
snapform_process_res_factor: bpy.props.FloatProperty(
name="Resolution Scale",
description="Scale applied to the source image long-edge resolution",
default=1.0,
min=0.1,
max=1.0,
step=10,
precision=2,
)
snapform_delete_original: bpy.props.BoolProperty(
name="Delete Original",
description="Delete the source Image Empty after conversion",
default=True,
snapform_process_res_factor: process_resolution_scale_property()
snapform_keep_original: bpy.props.BoolProperty(
name="Keep Original",
description="Keep the source Image Empty after conversion",
default=False,
)
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy, model_ready
from .models import BEN2_MODEL, MODELS
from .models import 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):
@@ -1810,28 +1839,17 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
else:
estimate_box.prop(self, "snapform_model")
estimate_box.prop(self, "snapform_use_ray_pose")
process_res = scaled_process_resolution(
context.object.data,
self.snapform_process_res_factor,
)
estimate_box.prop(
self,
"snapform_process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
settings_header, settings_body = layout.panel(
"snapform_depth_plane_options",
default_closed=True,
)
settings_header.label(text="Options")
if settings_body is not None:
settings_body.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_delete_original")
draw_process_resolution(estimate_box, self, context.object.data)
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_unshaded")
layout.separator(type="LINE")
layout.prop(self, "snapform_keep_original")
def execute(self, context):
from .operators import (
IMAGE_DEPTH_JOB,
ORIGINAL_BACKGROUND_JOB,
ORIGINAL_FRAMES_JOB,
)
source_object = context.object
@@ -1849,10 +1867,14 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.run_job(
job_type=ORIGINAL_BACKGROUND_JOB,
job_type=(
ORIGINAL_BACKGROUND_JOB
if self.snapform_remove_background
else ORIGINAL_FRAMES_JOB
),
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_delete_source=self.snapform_delete_original,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
)
@@ -1888,10 +1910,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_delete_source=self.snapform_delete_original,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
snapform_remove_background=True,
snapform_remove_background=self.snapform_remove_background,
)
if "CANCELLED" in result:
if delete_input: