feat: Add BEN2 background removal functionality and integrate into SnapForm
- 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:
@@ -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)
|
||||
|
||||
|
||||
@@ -31,6 +31,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
runtime_dir=str(runtime_dir),
|
||||
backend="cu128",
|
||||
da3_commit="test-commit",
|
||||
ben2_commit="test-ben2-commit",
|
||||
status=str(root / "status.json"),
|
||||
)
|
||||
with patch.object(bootstrap, "run") as mocked_run:
|
||||
@@ -58,6 +59,22 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
||||
)
|
||||
self.assertEqual(gsplat_call.args[1], {"BUILD_NO_CUDA": "1"})
|
||||
|
||||
ben2_command = next(
|
||||
command
|
||||
for command in commands
|
||||
if any("ben2 @" in str(value) for value in command)
|
||||
)
|
||||
self.assertTrue(
|
||||
any("test-ben2-commit" in str(value) for value in ben2_command)
|
||||
)
|
||||
check_command = commands[-1]
|
||||
self.assertTrue(
|
||||
any(
|
||||
"from ben2 import BEN_Base" in str(value)
|
||||
for value in check_command
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -298,9 +298,52 @@ class ImageDepthTest(unittest.TestCase):
|
||||
original_source = inspect.getsource(
|
||||
self.image_depth.build_image_depth_plane
|
||||
)
|
||||
self.assertIn("depth_name = depth_data_name(source_object)", estimate_source)
|
||||
self.assertIn(
|
||||
"depth_name = depth_name or depth_data_name(source_object)",
|
||||
estimate_source,
|
||||
)
|
||||
self.assertIn("depth_data_name(source_object)", original_source)
|
||||
|
||||
def test_removed_background_name_uses_source_image_base_name(self):
|
||||
source_object = SimpleNamespace(
|
||||
data=SimpleNamespace(name="img.png"),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
self.image_depth.removed_background_data_name(source_object),
|
||||
"img.rb",
|
||||
)
|
||||
source_object.data.name = "img.png.001"
|
||||
self.assertEqual(
|
||||
self.image_depth.removed_background_data_name(source_object),
|
||||
"img.001.rb",
|
||||
)
|
||||
|
||||
def test_depth_name_override_keeps_original_base_name(self):
|
||||
source = inspect.getsource(self.image_depth.build_depth_plane)
|
||||
|
||||
self.assertIn(
|
||||
"depth_name = depth_name or depth_data_name(source_object)",
|
||||
source,
|
||||
)
|
||||
self.assertIn(
|
||||
"color_image if color_image is not None else source_object.data",
|
||||
source,
|
||||
)
|
||||
|
||||
def test_original_depth_background_only_changes_material_color(self):
|
||||
source = inspect.getsource(self.image_depth.build_image_depth_plane)
|
||||
|
||||
self.assertIn("texture.image = source_object.data", source)
|
||||
self.assertIn(
|
||||
"source_object.data if displacement_via == \"MATERIAL\" else None",
|
||||
source,
|
||||
)
|
||||
self.assertIn(
|
||||
"color_image if color_image is not None else source_object.data",
|
||||
source,
|
||||
)
|
||||
|
||||
def test_material_uses_source_object_name(self):
|
||||
da3_source = inspect.getsource(self.image_depth.build_depth_plane)
|
||||
image_source = inspect.getsource(
|
||||
@@ -319,7 +362,10 @@ class ImageDepthTest(unittest.TestCase):
|
||||
self.image_depth.build_depth_plane
|
||||
)
|
||||
|
||||
self.assertIn("depth_name = depth_data_name(source_object)", source)
|
||||
self.assertIn(
|
||||
"depth_name = depth_name or depth_data_name(source_object)",
|
||||
source,
|
||||
)
|
||||
self.assertIn("create_depth_image(\n depth_name,", source)
|
||||
self.assertIn("bpy.data.textures.new(\n depth_name,", source)
|
||||
self.assertIn('texture.extension = "EXTEND"', source)
|
||||
@@ -385,6 +431,42 @@ class ImageDepthTest(unittest.TestCase):
|
||||
)
|
||||
self.assertIn("set_displacement_only(material)", source)
|
||||
|
||||
def test_unshaded_material_uses_alpha_mixed_emission(self):
|
||||
source = inspect.getsource(self.image_depth.create_material)
|
||||
|
||||
unshaded_start = source.index("if unshaded:")
|
||||
shaded_start = source.index("else:", unshaded_start)
|
||||
displacement_start = source.index(
|
||||
"if displacement_image is not None:",
|
||||
shaded_start,
|
||||
)
|
||||
unshaded_source = source[unshaded_start:shaded_start]
|
||||
shaded_source = source[shaded_start:displacement_start]
|
||||
|
||||
self.assertIn('nodes.new("ShaderNodeEmission")', unshaded_source)
|
||||
self.assertIn(
|
||||
'nodes.new("ShaderNodeBsdfTransparent")',
|
||||
unshaded_source,
|
||||
)
|
||||
self.assertIn('nodes.new("ShaderNodeMixShader")', unshaded_source)
|
||||
self.assertIn(
|
||||
'color_texture.outputs["Color"]',
|
||||
unshaded_source,
|
||||
)
|
||||
self.assertIn('emission.inputs["Color"]', unshaded_source)
|
||||
self.assertIn('color_texture.outputs["Alpha"]', unshaded_source)
|
||||
self.assertIn('mix_shader.inputs["Fac"]', unshaded_source)
|
||||
self.assertIn('transparent.outputs["BSDF"]', unshaded_source)
|
||||
self.assertIn('mix_shader.inputs[1]', unshaded_source)
|
||||
self.assertIn('emission.outputs["Emission"]', unshaded_source)
|
||||
self.assertIn('mix_shader.inputs[2]', unshaded_source)
|
||||
self.assertIn('mix_shader.outputs["Shader"]', unshaded_source)
|
||||
self.assertNotIn("ShaderNodeBsdfPrincipled", unshaded_source)
|
||||
self.assertNotIn("ShaderNodeBump", unshaded_source)
|
||||
self.assertNotIn("ShaderNodeFloatCurve", unshaded_source)
|
||||
self.assertIn("ShaderNodeBsdfPrincipled", shaded_source)
|
||||
self.assertGreater(displacement_start, unshaded_start)
|
||||
|
||||
def test_displacement_values_follow_depth_source(self):
|
||||
estimate_source = inspect.getsource(
|
||||
self.image_depth.build_depth_plane
|
||||
@@ -434,6 +516,7 @@ class ImageDepthTest(unittest.TestCase):
|
||||
'if displacement_via == "MATERIAL" else None',
|
||||
source,
|
||||
)
|
||||
self.assertIn("unshaded=unshaded", source)
|
||||
|
||||
def test_surface_displacement_is_set_to_displacement_only(self):
|
||||
class Material:
|
||||
|
||||
@@ -43,3 +43,14 @@ class ModelTest(unittest.TestCase):
|
||||
def test_unknown_model_is_rejected(self):
|
||||
with self.assertRaisesRegex(ValueError, "Unknown DA3 model"):
|
||||
self.models.get_model("UNKNOWN")
|
||||
|
||||
def test_ben2_download_only_includes_required_hugging_face_files(self):
|
||||
model = self.models.get_downloadable_model("BEN2_BASE")
|
||||
|
||||
self.assertEqual(model.repository, "PramaLLC/BEN2")
|
||||
self.assertEqual(model.license_name, "MIT")
|
||||
self.assertEqual(
|
||||
model.download_patterns,
|
||||
("config.json", "model.safetensors"),
|
||||
)
|
||||
self.assertNotIn("BEN2_BASE", self.models.MODELS)
|
||||
|
||||
@@ -29,7 +29,7 @@ class PackTest(unittest.TestCase):
|
||||
finally:
|
||||
pack.DIST_DIR = original_dist
|
||||
|
||||
self.assertEqual(output_path.name, "SnapForm.v0.9.27.zip")
|
||||
self.assertEqual(output_path.name, "SnapForm.v0.9.28.zip")
|
||||
with ZipFile(output_path) as archive:
|
||||
names = set(archive.namelist())
|
||||
|
||||
|
||||
@@ -47,6 +47,52 @@ class WorkerTest(unittest.TestCase):
|
||||
source.index("import torch"),
|
||||
)
|
||||
|
||||
def test_optional_background_removal_runs_before_da3_model_load(self):
|
||||
source = inspect.getsource(self.worker.reconstruct)
|
||||
|
||||
self.assertLess(
|
||||
source.index("infer_foreground("),
|
||||
source.index("DepthAnything3.from_pretrained"),
|
||||
)
|
||||
args = self.worker.parser().parse_args(
|
||||
[
|
||||
"reconstruct",
|
||||
"--input",
|
||||
"input.png",
|
||||
"--model-dir",
|
||||
"da3",
|
||||
"--output-dir",
|
||||
"result",
|
||||
"--export-format",
|
||||
"mini_npz",
|
||||
"--status",
|
||||
"status.json",
|
||||
"--remove-background-model-dir",
|
||||
"ben2",
|
||||
"--process-res",
|
||||
"1024",
|
||||
"--process-res-method",
|
||||
"upper_bound_resize",
|
||||
"--conf-thresh-percentile",
|
||||
"40",
|
||||
"--num-max-points",
|
||||
"1000",
|
||||
"--show-cameras",
|
||||
"0",
|
||||
"--gs-views-interval",
|
||||
"1",
|
||||
"--video-frames",
|
||||
"16",
|
||||
"--device",
|
||||
"auto",
|
||||
"--use-ray-pose",
|
||||
"1",
|
||||
"--ref-view-strategy",
|
||||
"first",
|
||||
]
|
||||
)
|
||||
self.assertEqual(args.remove_background_model_dir, "ben2")
|
||||
|
||||
def test_prepares_a_single_image_file(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
image_path = Path(directory) / "input.png"
|
||||
@@ -69,6 +115,117 @@ class WorkerTest(unittest.TestCase):
|
||||
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
|
||||
self.worker.prepare_input(input_path, 16, directory)
|
||||
|
||||
def test_download_passes_optional_file_patterns(self):
|
||||
calls = []
|
||||
fake_hugging_face = ModuleType("huggingface_hub")
|
||||
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
||||
options
|
||||
)
|
||||
args = SimpleNamespace(
|
||||
repository="PramaLLC/BEN2",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
allow_pattern=["config.json", "model.safetensors"],
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||
sys.modules,
|
||||
{"huggingface_hub": fake_hugging_face},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"create_download_progress_class",
|
||||
return_value=object,
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
self.worker.download_model(args)
|
||||
|
||||
self.assertEqual(
|
||||
calls[0]["allow_patterns"],
|
||||
["config.json", "model.safetensors"],
|
||||
)
|
||||
|
||||
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
||||
inference_options = []
|
||||
save_options = []
|
||||
|
||||
class FakeImage:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_args):
|
||||
return None
|
||||
|
||||
def convert(self, mode):
|
||||
self.mode = mode
|
||||
return self
|
||||
|
||||
def save(self, path, **options):
|
||||
save_options.append(options)
|
||||
Path(path).touch()
|
||||
|
||||
class FakeModel:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_dir):
|
||||
self.assertEqual(model_dir, "model")
|
||||
return cls()
|
||||
|
||||
def to(model_self, device):
|
||||
self.assertEqual(device.type, "cpu")
|
||||
return model_self
|
||||
|
||||
def eval(self):
|
||||
return None
|
||||
|
||||
def inference(self, image, **options):
|
||||
inference_options.append((image, options))
|
||||
return FakeImage()
|
||||
|
||||
fake_torch = ModuleType("torch")
|
||||
fake_torch.cuda = SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
synchronize=lambda: None,
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("ben2")
|
||||
fake_ben2.BEN_Base = FakeModel
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"ben2": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
root = Path(directory)
|
||||
input_path = root / "input.png"
|
||||
input_path.touch()
|
||||
args = SimpleNamespace(
|
||||
input=str(input_path),
|
||||
model_dir="model",
|
||||
output_dir=str(root / "result"),
|
||||
status=str(root / "status.json"),
|
||||
device="cpu",
|
||||
)
|
||||
self.worker.remove_background(args)
|
||||
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
||||
|
||||
self.assertTrue((root / "result" / "foreground.png").is_file())
|
||||
|
||||
self.assertEqual(
|
||||
inference_options[0][1],
|
||||
{"refine_foreground": False},
|
||||
)
|
||||
self.assertEqual(
|
||||
save_options[0],
|
||||
{"format": "PNG", "compress_level": 1},
|
||||
)
|
||||
self.assertIn("Background removed in", status["message"])
|
||||
|
||||
def test_uses_official_da3_export_parameters_for_all_four_formats(self):
|
||||
inference_options = []
|
||||
export_calls = []
|
||||
|
||||
Reference in New Issue
Block a user