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

feat: Add BEN2 background removal functionality and integrate into SnapForm

Browse Source 
- Introduced BEN2 model for background removal in the panel and properties.
- Updated runtime to handle background removal commands and model downloads.
- Implemented background removal logic in worker, including image processing and saving results.
- Enhanced tests to cover new background removal features and ensure proper integration.
- Updated version to 0.9.28 to reflect new features and improvements.
This commit is contained in:
马南
2026-08-03 20:18:00 +08:00
parent 69eaefa016
commit d694980da7
19 changed files with 1523 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

375
src/snapform/image_depth.py
View File

@@ -80,6 +80,10 @@ def depth_data_name(source_object):
return f"{image_base_name(source_object.data)}.depth"
def removed_background_data_name(source_object):
return f"{image_base_name(source_object.data)}.rb"
def scaled_process_resolution(image, scale):
width, height = image.size
if width <= 0 or height <= 0:
@@ -108,6 +112,49 @@ def save_packed_image(image, scene):
return output_path
def prepare_image_input(image, scene):
source_path = image_path(image)
if image_is_packed(image):
return save_packed_image(image, scene), True
if not source_path.is_file():
raise RuntimeError(
"The image empty must reference an existing image file"
)
return source_path, False
def create_removed_background_image(
source_object,
output_path,
assign_to_source=True,
):
output_path = Path(output_path)
if not output_path.is_file():
raise RuntimeError(
f"Background removal result does not exist: {output_path}"
)
try:
result_image = bpy.data.images.load(
str(output_path),
check_existing=False,
)
except RuntimeError as error:
raise RuntimeError("Unable to load the background result") from error
source_image = source_object.data
original_fake_user = source_image.use_fake_user
try:
result_image.name = removed_background_data_name(source_object)
result_image.pack()
if assign_to_source:
source_image.use_fake_user = True
source_object.data = result_image
except (OSError, RuntimeError) as error:
source_image.use_fake_user = original_fake_user
bpy.data.images.remove(result_image)
raise RuntimeError("Unable to pack the background result") from error
return result_image
def plane_bounds(source_object):
image_width, image_height = source_object.data.size
if image_width <= 0 or image_height <= 0:
@@ -422,6 +469,7 @@ def create_material(
displacement_image=None,
displacement_midlevel=0.0,
displacement_scale=1.0,
unshaded=False,
):
material = bpy.data.materials.new(name)
material.use_nodes = True
@@ -430,42 +478,74 @@ def create_material(
nodes = material.node_tree.nodes
nodes.clear()
output = nodes.new("ShaderNodeOutputMaterial")
shader = nodes.new("ShaderNodeBsdfPrincipled")
color_texture = nodes.new("ShaderNodeTexImage")
bump = nodes.new("ShaderNodeBump")
roughness_curve = nodes.new("ShaderNodeFloatCurve")
output = nodes.new("ShaderNodeOutputMaterial")
color_texture.image = color_image
color_texture.interpolation = "Linear"
color_texture.extension = "EXTEND"
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()
color_texture.location = (-640.0, 180.0)
roughness_curve.location = (-300.0, 320.0)
bump.location = (-300.0, -80.0)
shader.location = (80.0, 180.0)
output.location = (460.0, 180.0)
links = material.node_tree.links
links.new(color_texture.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(
color_texture.outputs["Color"],
roughness_curve.inputs["Value"],
)
links.new(
roughness_curve.outputs["Value"],
shader.inputs["Roughness"],
)
links.new(shader.outputs["BSDF"], output.inputs["Surface"])
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)
links.new(
color_texture.outputs["Color"],
emission.inputs["Color"],
)
links.new(
color_texture.outputs["Alpha"],
mix_shader.inputs["Fac"],
)
links.new(
transparent.outputs["BSDF"],
mix_shader.inputs[1],
)
links.new(
emission.outputs["Emission"],
mix_shader.inputs[2],
)
links.new(
mix_shader.outputs["Shader"],
output.inputs["Surface"],
)
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)
links.new(
color_texture.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(
color_texture.outputs["Color"],
roughness_curve.inputs["Value"],
)
links.new(
roughness_curve.outputs["Value"],
shader.inputs["Roughness"],
)
links.new(shader.outputs["BSDF"], output.inputs["Surface"])
if displacement_image is not None:
depth_texture = nodes.new("ShaderNodeTexImage")
@@ -525,6 +605,8 @@ def build_image_depth_plane(
subdivision_level,
delete_source,
displacement_via="MODIFIER",
unshaded=False,
color_image=None,
):
source_name = source_object.name
mesh = create_plane_mesh(
@@ -554,13 +636,14 @@ def build_image_depth_plane(
depth_scale = abs(float(plane_object.scale.z))
material = create_material(
source_name,
source_object.data,
color_image if color_image is not None else source_object.data,
depth_scale,
displacement_image=(
source_object.data if displacement_via == "MATERIAL" else None
),
displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL,
displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH,
unshaded=unshaded,
)
mesh.materials.append(material)
return finish_depth_plane(
@@ -578,6 +661,9 @@ def build_depth_plane(
subdivision_level,
delete_source,
displacement_via="MODIFIER",
unshaded=False,
depth_name="",
color_image=None,
):
source_name = source_object.name
depth = extract_depth_array(npz_path)
@@ -589,7 +675,7 @@ def build_depth_plane(
intrinsics,
DISPLACEMENT_STRENGTH,
)
depth_name = depth_data_name(source_object)
depth_name = depth_name or depth_data_name(source_object)
depth_image = create_depth_image(
depth_name,
depth_values,
@@ -627,13 +713,14 @@ def build_depth_plane(
material = create_material(
source_name,
source_object.data,
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_midlevel=DISPLACEMENT_MIDLEVEL,
displacement_scale=DISPLACEMENT_STRENGTH,
unshaded=unshaded,
)
mesh.materials.append(material)
plane_object["snapform_reference_depth"] = reference_depth
@@ -660,6 +747,12 @@ class BuildDepthPlane(bpy.types.Operator):
snapform_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_data_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_subdivision_level: bpy.props.IntProperty(
default=5,
min=0,
@@ -678,6 +771,10 @@ class BuildDepthPlane(bpy.types.Operator):
default="MODIFIER",
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
@@ -697,13 +794,23 @@ class BuildDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
try:
color_image = None
if self.snapform_background_path:
color_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
build_depth_plane(
context,
source_object,
npz_path,
self.snapform_subdivision_level,
self.snapform_delete_source,
self.snapform_displacement_via,
displacement_via=self.snapform_displacement_via,
unshaded=self.snapform_unshaded,
depth_name=self.snapform_depth_data_name,
color_image=color_image,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
@@ -711,6 +818,138 @@ class BuildDepthPlane(bpy.types.Operator):
return {"FINISHED"}
class BuildOriginalDepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_original_depth_plane"
bl_label = "Build Original Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_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"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
try:
color_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
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,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class ApplyRemovedBackground(bpy.types.Operator):
bl_idname = "snapform.apply_removed_background"
bl_label = "Apply Removed Background"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_output_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, _context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
try:
create_removed_background_image(
source_object,
self.snapform_output_path,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class RemoveImageBackground(bpy.types.Operator):
bl_idname = "snapform.remove_image_background"
bl_label = "Remove Background"
bl_description = "Remove the background from this Image Empty with BEN2"
bl_options = {"REGISTER"}
@classmethod
def poll(cls, context):
from .operators import active_job
return is_image_empty(context.object) and active_job() is None
def execute(self, context):
from .operators import IMAGE_BACKGROUND_JOB
source_object = context.object
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=IMAGE_BACKGROUND_JOB,
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_delete_input=delete_input,
)
if "CANCELLED" not in result:
return {"FINISHED"}
if delete_input:
source_path.unlink(missing_ok=True)
return {"CANCELLED"}
class ConvertImageToDepthPlane(bpy.types.Operator):
bl_idname = "snapform.convert_image_to_depth_plane"
bl_label = "Convert to Depth Plane"
@@ -739,6 +978,16 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
),
default="MODIFIER",
)
snapform_unshaded: bpy.props.BoolProperty(
name="Unshaded",
description="Display the source color through an Emission shader",
default=False,
)
snapform_remove_background: bpy.props.BoolProperty(
name="Remove Background",
description="Remove the color image background with BEN2",
default=False,
)
snapform_model: bpy.props.EnumProperty(
name="Model",
description="Model used to estimate image depth",
@@ -800,15 +1049,20 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
"snapform_process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_displacement_via")
layout.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_subdivision_level")
layout.prop(self, "snapform_delete_original")
def execute(self, context):
from .operators import IMAGE_DEPTH_JOB
from .operators import IMAGE_DEPTH_JOB, ORIGINAL_BACKGROUND_JOB
source_object = context.object
if self.snapform_depth_source == "ORIGINAL":
if (
self.snapform_depth_source == "ORIGINAL"
and not self.snapform_remove_background
):
try:
build_image_depth_plane(
context,
@@ -816,6 +1070,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
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))
@@ -831,6 +1086,31 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
return {"FINISHED"}
if self.snapform_depth_source == "ORIGINAL":
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_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,
)
if "CANCELLED" not in result:
return {"FINISHED"}
if delete_input:
source_path.unlink(missing_ok=True)
return {"CANCELLED"}
image = source_object.data
try:
process_res = scaled_process_resolution(
@@ -840,26 +1120,20 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
except ValueError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
source_path = image_path(image)
delete_input = False
if image_is_packed(image):
try:
source_path = save_packed_image(image, context.scene)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
delete_input = True
if not source_path.is_file():
self.report(
{"ERROR"},
"The image empty must reference an existing image file",
try:
source_path, delete_input = prepare_image_input(
image,
context.scene,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.run_job(
job_type=IMAGE_DEPTH_JOB,
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_depth_data_name=depth_data_name(source_object),
snapform_model_key=self.snapform_model,
snapform_use_ray_pose=self.snapform_use_ray_pose,
snapform_process_res=process_res,
@@ -867,6 +1141,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
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_remove_background=self.snapform_remove_background,
)
if "CANCELLED" in result:
if delete_input:
@@ -889,3 +1165,10 @@ def draw_image_context_menu(self, context):
text="Convert to Depth Plane",
icon="MOD_DISPLACE",
)
row = self.layout.row()
row.operator_context = "INVOKE_DEFAULT"
row.operator(
"snapform.remove_image_background",
text="Remove Background",
icon="IMAGE_ALPHA",
)