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

6
src/snapform/__init__.py
View File

@@ -1,8 +1,11 @@
import bpy
from .image_depth import (
ApplyRemovedBackground,
BuildDepthPlane,
BuildOriginalDepthPlane,
ConvertImageToDepthPlane,
RemoveImageBackground,
draw_image_context_menu,
)
from .operators import (
@@ -22,7 +25,10 @@ CLASSES = (
SnapFormSettings,
RunJob,
CancelJob,
ApplyRemovedBackground,
RemoveImageBackground,
BuildDepthPlane,
BuildOriginalDepthPlane,
ConvertImageToDepthPlane,
RuntimeModelsPanel,
DebugPanel,

2
src/snapform/blender_manifest.toml
View File

@@ -1,7 +1,7 @@
schema_version = "1.0.0"
id = "snapform"
version = "0.9.27"
version = "0.9.28"
name = "SnapForm"
tagline = "Build spatial structures from images, fast"
maintainer = "OmooLab"

21
src/snapform/bootstrap_runtime.py
View File

@@ -11,7 +11,7 @@ TORCHVISION_VERSION = "0.24.1"
XFORMERS_VERSION = "0.0.33.post2"
ADDICT_VERSION = "2.4.0"
GSPLAT_COMMIT = "0b4dddf04cb687367602c01196913cde6a743d70"
RUNTIME_SCHEMA = 2
RUNTIME_SCHEMA = 3
TORCH_INDEXES = {
"cu128": "https://download.pytorch.org/whl/cu128",
"cpu": "https://download.pytorch.org/whl/cpu",
@@ -113,6 +113,22 @@ def install(args):
]
)
write_status(args.status, 0.68, "Installing BEN2")
ben2_requirement = (
"ben2 @ git+https://github.com/PramaLLC/BEN2.git"
f"@{args.ben2_commit}"
)
run(
[
args.uv,
"pip",
"install",
"--python",
python,
ben2_requirement,
]
)
write_status(args.status, 0.78, "Installing gsplat build tools")
run(
[
@@ -154,6 +170,7 @@ def install(args):
(
"import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"import gsplat; "
"print(torch.__version__, torch.cuda.is_available())"
),
@@ -170,6 +187,7 @@ def install(args):
"xformers": XFORMERS_VERSION,
"backend": args.backend,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
"gsplat_commit": GSPLAT_COMMIT,
"gsplat_cuda": False,
},
@@ -186,6 +204,7 @@ def parse_args():
parser.add_argument("--runtime-dir", required=True)
parser.add_argument("--backend", choices=tuple(TORCH_INDEXES), required=True)
parser.add_argument("--da3-commit", required=True)
parser.add_argument("--ben2-commit", required=True)
parser.add_argument("--status", required=True)
return parser.parse_args()

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",
)

19
src/snapform/models.py
View File

@@ -9,6 +9,7 @@ class ModelSpec:
license_name: str
description: str
supports_gs: bool
download_patterns: tuple[str, ...] = ()
MODELS = {
@@ -53,6 +54,17 @@ MODELS = {
supports_gs=True,
),
}
BEN2_MODEL = ModelSpec(
key="BEN2_BASE",
label="BEN2 Base",
repository="PramaLLC/BEN2",
license_name="MIT",
description="Fast high-quality background removal model",
supports_gs=False,
download_patterns=("config.json", "model.safetensors"),
)
DOWNLOADABLE_MODELS = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
MODEL_NUMBERS = {
model_key: number for number, model_key in enumerate(MODELS)
}
@@ -123,3 +135,10 @@ def get_model(model_key):
return MODELS[model_key]
except KeyError as error:
raise ValueError(f"Unknown DA3 model: {model_key}") from error
def get_downloadable_model(model_key):
try:
return DOWNLOADABLE_MODELS[model_key]
except KeyError as error:
raise ValueError(f"Unknown model: {model_key}") from error

178
src/snapform/operators.py
View File

@@ -5,11 +5,19 @@ from pathlib import Path
import bpy
from . import runtime
from .models import get_model, selected_model_key
from .models import (
BEN2_MODEL,
get_downloadable_model,
get_model,
selected_model_key,
)
RECONSTRUCTION_JOB = "RECONSTRUCT"
IMAGE_DEPTH_JOB = "IMAGE_DEPTH"
BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND"
IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND"
ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND"
IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024
IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize"
@@ -33,10 +41,13 @@ class ActiveJob:
result_directory: Path
export_format: str = ""
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
temporary_input_path: Path | None = None
cancelled: bool = False
@@ -83,6 +94,9 @@ class RunJob(bpy.types.Operator):
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
(BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""),
(IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""),
(ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""),
),
options={"HIDDEN", "SKIP_SAVE"},
)
@@ -95,6 +109,9 @@ class RunJob(bpy.types.Operator):
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_data_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_model_key: bpy.props.StringProperty(
default="BASE",
options={"HIDDEN", "SKIP_SAVE"},
@@ -131,6 +148,14 @@ class RunJob(bpy.types.Operator):
default="MODIFIER",
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_remove_background: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
global _active_job
@@ -173,13 +198,20 @@ class RunJob(bpy.types.Operator):
else ""
),
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_remove_background,
temporary_input_path=(
Path(self.snapform_input_path)
if self.job_type == IMAGE_DEPTH_JOB
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}
and self.snapform_delete_input
else None
),
@@ -196,6 +228,9 @@ class RunJob(bpy.types.Operator):
def _prepare_command(self, _context, settings):
paths = runtime.new_job_paths(self.job_type)
remove_background = bool(
getattr(self, "snapform_remove_background", False)
)
if self.job_type == "INSTALL_RUNTIME":
self._require_online_access()
return runtime.install_runtime_command(paths["status"]), paths
@@ -205,20 +240,49 @@ class RunJob(bpy.types.Operator):
if self.job_type == "DOWNLOAD_MODEL":
self._require_online_access()
model = get_model(self.model_key)
model = get_downloadable_model(self.model_key)
return (
runtime.download_model_command(model.key, paths["status"]),
paths,
)
configured_input = (
self.snapform_input_path
if self.job_type == IMAGE_DEPTH_JOB
else settings.input_path
)
if self.job_type in {
IMAGE_DEPTH_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}:
configured_input = self.snapform_input_path
elif self.job_type == BACKGROUND_REMOVAL_JOB:
configured_input = settings.background_input_path
else:
configured_input = settings.input_path
input_path = Path(bpy.path.abspath(configured_input))
if not input_path.is_file():
raise ValueError("Select a valid image or video file")
if self.job_type in {
BACKGROUND_REMOVAL_JOB,
IMAGE_BACKGROUND_JOB,
ORIGINAL_BACKGROUND_JOB,
}:
if input_path.suffix.lower() not in {
".jpg",
".jpeg",
".png",
".webp",
}:
raise ValueError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
if not runtime.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
runtime.remove_background_command(
settings,
input_path,
paths["status"],
paths["result"],
BEN2_MODEL.key,
),
paths,
)
if self.job_type not in {RECONSTRUCTION_JOB, IMAGE_DEPTH_JOB}:
raise RuntimeError(f"Unsupported task type: {self.job_type}")
export_format = (
@@ -236,6 +300,12 @@ class RunJob(bpy.types.Operator):
model = get_model(model_key)
if not runtime.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
if (
self.job_type == IMAGE_DEPTH_JOB
and remove_background
and not runtime.model_ready(BEN2_MODEL.key)
):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
if export_format == "gs_ply" and not model.supports_gs:
raise RuntimeError(
"PLY export requires a DA3 Giant or Nested Giant model"
@@ -255,6 +325,10 @@ class RunJob(bpy.types.Operator):
else None
),
use_ray_pose=use_ray_pose,
remove_background=(
self.job_type == IMAGE_DEPTH_JOB
and remove_background
),
),
paths,
)
@@ -274,8 +348,15 @@ class RunJob(bpy.types.Operator):
return f"Starting {settings.export_format} export"
if self.job_type == IMAGE_DEPTH_JOB:
return "Starting depth estimation"
if self.job_type == BACKGROUND_REMOVAL_JOB:
return "Starting BEN2 background removal"
if self.job_type == IMAGE_BACKGROUND_JOB:
return "Starting Image Empty background removal"
if self.job_type == ORIGINAL_BACKGROUND_JOB:
return "Starting background removal"
if self.job_type == "DOWNLOAD_MODEL":
return f"Starting {get_model(self.model_key).label} download"
model = get_downloadable_model(self.model_key)
return f"Starting {model.label} download"
return labels[self.job_type]
def modal(self, context, event):
@@ -330,15 +411,18 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
if job.kind == RECONSTRUCTION_JOB:
if job.kind in {RECONSTRUCTION_JOB, BACKGROUND_REMOVAL_JOB}:
try:
open_job_directory(job.result_directory)
except Exception as open_error:
self.report({"WARNING"}, str(open_error))
self.report(
{"INFO"},
f"Exported {job.export_format} to {job.result_directory}",
)
if job.kind == RECONSTRUCTION_JOB:
message = (
f"Exported {job.export_format} to {job.result_directory}"
)
else:
message = f"Removed background to {job.result_directory}"
self.report({"INFO"}, message)
if job.kind == IMAGE_DEPTH_JOB:
try:
settings.status = "Building depth plane"
@@ -353,6 +437,34 @@ class RunJob(bpy.types.Operator):
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Depth plane is ready"
if job.kind == ORIGINAL_BACKGROUND_JOB:
try:
settings.status = "Building depth plane"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._build_original_depth_plane(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Depth plane is ready"
if job.kind == IMAGE_BACKGROUND_JOB:
try:
settings.status = "Updating Image Empty"
settings.progress = 0.96
self._redraw_ui(context, force=True)
self._apply_removed_background(job)
except (OSError, RuntimeError, ValueError) as import_error:
settings.status = "Image update failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status["message"] = "Image background is removed"
settings.status = status.get("message", "Task complete")
settings.progress = 1.0
@@ -370,9 +482,16 @@ class RunJob(bpy.types.Operator):
result = bpy.ops.snapform.build_depth_plane(
snapform_source_object_name=job.source_object_name,
snapform_npz_path=str(npz_path),
snapform_depth_data_name=job.depth_data_name,
snapform_background_path=(
str(job.result_directory / "foreground.png")
if job.remove_background
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:
raise RuntimeError("Unable to build the depth plane")
@@ -380,6 +499,37 @@ class RunJob(bpy.types.Operator):
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the depth plane Undo step")
def _build_original_depth_plane(self, job):
result = bpy.ops.snapform.build_original_depth_plane(
snapform_source_object_name=job.source_object_name,
snapform_background_path=str(
job.result_directory / "foreground.png"
),
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:
raise RuntimeError("Unable to build the depth plane")
undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the depth plane Undo step")
def _apply_removed_background(self, job, push_undo=True):
output_path = job.result_directory / "foreground.png"
result = bpy.ops.snapform.apply_removed_background(
snapform_source_object_name=job.source_object_name,
snapform_output_path=str(output_path),
)
if "FINISHED" not in result:
raise RuntimeError("Unable to update the Image Empty")
if not push_undo:
return
undo_result = bpy.ops.ed.undo_push(message="Remove Background")
if "FINISHED" not in undo_result:
raise RuntimeError("Unable to create the background removal Undo step")
def _remove_timer(self, context):
timer = getattr(self, "_timer", None)
if timer is None:

15
src/snapform/panel.py
View File

@@ -1,6 +1,6 @@
import bpy
from .models import MODELS
from .models import BEN2_MODEL, MODELS
from .operators import active_job
from .runtime import model_ready, runtime_ready
@@ -29,6 +29,7 @@ class RuntimeModelsPanel(bpy.types.Panel):
for model in MODELS.values():
self._draw_model(layout, model, running)
self._draw_model(layout, BEN2_MODEL, running)
def _draw_model(self, layout, model, running):
row = layout.row(align=True)
@@ -106,3 +107,15 @@ class DebugPanel(bpy.types.Panel):
icon="EXPORT",
)
operator.job_type = "RECONSTRUCT"
background_box = layout.box()
background_box.label(text="BEN2 Background Removal")
background_box.prop(settings, "background_input_path")
background_row = background_box.row()
background_row.enabled = active_job() is None
operator = background_row.operator(
"snapform.run_job",
text="Remove Background",
icon="IMAGE_ALPHA",
)
operator.job_type = "REMOVE_BACKGROUND"

5
src/snapform/properties.py
View File

@@ -9,6 +9,11 @@ class SnapFormSettings(bpy.types.PropertyGroup):
description="A single image or video file used for reconstruction",
subtype="FILE_PATH",
)
background_input_path: bpy.props.StringProperty(
name="Image",
description="A single image used to test BEN2 background removal",
subtype="FILE_PATH",
)
video_frames: bpy.props.IntProperty(
name="Video Frames",
description="Maximum number of frames sampled evenly from a video",

49
src/snapform/runtime.py
View File

@@ -8,11 +8,12 @@ from pathlib import Path
import bpy
from .models import get_model, selected_model_key
from .models import BEN2_MODEL, get_downloadable_model, selected_model_key
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
RUNTIME_SCHEMA = 2
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
RUNTIME_SCHEMA = 3
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
@@ -127,14 +128,16 @@ def install_runtime_command(status_path):
backend.lower(),
"--da3-commit",
DA3_COMMIT,
"--ben2-commit",
BEN2_COMMIT,
"--status",
str(status_path),
]
def download_model_command(model_key, status_path):
model = get_model(model_key)
return [
model = get_downloadable_model(model_key)
command = [
str(runtime_python()),
str(Path(__file__).with_name("worker.py")),
"download",
@@ -145,6 +148,33 @@ def download_model_command(model_key, status_path):
"--status",
str(status_path),
]
for pattern in model.download_patterns:
command.extend(["--allow-pattern", pattern])
return command
def remove_background_command(
settings,
input_path,
status_path,
result_directory,
model_key,
):
return [
str(runtime_python()),
str(Path(__file__).with_name("worker.py")),
"remove-background",
"--input",
str(input_path),
"--model-dir",
str(model_directory(model_key)),
"--output-dir",
str(result_directory),
"--status",
str(status_path),
"--device",
settings.device.lower(),
]
def reconstruct_command(
@@ -157,6 +187,7 @@ def reconstruct_command(
process_res=None,
process_res_method=None,
use_ray_pose=None,
remove_background=False,
):
snapform_selected_key = model_key or selected_model_key(settings, export_format)
snapform_selected_process_res = (
@@ -170,7 +201,7 @@ def reconstruct_command(
snapform_selected_ray_pose = (
settings.use_ray_pose if use_ray_pose is None else use_ray_pose
)
return [
command = [
str(runtime_python()),
str(Path(__file__).with_name("worker.py")),
"reconstruct",
@@ -209,6 +240,14 @@ def reconstruct_command(
"--ref-view-strategy",
settings.ref_view_strategy,
]
if remove_background:
command.extend(
[
"--remove-background-model-dir",
str(model_directory(BEN2_MODEL.key)),
]
)
return command
def creation_flags():

112
src/snapform/worker.py
View File

@@ -188,6 +188,7 @@ def download_model(args):
repo_id=args.repository,
local_dir=model_dir,
tqdm_class=progress_class,
allow_patterns=args.allow_pattern or None,
)
write_status(args.status, 1.0, "Model is ready")
@@ -202,6 +203,76 @@ def select_device(requested, torch):
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
def infer_foreground(input_path, model_dir, device, torch):
from ben2 import BEN_Base
from PIL import Image
input_path = Path(input_path)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
load_started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
load_ms = (time.perf_counter() - load_started) * 1000.0
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
if device.type == "cuda":
torch.cuda.synchronize()
inference_started = time.perf_counter()
with torch.inference_mode():
foreground = model.inference(
image,
refine_foreground=False,
)
if device.type == "cuda":
torch.cuda.synchronize()
inference_ms = (time.perf_counter() - inference_started) * 1000.0
del model
if device.type == "cuda":
torch.cuda.empty_cache()
return foreground, load_ms, inference_ms
def remove_background(args):
total_started = time.perf_counter()
write_status(args.status, 0.02, "Starting BEN2 Runtime")
import torch
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=False)
device = select_device(args.device, torch)
write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
write_status(args.status, 0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
args.input,
args.model_dir,
device,
torch,
)
write_status(args.status, 0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
total_ms = (time.perf_counter() - total_started) * 1000.0
message = (
f"Background removed in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
)
print(message)
write_status(args.status, 1.0, message)
def exported_paths(output_directory, export_format):
output_directory = Path(output_directory)
if export_format == "glb":
@@ -246,6 +317,30 @@ def reconstruct(args):
try:
device = select_device(args.device, torch)
remove_background_model_dir = getattr(
args,
"remove_background_model_dir",
"",
)
if remove_background_model_dir:
write_status(args.status, 0.12, "Removing image background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
remove_background_model_dir,
device,
torch,
)
foreground.convert("RGBA").save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
if not (output_directory / "foreground.png").is_file():
raise RuntimeError("BEN2 did not write a foreground image")
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)
write_status(args.status, 0.15, f"Loading model on {device.type}")
model = DepthAnything3.from_pretrained(args.model_dir)
model = model.to(device=device)
@@ -311,6 +406,18 @@ def parser():
download.add_argument("--repository", required=True)
download.add_argument("--model-dir", required=True)
download.add_argument("--status", required=True)
download.add_argument("--allow-pattern", action="append", default=[])
background = commands.add_parser("remove-background")
background.add_argument("--input", required=True)
background.add_argument("--model-dir", required=True)
background.add_argument("--output-dir", required=True)
background.add_argument("--status", required=True)
background.add_argument(
"--device",
choices=("auto", "cuda", "cpu"),
required=True,
)
reconstruction = commands.add_parser("reconstruct")
reconstruction.add_argument("--input", required=True)
@@ -322,6 +429,7 @@ def parser():
required=True,
)
reconstruction.add_argument("--status", required=True)
reconstruction.add_argument("--remove-background-model-dir", default="")
reconstruction.add_argument("--process-res", type=int, required=True)
reconstruction.add_argument(
"--process-res-method",
@@ -370,8 +478,10 @@ def main():
try:
if args.command == "download":
download_model(args)
else:
elif args.command == "reconstruct":
reconstruct(args)
else:
remove_background(args)
except Exception as error:
write_status(args.status, 1.0, "Task failed", str(error))
raise