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

494
tests/test_addon.py
View File

@@ -118,12 +118,26 @@ class AddonTest(unittest.TestCase):
)
self.assertFalse(hasattr(self.addon, "ImportDepthPlane"))
def test_registers_image_background_operators(self):
self.assertIn(self.addon.RemoveImageBackground, self.addon.CLASSES)
self.assertIn(self.addon.ApplyRemovedBackground, self.addon.CLASSES)
self.assertIn(self.addon.BuildOriginalDepthPlane, self.addon.CLASSES)
self.assertEqual(
self.addon.ApplyRemovedBackground.bl_options,
{"INTERNAL"},
)
def test_completed_depth_job_uses_undoable_build_operator(self):
source = inspect.getsource(
self.addon.RunJob._build_depth_plane
)
self.assertIn("bpy.ops.snapform.build_depth_plane(", source)
self.assertIn("snapform_unshaded=job.unshaded", source)
self.assertIn(
"snapform_depth_data_name=job.depth_data_name",
source,
)
self.assertIn('if "FINISHED" not in result:', source)
self.assertIn(
'bpy.ops.ed.undo_push(message="Convert to Depth Plane")',
@@ -131,7 +145,7 @@ class AddonTest(unittest.TestCase):
)
self.assertIn('if "FINISHED" not in undo_result:', source)
def test_image_context_menu_only_exposes_conversion(self):
def test_image_context_menu_exposes_depth_and_background_actions(self):
draw_source = inspect.getsource(
self.addon.draw_image_context_menu
)
@@ -140,10 +154,185 @@ class AddonTest(unittest.TestCase):
'"snapform.convert_image_to_depth_plane"',
draw_source,
)
self.assertIn(
'"snapform.remove_image_background"',
draw_source,
)
self.assertIn('text="Remove Background"', draw_source)
self.assertEqual(draw_source.count("row = self.layout.row()"), 2)
self.assertIn('operator_context = "INVOKE_DEFAULT"', draw_source)
self.assertNotIn("import_depth_plane", draw_source)
self.assertNotIn("Rebuild", draw_source)
def test_background_result_creates_new_packed_image(self):
original_image = SimpleNamespace(
name="img.png",
use_fake_user=False,
)
transform = object()
source_object = SimpleNamespace(
data=original_image,
matrix_world=transform,
empty_display_size=2.5,
empty_image_offset=(-0.25, -0.75),
)
pack_calls = []
result_image = SimpleNamespace(
name="foreground.png",
pack=lambda: pack_calls.append(True),
)
load_calls = []
remove_calls = []
self.fake_bpy.data = SimpleNamespace(
images=SimpleNamespace(
load=lambda *args, **options: load_calls.append(
(args, options)
)
or result_image,
remove=remove_calls.append,
),
)
with TemporaryDirectory() as directory:
output_path = Path(directory) / "foreground.png"
output_path.touch()
created = self.addon.image_depth.create_removed_background_image(
source_object,
output_path,
)
self.assertIs(created, result_image)
self.assertIs(source_object.data, result_image)
self.assertIsNot(source_object.data, original_image)
self.assertTrue(original_image.use_fake_user)
self.assertIs(source_object.matrix_world, transform)
self.assertEqual(source_object.empty_display_size, 2.5)
self.assertEqual(
source_object.empty_image_offset,
(-0.25, -0.75),
)
self.assertEqual(result_image.name, "img.rb")
self.assertEqual(pack_calls, [True])
self.assertEqual(remove_calls, [])
self.assertEqual(
load_calls[0][1],
{"check_existing": False},
)
def test_background_result_name_handles_numbered_format_suffix(self):
source_object = SimpleNamespace(
data=SimpleNamespace(name="img.png.001"),
)
self.assertEqual(
self.addon.image_depth.removed_background_data_name(
source_object
),
"img.001.rb",
)
def test_depth_background_image_does_not_replace_source_empty(self):
original_image = SimpleNamespace(
name="img.png",
use_fake_user=False,
)
source_object = SimpleNamespace(data=original_image)
pack_calls = []
result_image = SimpleNamespace(
name="foreground.png",
pack=lambda: pack_calls.append(True),
)
self.fake_bpy.data = SimpleNamespace(
images=SimpleNamespace(
load=lambda *_args, **_options: result_image,
remove=lambda _image: None,
),
)
with TemporaryDirectory() as directory:
output_path = Path(directory) / "foreground.png"
output_path.touch()
created = self.addon.image_depth.create_removed_background_image(
source_object,
output_path,
assign_to_source=False,
)
self.assertIs(created, result_image)
self.assertIs(source_object.data, original_image)
self.assertFalse(original_image.use_fake_user)
self.assertEqual(result_image.name, "img.rb")
self.assertEqual(pack_calls, [True])
def test_image_empty_background_action_starts_production_job(self):
source_object = SimpleNamespace(
name="Reference",
data=SimpleNamespace(name="img.png"),
)
context = SimpleNamespace(object=source_object, scene=object())
calls = []
self.fake_bpy.ops = SimpleNamespace(
snapform=SimpleNamespace(
run_job=lambda **options: calls.append(options)
or {"RUNNING_MODAL"},
),
)
operator = self.addon.RemoveImageBackground()
with TemporaryDirectory() as directory:
input_path = Path(directory) / "input.png"
input_path.touch()
with patch.object(
self.addon.image_depth,
"prepare_image_input",
return_value=(input_path, True),
):
result = operator.execute(context)
self.assertEqual(result, {"FINISHED"})
self.assertEqual(calls[0]["job_type"], "IMAGE_BACKGROUND")
self.assertEqual(calls[0]["snapform_input_path"], str(input_path))
self.assertEqual(
calls[0]["snapform_source_object_name"],
"Reference",
)
self.assertTrue(calls[0]["snapform_delete_input"])
def test_original_depth_background_starts_independent_job(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ORIGINAL"
operator.snapform_remove_background = True
operator.snapform_displacement_via = "MODIFIER"
operator.snapform_unshaded = False
operator.snapform_subdivision_level = 5
operator.snapform_delete_original = False
source_object = SimpleNamespace(
name="Reference",
data=SimpleNamespace(name="img.png"),
)
context = SimpleNamespace(object=source_object, scene=object())
calls = []
self.fake_bpy.ops = SimpleNamespace(
snapform=SimpleNamespace(
run_job=lambda **options: calls.append(options)
or {"RUNNING_MODAL"},
),
)
with TemporaryDirectory() as directory:
input_path = Path(directory) / "input.png"
input_path.touch()
with patch.object(
self.addon.image_depth,
"prepare_image_input",
return_value=(input_path, False),
):
result = operator.execute(context)
self.assertEqual(result, {"FINISHED"})
self.assertEqual(calls[0]["job_type"], "ORIGINAL_BACKGROUND")
self.assertFalse(calls[0]["snapform_delete_source"])
def test_conversion_dialog_exposes_requested_settings(self):
operator = self.addon.ConvertImageToDepthPlane
@@ -159,6 +348,12 @@ class AddonTest(unittest.TestCase):
self.assertTrue(
operator.__annotations__["snapform_delete_original"]["default"]
)
self.assertFalse(
operator.__annotations__["snapform_unshaded"]["default"]
)
self.assertFalse(
operator.__annotations__["snapform_remove_background"]["default"]
)
self.assertEqual(
operator.__annotations__["snapform_subdivision_level"]["default"],
5,
@@ -168,6 +363,8 @@ class AddonTest(unittest.TestCase):
{
"snapform_depth_source",
"snapform_displacement_via",
"snapform_unshaded",
"snapform_remove_background",
"snapform_model",
"snapform_use_ray_pose",
"snapform_process_res_factor",
@@ -192,6 +389,10 @@ class AddonTest(unittest.TestCase):
[item[0] for item in displacement_via["items"]],
["MODIFIER", "MATERIAL"],
)
self.assertEqual(
operator.__annotations__["snapform_unshaded"]["name"],
"Unshaded",
)
draw_source = inspect.getsource(operator.draw)
self.assertIn(
'if self.snapform_depth_source == "ESTIMATE":',
@@ -218,8 +419,16 @@ class AddonTest(unittest.TestCase):
'estimate_box.prop(self, "snapform_use_ray_pose")',
draw_source,
)
self.assertIn(
'layout.prop(self, "snapform_remove_background")',
draw_source,
)
self.assertNotIn("layout.separator", draw_source)
self.assertNotIn("layout.label", draw_source)
self.assertIn(
'layout.prop(self, "snapform_unshaded")',
draw_source,
)
model_items = operator.__annotations__["snapform_model"]["items"]
model_items = model_items(None, None)
self.assertEqual(
@@ -266,6 +475,12 @@ class AddonTest(unittest.TestCase):
self.assertNotIn("DA3 image depth", starting_source)
self.assertNotIn("DA3 depth plane", finish_source)
def test_depth_background_is_loaded_without_changing_source_empty(self):
source = inspect.getsource(self.addon.RunJob._build_depth_plane)
self.assertIn("snapform_background_path=", source)
self.assertNotIn("_apply_removed_background", source)
def test_model_options_mark_only_models_that_are_not_downloaded(self):
items_callback = self.addon.SnapFormSettings.__annotations__["model"][
"items"
@@ -288,9 +503,12 @@ class AddonTest(unittest.TestCase):
operator.snapform_use_ray_pose = True
operator.snapform_process_res_factor = 0.5
operator.snapform_displacement_via = "MATERIAL"
operator.snapform_unshaded = True
operator.snapform_remove_background = True
operator.snapform_subdivision_level = 5
operator.snapform_delete_original = True
image = SimpleNamespace(
name="img.png",
filepath="",
packed_file=object(),
size=(1920, 1080),
@@ -326,6 +544,9 @@ class AddonTest(unittest.TestCase):
calls[0]["snapform_displacement_via"],
"MATERIAL",
)
self.assertTrue(calls[0]["snapform_unshaded"])
self.assertTrue(calls[0]["snapform_remove_background"])
self.assertEqual(calls[0]["snapform_depth_data_name"], "img.depth")
self.assertTrue(calls[0]["snapform_delete_input"])
self.assertEqual(
self.addon.image_depth._last_depth_plane_model,
@@ -454,6 +675,81 @@ class AddonTest(unittest.TestCase):
snapform_ray_pose,
)
def test_image_depth_background_option_requires_ben2(self):
settings = SimpleNamespace(
input_path="debug-input.png",
export_format="mini_npz",
model="BASE",
gs_model="GIANT_1_1",
)
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB
operator.snapform_model_key = "BASE"
operator.snapform_use_ray_pose = True
operator.snapform_process_res = 1024
operator.snapform_remove_background = True
with TemporaryDirectory() as directory:
input_path = Path(directory) / "input.png"
input_path.touch()
operator.snapform_input_path = str(input_path)
with patch.object(
self.addon.runtime,
"new_job_paths",
return_value={
"log": Path(directory) / "job.log",
"status": Path(directory) / "job.json",
"result": Path(directory) / "result",
},
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
side_effect=lambda key: key == "BASE",
):
with self.assertRaisesRegex(RuntimeError, "BEN2 Base"):
operator._prepare_command(None, settings)
def test_image_depth_background_option_reaches_worker_command(self):
runtime = importlib.import_module("snapform.runtime")
settings = SimpleNamespace(
input_path="input.png",
model="BASE",
gs_model="GIANT_1_1",
process_res=504,
process_res_method="upper_bound_resize",
conf_thresh_percentile=40.0,
num_max_points=1_000_000,
show_cameras=True,
gs_views_interval=1,
video_frames=16,
device="AUTO",
use_ray_pose=False,
ref_view_strategy="saddle_balanced",
)
with patch.object(
runtime,
"runtime_python",
return_value=Path("python"),
), patch.object(
runtime,
"model_directory",
side_effect=lambda key: Path(key.lower()),
):
command = runtime.reconstruct_command(
settings,
"mini_npz",
Path("status.json"),
Path("result"),
remove_background=True,
)
option_index = command.index("--remove-background-model-dir")
self.assertEqual(command[option_index + 1], str(Path("ben2_base")))
def test_preferences_use_package_identifier(self):
self.assertEqual(
self.addon.SnapFormPreferences.bl_idname,
@@ -625,6 +921,13 @@ class AddonTest(unittest.TestCase):
self.assertNotIn('text="Missing"', draw_source)
self.assertNotIn("License:", draw_source)
self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source)
self.assertIn('text="BEN2 Background Removal"', draw_source)
self.assertIn(
'background_box.prop(settings, "background_input_path")',
draw_source,
)
self.assertIn('text="Remove Background"', draw_source)
self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', draw_source)
def test_runtime_models_panel_assigns_each_download_target(self):
draw_source = inspect.getsource(
@@ -635,6 +938,10 @@ class AddonTest(unittest.TestCase):
)
self.assertIn("for model in MODELS.values()", draw_source)
self.assertIn(
"self._draw_model(layout, BEN2_MODEL, running)",
draw_source,
)
self.assertIn("operator.model_key = model.key", model_source)
self.assertIn(
"download_row.enabled = runtime_ready() and not running",
@@ -648,6 +955,191 @@ class AddonTest(unittest.TestCase):
"Setup",
)
def test_ben2_debug_job_uses_its_image_and_model(self):
settings = SimpleNamespace(
background_input_path="input.png",
device="CUDA",
)
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.BACKGROUND_REMOVAL_JOB
with TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
settings.background_input_path = str(input_path)
paths = {
"log": root / "job.log",
"status": root / "job.json",
"result": root / "result",
}
with patch.object(
self.addon.runtime,
"new_job_paths",
return_value=paths,
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
) as model_ready, patch.object(
self.addon.runtime,
"remove_background_command",
return_value=["worker"],
) as remove_background:
command, prepared_paths = operator._prepare_command(
None,
settings,
)
self.assertEqual(command, ["worker"])
self.assertEqual(prepared_paths, paths)
model_ready.assert_called_once_with("BEN2_BASE")
self.assertEqual(remove_background.call_args.args[1], input_path)
self.assertEqual(remove_background.call_args.args[4], "BEN2_BASE")
def test_image_background_job_uses_operator_input(self):
settings = SimpleNamespace(device="AUTO")
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB
with TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
operator.snapform_input_path = str(input_path)
paths = {
"log": root / "job.log",
"status": root / "job.json",
"result": root / "result",
}
with patch.object(
self.addon.runtime,
"new_job_paths",
return_value=paths,
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"remove_background_command",
return_value=["worker"],
) as remove_background:
command, prepared_paths = operator._prepare_command(
None,
settings,
)
self.assertEqual(command, ["worker"])
self.assertEqual(prepared_paths, paths)
self.assertEqual(remove_background.call_args.args[1], input_path)
def test_completed_image_background_job_applies_result_with_undo(self):
calls = []
self.fake_bpy.ops = SimpleNamespace(
snapform=SimpleNamespace(
apply_removed_background=lambda **options: calls.append(
("apply", options)
)
or {"FINISHED"},
),
ed=SimpleNamespace(
undo_push=lambda **options: calls.append(("undo", options))
or {"FINISHED"},
),
)
operator = self.addon.RunJob()
job = SimpleNamespace(
result_directory=Path("result"),
source_object_name="Reference",
)
operator._apply_removed_background(job)
self.assertEqual(
calls,
[
(
"apply",
{
"snapform_source_object_name": "Reference",
"snapform_output_path": str(
Path("result") / "foreground.png"
),
},
),
("undo", {"message": "Remove Background"}),
],
)
def test_completed_original_background_job_builds_with_one_undo(self):
calls = []
self.fake_bpy.ops = SimpleNamespace(
snapform=SimpleNamespace(
build_original_depth_plane=lambda **options: calls.append(
("build", options)
)
or {"FINISHED"},
),
ed=SimpleNamespace(
undo_push=lambda **options: calls.append(("undo", options))
or {"FINISHED"},
),
)
operator = self.addon.RunJob()
job = SimpleNamespace(
result_directory=Path("result"),
source_object_name="Reference",
subdivision_level=5,
delete_source=False,
displacement_via="MODIFIER",
unshaded=False,
)
operator._build_original_depth_plane(job)
self.assertEqual(calls[0][0], "build")
self.assertEqual(
calls[0][1]["snapform_background_path"],
str(Path("result") / "foreground.png"),
)
self.assertFalse(calls[0][1]["snapform_delete_source"])
self.assertEqual(
calls[1],
("undo", {"message": "Convert to Depth Plane"}),
)
def test_ben2_download_command_filters_large_alternative_weights(self):
runtime = importlib.import_module("snapform.runtime")
with patch.object(
runtime,
"runtime_python",
return_value=Path("python"),
), patch.object(
runtime,
"model_directory",
return_value=Path("model"),
):
command = runtime.download_model_command(
"BEN2_BASE",
Path("status.json"),
)
patterns = [
command[index + 1]
for index, value in enumerate(command)
if value == "--allow-pattern"
]
self.assertEqual(patterns, ["config.json", "model.safetensors"])
def test_preferences_install_button_is_disabled_while_busy(self):
draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw)