Files
AnyImage/tests/test_image_depth.py
nan 13ef705226 Add operators for model downloading and surface generation
- Implemented `DownloadModel` and `DownloadRequiredModels` operators for downloading models.
- Created `ClearModels` operator to clear model resources.
- Added `ConvertToSurface` operator for converting image empties into textured surfaces.
- Introduced `GenerateDa3Surface` and `GenerateMoge2Surface` operators for generating surfaces using DA3 and MoGe-2 models respectively.
- Developed utility functions for handling image mesh data, assigning attributes, and managing geometry nodes.
- Established runtime configuration and post-installation model downloading in `runtime.py`.
2026-08-20 01:56:04 +08:00

761 lines
24 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(Menu=object, Operator=object)
fake_bpy.path = SimpleNamespace(abspath=lambda path: path)
fake_bpy.context = SimpleNamespace(scene=None)
module_paths = (
(
"snapform.operators.image",
PROJECT_ROOT / "src" / "snapform" / "operators" / "image.py",
),
(
"snapform.operators.material",
PROJECT_ROOT / "src" / "snapform" / "operators" / "material.py",
),
(
"snapform.operators.depth_plane.data",
PROJECT_ROOT / "src" / "snapform" / "operators" / "depth_plane" / "data.py",
),
(
"snapform.operators.depth_plane.build",
PROJECT_ROOT / "src" / "snapform" / "operators" / "depth_plane" / "build.py",
),
)
packages = {
name: ModuleType(name)
for name in (
"snapform",
"snapform.operators",
"snapform.operators.depth_plane",
)
}
for package in packages.values():
package.__path__ = []
loaded = []
with patch.dict(sys.modules, {"bpy": fake_bpy, **packages}):
for module_name, module_path in module_paths:
specification = importlib.util.spec_from_file_location(
module_name,
module_path,
)
module = importlib.util.module_from_spec(specification)
sys.modules[module_name] = module
specification.loader.exec_module(module)
loaded.append(module)
namespace = SimpleNamespace()
for module in loaded:
for name, value in vars(module).items():
if not name.startswith("__"):
setattr(namespace, name, value)
return namespace
class ImageDepthTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.image_depth = load_image_depth()
def test_material_asset_path_survives_package_split(self):
self.assertEqual(
self.image_depth.node_asset_path(),
PROJECT_ROOT / "src" / "snapform" / "assets" / "O_SnapForm.blend",
)
def test_material_color_space_matches_scene_view_transform(self):
mappings = {
"ACES 1.3": "ACES 1.3 sRGB",
"ACES 2.0": "ACES 2.0 sRGB",
"AgX": "AgX Base sRGB",
"Filmic": "Filmic sRGB",
"Khronos PBR Neutral": "Khronos PBR Neutral sRGB",
"Standard": "sRGB",
}
for view_transform, expected in mappings.items():
with self.subTest(view_transform=view_transform):
scene = SimpleNamespace(
view_settings=SimpleNamespace(
view_transform=view_transform,
)
)
self.assertEqual(
self.image_depth.material_color_space_name(scene),
expected,
)
def test_material_color_space_falls_back_when_unsupported(self):
class ColorSpaceSettings:
def __init__(self):
self._name = "sRGB"
@property
def name(self):
return self._name
@name.setter
def name(self, value):
if value != "sRGB":
raise TypeError("unsupported color space")
self._name = value
image = SimpleNamespace(colorspace_settings=ColorSpaceSettings())
scene = SimpleNamespace(
view_settings=SimpleNamespace(view_transform="ACES 2.0")
)
result = self.image_depth.set_material_color_space(image, scene)
self.assertEqual(result, "sRGB")
self.assertEqual(image.colorspace_settings.name, "sRGB")
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_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_color_name_strips_video_extension(self):
source_object = SimpleNamespace(data=SimpleNamespace(name="clip.mp4"))
self.assertEqual(
self.image_depth.color_data_name(source_object),
"clip.color",
)
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_raw_da3_depth_is_converted_only_for_displacement(self):
raw_depth = np.array(
((0.25, 1.0), (1.5, 2.0)),
dtype=np.float32,
)
displacement = self.image_depth.depth_displacement_values(raw_depth)
np.testing.assert_array_equal(
displacement,
np.array(((0.75, 0.0), (-0.5, -1.0)), dtype=np.float32),
)
np.testing.assert_array_equal(
raw_depth,
np.array(((0.25, 1.0), (1.5, 2.0)), dtype=np.float32),
)
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_long_edge_and_aligns_to_patch_size(self):
image = SimpleNamespace(size=(1080, 1920))
self.assertEqual(
self.image_depth.scaled_process_resolution(image, 25.0),
476,
)
self.assertEqual(
self.image_depth.scaled_process_resolution(image, 50.0),
966,
)
def test_process_resolution_clamps_to_supported_range(self):
large_image = SimpleNamespace(size=(4000, 3000))
small_image = SimpleNamespace(size=(200, 100))
self.assertEqual(
self.image_depth.scaled_process_resolution(large_image, 100.0),
1008,
)
self.assertEqual(
self.image_depth.scaled_process_resolution(small_image, 1.0),
252,
)
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_moge_normal_name_uses_source_image_name(self):
source_object = SimpleNamespace(
name="obj_img",
data=SimpleNamespace(name="img.png"),
)
self.assertEqual(
self.image_depth.normal_data_name(source_object),
"img.normal",
)
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_color_name_uses_source_image_base_name(self):
source_object = SimpleNamespace(
data=SimpleNamespace(name="img.png"),
)
self.assertEqual(
self.image_depth.color_data_name(source_object),
"img.color",
)
source_object.data.name = "img.png.001"
self.assertEqual(
self.image_depth.color_data_name(source_object),
"img.001.color",
)
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()