nan
d694980da7
- 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.
553 lines
19 KiB
Python
553 lines
19 KiB
Python
import importlib.util
|
|
import inspect
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_image_depth():
|
|
fake_bpy = ModuleType("bpy")
|
|
fake_bpy.props = SimpleNamespace(
|
|
BoolProperty=lambda **options: options,
|
|
EnumProperty=lambda **options: options,
|
|
FloatProperty=lambda **options: options,
|
|
IntProperty=lambda **options: options,
|
|
StringProperty=lambda **options: options,
|
|
)
|
|
fake_bpy.types = SimpleNamespace(Operator=object)
|
|
fake_bpy.path = SimpleNamespace(abspath=lambda path: path)
|
|
with patch.dict(sys.modules, {"bpy": fake_bpy}):
|
|
specification = importlib.util.spec_from_file_location(
|
|
"snapform_image_depth_test",
|
|
PROJECT_ROOT / "src" / "snapform" / "image_depth.py",
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class ImageDepthTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.image_depth = load_image_depth()
|
|
|
|
def test_extracts_depth_and_uses_first_view(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
npz_path = Path(directory) / "results.npz"
|
|
np.savez_compressed(
|
|
npz_path,
|
|
depth=np.array(
|
|
(
|
|
((1.0, 2.0), (3.0, 4.0)),
|
|
((5.0, 6.0), (7.0, 8.0)),
|
|
),
|
|
dtype=np.float32,
|
|
),
|
|
)
|
|
|
|
depth = self.image_depth.extract_depth_array(npz_path)
|
|
|
|
np.testing.assert_array_equal(
|
|
depth,
|
|
np.array(((1.0, 2.0), (3.0, 4.0)), dtype=np.float32),
|
|
)
|
|
|
|
def test_accepts_depths_key_with_single_channel(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
npz_path = Path(directory) / "results.npz"
|
|
np.savez_compressed(
|
|
npz_path,
|
|
depths=np.ones((1, 2, 3, 1), dtype=np.float32),
|
|
)
|
|
|
|
depth = self.image_depth.extract_depth_array(npz_path)
|
|
|
|
self.assertEqual(depth.shape, (2, 3))
|
|
|
|
def test_accepts_single_depth_map_with_channel_axis(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
npz_path = Path(directory) / "results.npz"
|
|
np.savez_compressed(
|
|
npz_path,
|
|
depth=np.ones((2, 3, 1), dtype=np.float32),
|
|
)
|
|
|
|
depth = self.image_depth.extract_depth_array(npz_path)
|
|
|
|
self.assertEqual(depth.shape, (2, 3))
|
|
|
|
def test_extracts_camera_intrinsics(self):
|
|
expected = np.array(
|
|
((100.0, 0.0, 50.0), (0.0, 120.0, 40.0), (0.0, 0.0, 1.0)),
|
|
dtype=np.float32,
|
|
)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
npz_path = Path(directory) / "results.npz"
|
|
np.savez_compressed(
|
|
npz_path,
|
|
intrinsics=expected[None, ...],
|
|
)
|
|
|
|
intrinsics = self.image_depth.extract_intrinsics_array(
|
|
npz_path
|
|
)
|
|
|
|
np.testing.assert_array_equal(intrinsics, expected)
|
|
|
|
def test_camera_plane_bounds_use_intrinsics_and_median_depth(self):
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(200, 100)),
|
|
empty_image_offset=(-0.5, -0.5),
|
|
)
|
|
depth = np.full((100, 200), 2.0, dtype=np.float32)
|
|
intrinsics = np.array(
|
|
((100.0, 0.0, 100.0), (0.0, 100.0, 50.0), (0.0, 0.0, 1.0)),
|
|
dtype=np.float32,
|
|
)
|
|
|
|
bounds, reference_depth = self.image_depth.camera_plane_bounds(
|
|
source,
|
|
depth,
|
|
intrinsics,
|
|
-1.0,
|
|
)
|
|
|
|
self.assertEqual(bounds, (-2.0, 2.0, -1.0, 1.0))
|
|
self.assertEqual(reference_depth, 2.0)
|
|
|
|
def test_camera_plane_bounds_follow_displacement_unit_scale(self):
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(200, 100)),
|
|
empty_image_offset=(-0.5, -0.5),
|
|
)
|
|
depth = np.full((100, 200), 2.0, dtype=np.float32)
|
|
intrinsics = np.array(
|
|
((100.0, 0.0, 100.0), (0.0, 100.0, 50.0), (0.0, 0.0, 1.0)),
|
|
dtype=np.float32,
|
|
)
|
|
|
|
bounds, _reference_depth = (
|
|
self.image_depth.camera_plane_bounds(
|
|
source,
|
|
depth,
|
|
intrinsics,
|
|
-0.5,
|
|
)
|
|
)
|
|
|
|
self.assertEqual(bounds, (-1.0, 1.0, -0.5, 0.5))
|
|
|
|
def test_camera_plane_bounds_preserve_square_source_aspect(self):
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(447, 447)),
|
|
empty_image_offset=(-0.5, -0.5),
|
|
)
|
|
depth = np.full((140, 140), 0.7366, dtype=np.float32)
|
|
intrinsics = np.array(
|
|
(
|
|
(460.758, 0.0, 70.0),
|
|
(0.0, 393.372, 70.0),
|
|
(0.0, 0.0, 1.0),
|
|
),
|
|
dtype=np.float32,
|
|
)
|
|
|
|
bounds, _reference_depth = (
|
|
self.image_depth.camera_plane_bounds(
|
|
source,
|
|
depth,
|
|
intrinsics,
|
|
-1.0,
|
|
)
|
|
)
|
|
width = bounds[1] - bounds[0]
|
|
height = bounds[3] - bounds[2]
|
|
|
|
self.assertAlmostEqual(width / height, 1.0)
|
|
|
|
def test_raw_distance_preserves_original_values(self):
|
|
depth = np.array(
|
|
((1.25, 2.5), (5.0, np.nan)),
|
|
dtype=np.float32,
|
|
)
|
|
|
|
mapped = self.image_depth.prepare_depth_values(depth)
|
|
|
|
self.assertEqual(float(mapped[0, 0]), 1.25)
|
|
self.assertEqual(float(mapped[0, 1]), 2.5)
|
|
self.assertEqual(float(mapped[1, 0]), 5.0)
|
|
self.assertEqual(float(mapped[1, 1]), 0.0)
|
|
|
|
def test_depth_plane_uses_fixed_import_settings(self):
|
|
self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, -1.0)
|
|
self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 1.0)
|
|
self.assertEqual(self.image_depth.BUMP_BASE_DISTANCE, 0.001)
|
|
|
|
def test_grid_segments_follow_aspect_ratio_from_subdivision_power(self):
|
|
self.assertEqual(
|
|
self.image_depth.grid_segments((0.0, 3.0, 0.0, 2.0), 1),
|
|
(3, 2),
|
|
)
|
|
self.assertEqual(
|
|
self.image_depth.grid_segments((0.0, 2.0, 0.0, 3.0), 1),
|
|
(2, 3),
|
|
)
|
|
self.assertEqual(
|
|
self.image_depth.grid_segments((0.0, 3.0, 0.0, 2.0), 3),
|
|
(12, 8),
|
|
)
|
|
|
|
def test_modifiers_use_native_names_and_fixed_levels(self):
|
|
source = inspect.getsource(self.image_depth.build_depth_plane)
|
|
|
|
self.assertIn('modifiers.new("Subdivision", "SUBSURF")', source)
|
|
self.assertIn('modifiers.new("Displace", "DISPLACE")', source)
|
|
self.assertNotIn("DA3 Subdivision", source)
|
|
self.assertNotIn("DA3 Displace", source)
|
|
self.assertIn("subdivision.levels = 1", source)
|
|
self.assertIn("subdivision.render_levels = 2", source)
|
|
|
|
def test_plane_fit_transform_matches_source_size_and_offset(self):
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(800, 400)),
|
|
empty_display_size=2.0,
|
|
empty_image_offset=(-0.25, -0.75),
|
|
)
|
|
|
|
transform = self.image_depth.plane_fit_transform(
|
|
source,
|
|
(-0.1, 0.3, -0.1, 0.1),
|
|
)
|
|
|
|
for value, expected in zip(
|
|
transform,
|
|
(5.0, 5.0, 5.0, 0.0, -0.25),
|
|
):
|
|
self.assertAlmostEqual(value, expected)
|
|
|
|
def test_source_empty_is_only_removed_when_requested(self):
|
|
source = inspect.getsource(self.image_depth.finish_depth_plane)
|
|
|
|
self.assertIn("source_name = source_object.name", source)
|
|
self.assertIn("if delete_source:", source)
|
|
self.assertIn(
|
|
"bpy.data.objects.remove(source_object, do_unlink=True)",
|
|
source,
|
|
)
|
|
self.assertNotIn("source_object.name =", source)
|
|
self.assertIn(
|
|
"bpy.data.objects.remove(source_object, do_unlink=True)\n"
|
|
" plane_object.name = source_name",
|
|
source,
|
|
)
|
|
|
|
def test_original_depth_mode_uses_uv_and_zero_midlevel(self):
|
|
source = inspect.getsource(self.image_depth.build_image_depth_plane)
|
|
|
|
self.assertIn("texture.image = source_object.data", source)
|
|
self.assertIn('modifiers.new("Displace", "DISPLACE")', source)
|
|
self.assertIn("displace.texture = texture", source)
|
|
self.assertIn('displace.texture_coords = "UV"', source)
|
|
self.assertNotIn("displace.direction", source)
|
|
self.assertNotIn("displace.strength", source)
|
|
self.assertIn(
|
|
"displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL",
|
|
source,
|
|
)
|
|
self.assertIn("plane_object.matrix_world =", source)
|
|
|
|
def test_process_resolution_scales_the_image_long_edge(self):
|
|
image = SimpleNamespace(size=(1080, 1920))
|
|
|
|
self.assertEqual(
|
|
self.image_depth.scaled_process_resolution(image, 0.5),
|
|
960,
|
|
)
|
|
|
|
def test_mesh_uses_source_object_name(self):
|
|
source = inspect.getsource(self.image_depth.create_plane_mesh)
|
|
|
|
self.assertIn("bpy.data.meshes.new(source_object.name)", source)
|
|
|
|
def test_depth_data_uses_source_image_name(self):
|
|
source_object = SimpleNamespace(
|
|
name="obj_img",
|
|
data=SimpleNamespace(name="img.png"),
|
|
)
|
|
|
|
self.assertEqual(
|
|
self.image_depth.depth_data_name(source_object),
|
|
"img.depth",
|
|
)
|
|
source_object.data.name = "img.png.001"
|
|
self.assertEqual(
|
|
self.image_depth.depth_data_name(source_object),
|
|
"img.001.depth",
|
|
)
|
|
estimate_source = inspect.getsource(
|
|
self.image_depth.build_depth_plane
|
|
)
|
|
original_source = inspect.getsource(
|
|
self.image_depth.build_image_depth_plane
|
|
)
|
|
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(
|
|
self.image_depth.build_image_depth_plane
|
|
)
|
|
|
|
self.assertIn(
|
|
"create_material(\n source_name,",
|
|
da3_source,
|
|
)
|
|
self.assertIn("create_material(\n source_name,", image_source)
|
|
self.assertNotIn("DA3DepthMaterial", da3_source)
|
|
|
|
def test_displace_texture_uses_extend(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("create_depth_image(\n depth_name,", source)
|
|
self.assertIn("bpy.data.textures.new(\n depth_name,", source)
|
|
self.assertIn('texture.extension = "EXTEND"', source)
|
|
|
|
def test_material_supports_bump_and_optional_displacement(self):
|
|
source = inspect.getsource(self.image_depth.create_material)
|
|
|
|
self.assertIn('nodes.new("ShaderNodeBump")', source)
|
|
self.assertIn(
|
|
'color_texture.outputs["Color"], bump.inputs["Height"]',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'bump.outputs["Normal"], shader.inputs["Normal"]',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'bump.inputs["Distance"].default_value = '
|
|
"BUMP_BASE_DISTANCE * depth_scale",
|
|
source,
|
|
)
|
|
self.assertIn('nodes.new("ShaderNodeFloatCurve")', source)
|
|
self.assertIn(
|
|
'roughness_curve.inputs["Value"]',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'shader.inputs["Roughness"]',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'curve_points[0].location = (0.0, 0.5)',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'curve_points[-1].location = (1.0, 0.1)',
|
|
source,
|
|
)
|
|
self.assertNotIn('shader.inputs["IOR"]', source)
|
|
self.assertIn(
|
|
'roughness_curve.location = (-300.0, 320.0)',
|
|
source,
|
|
)
|
|
self.assertIn('bump.location = (-300.0, -80.0)', source)
|
|
self.assertIn('nodes.new("ShaderNodeDisplacement")', source)
|
|
self.assertIn(
|
|
'depth_texture.outputs["Color"],\n'
|
|
' displacement.inputs["Height"]',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'displacement.outputs["Displacement"],\n'
|
|
' output.inputs["Displacement"]',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'displacement.inputs["Midlevel"].default_value',
|
|
source,
|
|
)
|
|
self.assertIn(
|
|
'displacement.inputs["Scale"].default_value',
|
|
source,
|
|
)
|
|
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
|
|
)
|
|
original_source = inspect.getsource(
|
|
self.image_depth.build_image_depth_plane
|
|
)
|
|
|
|
self.assertIn(
|
|
"displacement_midlevel=DISPLACEMENT_MIDLEVEL",
|
|
estimate_source,
|
|
)
|
|
self.assertIn(
|
|
"displacement_scale=DISPLACEMENT_STRENGTH",
|
|
estimate_source,
|
|
)
|
|
self.assertIn(
|
|
"displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL",
|
|
original_source,
|
|
)
|
|
self.assertIn(
|
|
"displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH",
|
|
original_source,
|
|
)
|
|
self.assertEqual(self.image_depth.DISPLACEMENT_MIDLEVEL, 1.0)
|
|
self.assertEqual(self.image_depth.DISPLACEMENT_STRENGTH, -1.0)
|
|
self.assertEqual(
|
|
self.image_depth.ORIGINAL_DISPLACEMENT_MIDLEVEL,
|
|
0.0,
|
|
)
|
|
self.assertEqual(
|
|
self.image_depth.ORIGINAL_DISPLACEMENT_STRENGTH,
|
|
1.0,
|
|
)
|
|
|
|
def test_material_displacement_mode_skips_displace_modifier(self):
|
|
estimate_source = inspect.getsource(
|
|
self.image_depth.build_depth_plane
|
|
)
|
|
original_source = inspect.getsource(
|
|
self.image_depth.build_image_depth_plane
|
|
)
|
|
|
|
for source in (estimate_source, original_source):
|
|
self.assertIn('if displacement_via == "MODIFIER":', source)
|
|
self.assertIn(
|
|
'if displacement_via == "MATERIAL" else None',
|
|
source,
|
|
)
|
|
self.assertIn("unshaded=unshaded", source)
|
|
|
|
def test_surface_displacement_is_set_to_displacement_only(self):
|
|
class Material:
|
|
surface_displacement_method = "BUMP"
|
|
|
|
material = Material()
|
|
self.image_depth.set_displacement_only(material)
|
|
|
|
self.assertEqual(
|
|
material.surface_displacement_method,
|
|
"DISPLACEMENT",
|
|
)
|
|
|
|
def test_plane_transform_scales_depth_axis(self):
|
|
source = inspect.getsource(self.image_depth.apply_plane_transform)
|
|
|
|
self.assertIn("(0.0, 0.0, scale_z, 0.0)", source)
|
|
self.assertIn("plane_object.scale.z", source)
|
|
|
|
def test_plane_bounds_keep_image_aspect_and_empty_offset(self):
|
|
source = SimpleNamespace(
|
|
data=SimpleNamespace(size=(800, 400)),
|
|
empty_display_size=2.0,
|
|
empty_image_offset=(-0.5, -0.5),
|
|
)
|
|
|
|
bounds = self.image_depth.plane_bounds(source)
|
|
|
|
self.assertEqual(bounds, (-1.0, 1.0, -0.5, 0.5))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|