Files
AnyImage/tests/test_image_depth.py
nan 8aac4822c6 feat: Implement background depth distance feature for depth processing
- Added `background_depth_distance` property to `ActiveJob` and `RunJob` classes to manage background depth settings.
- Updated job execution to handle background depth distance, allowing for fixed negative displacement in depth calculations.
- Introduced new functions in the worker module to apply background depth distance to depth frames and manage original background depth.
- Enhanced tests to cover new functionality, ensuring correct behavior of background depth processing and reference depth extraction.
- Refactored job submission to run in a separate thread, improving responsiveness during execution.
2026-08-09 16:34:41 +08:00

640 lines
20 KiB
Python

import importlib.util
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_extract_depth_arrays_returns_every_frame(self):
with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz"
np.savez_compressed(
npz_path,
depth=np.arange(12, dtype=np.float32).reshape(3, 2, 2),
)
depth_frames = self.image_depth.extract_depth_arrays(npz_path)
self.assertEqual(len(depth_frames), 3)
np.testing.assert_array_equal(
depth_frames[1],
np.array(((4.0, 5.0), (6.0, 7.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_extracts_optional_reference_depth(self):
with tempfile.TemporaryDirectory() as directory:
npz_path = Path(directory) / "results.npz"
np.savez_compressed(
npz_path,
depth=np.ones((1, 2, 2), dtype=np.float32),
snapform_reference_depth=np.float32(2.5),
)
reference_depth = self.image_depth.extract_reference_depth(
npz_path
)
self.assertEqual(reference_depth, 2.5)
def test_is_animated_image_accepts_movie_and_sequence(self):
self.assertTrue(
self.image_depth.is_animated_image(
SimpleNamespace(source="MOVIE")
)
)
self.assertTrue(
self.image_depth.is_animated_image(
SimpleNamespace(source="SEQUENCE")
)
)
self.assertFalse(
self.image_depth.is_animated_image(
SimpleNamespace(source="FILE")
)
)
def test_video_total_frames_reads_movie_duration(self):
image = SimpleNamespace(source="MOVIE", frame_duration=97)
self.assertEqual(self.image_depth.video_total_frames(image), 97)
def test_video_total_frames_counts_sequence_files(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "foreground-0001.png").touch()
(root / "foreground-0002.png").touch()
image = SimpleNamespace(
source="SEQUENCE",
filepath=str(root / "foreground-0001.png"),
)
self.assertEqual(self.image_depth.video_total_frames(image), 2)
def test_video_total_frames_returns_one_for_still_image(self):
image = SimpleNamespace(source="FILE", frame_duration=1)
self.assertEqual(self.image_depth.video_total_frames(image), 1)
def test_default_extract_frames_uses_empty_image_frames(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(frame_duration=24),
)
self.assertEqual(self.image_depth.default_extract_frames(source_object), 24)
def test_default_extract_frames_caps_at_video_total(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(frame_duration=200),
)
self.assertEqual(self.image_depth.default_extract_frames(source_object), 97)
def test_default_extract_frames_falls_back_to_video_total(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
)
self.assertEqual(self.image_depth.default_extract_frames(source_object), 97)
def test_extract_frame_range_honors_offset_within_clip(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=4,
frame_offset=70,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(71, 24),
)
def test_extract_frame_range_clamps_large_offset_to_clip_tail(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=90,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(74, 24),
)
def test_extract_frame_range_clamps_negative_offset_to_start(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=-5,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(1, 24),
)
def test_extract_frame_range_returns_full_clip_for_still(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="FILE", frame_duration=1),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=50,
frame_duration=24,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(1, 1),
)
def test_extract_frame_range_honors_single_frame_setting(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=70,
frame_duration=1,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(71, 1),
)
def test_configure_image_user_preserves_cyclic_setting(self):
image_user = SimpleNamespace(
frame_start=0,
frame_offset=5,
use_auto_refresh=False,
use_cyclic=False,
frame_duration=1,
)
self.image_depth.configure_image_user(
image_user,
97,
frame_start=4,
frame_offset=70,
)
self.assertFalse(image_user.use_cyclic)
self.assertTrue(image_user.use_auto_refresh)
self.assertEqual(image_user.frame_start, 4)
self.assertEqual(image_user.frame_offset, 70)
self.assertEqual(image_user.frame_duration, 97)
self.image_depth.configure_image_user(
image_user,
use_cyclic=True,
)
self.assertTrue(image_user.use_cyclic)
def test_image_user_settings_reads_empty_values(self):
source_object = SimpleNamespace(
image_user=SimpleNamespace(
frame_start=4,
frame_offset=70,
frame_duration=12,
)
)
self.assertEqual(
self.image_depth.image_user_settings(source_object),
(4, 70, 12, False),
)
source_object.image_user.use_cyclic = True
self.assertEqual(
self.image_depth.image_user_settings(source_object),
(4, 70, 12, True),
)
def test_image_user_settings_defaults_without_user(self):
self.assertEqual(
self.image_depth.image_user_settings(
SimpleNamespace(image_user=None)
),
(1, 0, 0, False),
)
def test_image_base_name_strips_video_extension(self):
image = SimpleNamespace(name="微信视频2026-07-23_192756_414.mp4")
self.assertEqual(
self.image_depth.image_base_name(image),
"微信视频2026-07-23_192756_414",
)
def test_removed_background_name_strips_video_extension(self):
source_object = SimpleNamespace(data=SimpleNamespace(name="clip.mp4"))
self.assertEqual(
self.image_depth.removed_background_data_name(source_object),
"clip.foreground",
)
def test_force_sequence_source_keeps_loaded_sequence(self):
image = SimpleNamespace(source="SEQUENCE")
result = self.image_depth.force_sequence_source(image)
self.assertIs(result, image)
self.assertEqual(image.source, "SEQUENCE")
def test_force_sequence_source_switches_single_image(self):
image = SimpleNamespace(source="FILE")
result = self.image_depth.force_sequence_source(image)
self.assertIs(result, image)
self.assertEqual(image.source, "SEQUENCE")
def test_force_sequence_source_reports_unloadable_sequence(self):
class BlockedImage:
@property
def source(self):
return "FILE"
@source.setter
def source(self, value):
raise RuntimeError("blocked")
with self.assertRaisesRegex(RuntimeError, "cannot load"):
self.image_depth.force_sequence_source(BlockedImage())
def test_collect_sequence_frames_returns_numbered_files_in_order(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "frame-0001.jpg").touch()
(root / "frame-0002.jpg").touch()
(root / "frame-0003.jpg").touch()
(root / "notes.txt").touch()
image = SimpleNamespace(
source="SEQUENCE",
filepath=str(root / "frame-0002.jpg"),
)
frames = self.image_depth.collect_sequence_frames(image)
self.assertEqual(
[path.name for path in frames],
["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"],
)
def test_sorted_frame_paths_uses_numeric_frame_order(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "frame-2.jpg").touch()
(root / "frame-10.jpg").touch()
(root / "frame-1.jpg").touch()
frames = self.image_depth.sorted_frame_paths(root)
self.assertEqual(
[path.name for path in frames],
["frame-1.jpg", "frame-2.jpg", "frame-10.jpg"],
)
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_use_stored_unmasked_reference_depth(self):
source = SimpleNamespace(
data=SimpleNamespace(size=(200, 100)),
empty_image_offset=(-0.5, -0.5),
)
depth = np.full((100, 200), 9.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,
reference_depth=2.0,
)
self.assertEqual(bounds, (-2.0, 2.0, -1.0, 1.0))
self.assertEqual(reference_depth, 2.0)
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), (0.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]), 0.0)
self.assertEqual(float(mapped[1, 1]), 0.0)
def test_depth_values_are_scaled_before_creating_depth_image(self):
depth = np.array(
((0.5, -1.0), (0.0, 2.0)),
dtype=np.float32,
)
scaled = self.image_depth.scale_depth_values(depth, 4.1)
np.testing.assert_allclose(
scaled,
((2.05, -4.1), (0.0, 8.2)),
)
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, 0.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_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_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_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",
)
def test_copy_depth_image_renames_still_source(self):
class FakeImage:
def __init__(self, name):
self.name = name
def copy(self):
return FakeImage("img.png.001")
source_object = SimpleNamespace(data=FakeImage("img.png"))
depth_image = self.image_depth.copy_depth_image(source_object)
self.assertEqual(depth_image.name, "img.depth")
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.foreground",
)
source_object.data.name = "img.png.001"
self.assertEqual(
self.image_depth.removed_background_data_name(source_object),
"img.001.foreground",
)
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_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()