import importlib.util import os import sys import tempfile import unittest from contextlib import nullcontext 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_server_module(name): addon_directory = PROJECT_ROOT / "src" / "snapform" sys.path.insert(0, str(addon_directory)) try: return __import__(f"server.{name}", fromlist=[name]) finally: sys.path.remove(str(addon_directory)) def _make_fake_av(frame_count): """Build a fake av module yielding frame_count integer-valued frames.""" class FakeImage: def __init__(self, value): self.value = value def save(self, path, format=None): Path(path).write_text(str(self.value), encoding="utf-8") class FakeFrame: def __init__(self, value): self.value = value def to_image(self): return FakeImage(self.value) class FakeContainer: def __init__(self): self.position = 0 self.closed = False self.streams = SimpleNamespace(video=[object()]) def decode(self, _stream): while self.position < frame_count: frame = FakeFrame(self.position) self.position += 1 yield frame def close(self): self.closed = True class FakeAv: def __init__(self): self.containers = [] def open(self, _path): container = FakeContainer() self.containers.append(container) return container return FakeAv() class FakeModels: def __init__(self, directory="model"): self.model_directory = Path(directory) def directory(self, _model): return self.model_directory def get_ben2(self, _model_directory, _device): from server.models.onnx_ben2 import create_session return create_session(self.model_directory, "cpu"), 0.0 def get_da3(self, _model_directory, _device): return object() class FakeJobContext: def __init__(self, directory, job_type, models=None): self.directory = Path(directory) self.job_type = job_type self.models = models or FakeModels() self.status = {} def resource(self, name): if name != "model_manager": raise KeyError(name) return self.models def progress(self, progress, message, **details): self.status = { "progress": progress, "message": message, **details, } def check_cancelled(self): return None class WorkerTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.media = load_server_module("media.input") cls.geometry = load_server_module("geometry") cls.surface = load_server_module("geometry.surface") cls.background = load_server_module("jobs.background") cls.ben2 = load_server_module("models.ben2") cls.da3_model = load_server_module("models.da3") cls.moge2_model = load_server_module("models.moge2") cls.surface_outputs = load_server_module("outputs.surface") cls.debug_da3_outputs = load_server_module("outputs.debug_da3") cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2") cls.common_job = load_server_module("jobs.common") cls.depth_job = load_server_module("jobs.depth") cls.surface_job = load_server_module("jobs.surface") cls.debug_da3_job = load_server_module("jobs.debug_da3") cls.debug_moge2_job = load_server_module("jobs.debug_moge2") cls.debug_ben2_job = load_server_module("jobs.debug_ben2") cls.frames_job = load_server_module("jobs.frames") cls.foreground_job = load_server_module("jobs.foreground") cls.download = load_server_module("model_download") def test_inference_progress_omits_single_frame_count(self): self.assertEqual( self.common_job.frame_progress_message( "Estimating MoGe-2 geometry", 0, 1 ), "Estimating MoGe-2 geometry", ) self.assertEqual( self.common_job.frame_progress_message( "Estimating DA3 geometry", 0, 1 ), "Estimating DA3 geometry", ) self.assertEqual( self.common_job.frame_progress_message( "Estimating DA3 geometry", 1, 3 ), "Estimating DA3 geometry (2/3)", ) self.assertEqual( self.common_job.frame_progress_message( "Estimating MoGe-2 depth", 0, 1 ), "Estimating MoGe-2 depth", ) self.assertEqual( self.common_job.frame_progress_message( "Estimating DA3 depth", 0, 1 ), "Estimating DA3 depth", ) def test_surface_has_one_model_independent_builder(self): self.assertTrue(callable(self.surface.build_surface)) self.assertTrue(callable(self.surface_outputs.write_surface_archive)) self.assertFalse(hasattr(self.surface_outputs, "write_surface_texture")) self.assertFalse(hasattr(self.surface, "export_da3_surface")) self.assertFalse(hasattr(self.surface, "export_moge2_surface")) def test_model_adapters_only_normalize_surface_input(self): self.assertTrue(callable(self.da3_model.surface_frame)) self.assertTrue(callable(self.moge2_model.surface_frame)) self.assertFalse(hasattr(self.da3_model, "generate_surface")) self.assertFalse(hasattr(self.moge2_model, "generate_surface")) def test_collect_images_is_sorted_and_filters_suffixes(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "b.PNG").touch() (root / "a.jpg").touch() (root / "notes.txt").touch() images = self.media.collect_images(root) self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"]) def test_ai_tasks_support_cooperative_cancellation(self): with self.assertRaises(self.media.TaskCancelled): self.media.check_cancelled(lambda: True) self.media.check_cancelled(lambda: False) def test_prepares_a_single_image_file(self): with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "input.png" image_path.touch() images, temporary_directory = self.media.prepare_input( image_path, 16, directory, ) self.assertEqual(images, [image_path]) self.assertIsNone(temporary_directory) def test_prepares_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "frame-0001.jpg").touch() (root / "frame-0002.png").touch() (root / "notes.txt").touch() images, temporary_directory = self.media.prepare_input( root, 16, directory, ) self.assertEqual( [path.name for path in images], ["frame-0001.jpg", "frame-0002.png"], ) self.assertIsNone(temporary_directory) def test_prepares_first_frames_of_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() images, temporary_directory = self.media.prepare_input( root, 3, directory, ) self.assertEqual( [path.name for path in images], ["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"], ) self.assertIsNone(temporary_directory) def test_prepares_first_frames_after_offset_in_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() images, temporary_directory = self.media.prepare_input( root, 3, directory, start_frame=2, ) self.assertEqual( [path.name for path in images], ["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"], ) self.assertIsNone(temporary_directory) def test_copy_color_images_numbers_a_still_image(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source = root / "source.png" source.write_bytes(b"frame") output = root / "result" output.mkdir() self.media.copy_color_images([source], output) self.assertEqual( (output / "color.0001.png").read_bytes(), b"frame", ) def test_copy_frames_skips_a_still_image_by_default(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source = root / "source.png" source.touch() output = root / "result" output.mkdir() self.media.copy_frames([source], output) self.assertEqual(list(output.iterdir()), []) def test_extract_video_frames_reads_frames_sequentially(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_av, ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) self.assertEqual(len(fake_av.containers), 1) self.assertEqual(fake_av.containers[0].position, 10) def test_extract_video_frames_starts_after_offset(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_av, start_frame=71, ) extracted = [ Path(path).read_text(encoding="utf-8") for path in image_paths ] self.assertEqual(extracted, [str(frame) for frame in range(70, 80)]) self.assertEqual(fake_av.containers[0].position, 80) def test_extract_video_frames_tolerates_short_video_tail(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, _make_fake_av(99), ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) def test_rejects_unsupported_input_file(self): with tempfile.TemporaryDirectory() as directory: input_path = Path(directory) / "input.txt" input_path.touch() with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): self.media.prepare_input(input_path, 16, directory) def test_download_model_uses_actual_request_without_preflight(self): calls = [] fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face.snapshot_download = lambda **options: calls.append( options ) args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=[], ) with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, {"huggingface_hub": fake_hugging_face}, ), patch.object( self.download, "create_download_progress_class", return_value=object, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( self.download, "_verify_model_complete", return_value=None, ): args.model_dir = str(Path(directory) / "models" / "DA3") self.download.download_model(args) self.assertEqual(len(calls), 1) def test_download_model_refuses_success_when_files_missing(self): fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face.snapshot_download = lambda **options: None args = SimpleNamespace( repository="PramaLLC/BEN2", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=[], ) with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, {"huggingface_hub": fake_hugging_face}, ), patch.object( self.download, "create_download_progress_class", return_value=object, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}): args.model_dir = str(Path(directory) / "model") with self.assertRaisesRegex( RuntimeError, "All download sources failed", ): self.download.download_model(args) def test_download_model_reports_all_failed_sources(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=[], r2_directory="DA3Small", r2_file=["config.json|100|abc"], ) with tempfile.TemporaryDirectory() as directory, patch.object( self.download, "download_from_hf", side_effect=OSError("direct down"), ), patch.object( self.download, "download_from_r2", side_effect=OSError("mirror down"), ): args.model_dir = str(Path(directory) / "model") with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"): self.download.download_model(args) def test_download_model_prefers_direct_hf_then_r2(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=["config.json"], r2_directory="DA3Small", r2_file=["config.json|100|abc"], ) order = [] def fake_download_from_hf(*_args, **_kwargs): order.append("hf") raise OSError("hf down") def fake_download_from_r2(*_args, **_kwargs): order.append("r2") with tempfile.TemporaryDirectory() as directory, patch.object( self.download, "download_from_hf", side_effect=fake_download_from_hf, ), patch.object( self.download, "download_from_r2", side_effect=fake_download_from_r2, ), patch.object( self.download, "_verify_model_complete", return_value=None, ): args.model_dir = str(Path(directory) / "models" / "DA3") self.download.download_model(args) self.assertEqual(order, ["hf", "r2"]) def test_download_from_r2_downloads_declared_files(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=["config.json"], r2_directory="DA3Small", r2_file=["config.json|100|abc"], ) downloaded = [] with tempfile.TemporaryDirectory() as directory, patch.object( self.download, "_download_file_from_url", side_effect=lambda *args, **kwargs: downloaded.append(args[1].name), ): args.model_dir = str(Path(directory) / "model") self.download.download_from_r2(args, args.model_dir) self.assertEqual(downloaded, ["config.json"]) def test_ben2_background_removal_writes_fast_transparent_png(self): inference_options = [] save_options = [] class FakeImage: size = (1, 1) def __enter__(self): return self def __exit__(self, *_args): return None def convert(self, mode): self.mode = mode return self def getchannel(self, _name): return self def putalpha(self, _alpha): return None 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("server.models.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() fake_ben2.infer = lambda _model, image: ( inference_options.append((image, {})) or FakeImage() ) fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, { "torch": fake_torch, "server.models.onnx_ben2": fake_ben2, "PIL": fake_pil, }, ): root = Path(directory) input_path = root / "input.png" input_path.touch() context = FakeJobContext(root / "result", "remove-background") self.background.run( context, {"input": str(input_path), "model": "BEN2", "device": "cpu"}, ) self.assertTrue((root / "result" / "color.0001.png").is_file()) self.assertEqual( inference_options[0][1], {}, ) self.assertEqual( save_options[0], {"format": "PNG", "compress_level": 1}, ) self.assertIn("Background removed in", context.status["message"]) def test_ben2_background_removal_processes_each_frame(self): inference_options = [] class FakeImage: size = (1, 1) def __enter__(self): return self def __exit__(self, *_args): return None def convert(self, mode): self.mode = mode return self def getchannel(self, _name): return self def putalpha(self, _alpha): return None def save(self, path, **options): Path(path).touch() class FakeModel: @classmethod def from_pretrained(cls, model_dir): return cls() def to(model_self, device): 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("server.models.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() fake_ben2.infer = lambda _model, image: ( inference_options.append((image, {})) or FakeImage() ) fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, { "torch": fake_torch, "server.models.onnx_ben2": fake_ben2, "PIL": fake_pil, }, ): root = Path(directory) input_directory = root / "frames" input_directory.mkdir() (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() context = FakeJobContext(root / "result", "remove-background") self.background.run( context, { "input": str(input_directory), "model": "BEN2", "device": "cpu", "video_frames": 8, }, ) result_directory = context.directory self.assertTrue( (result_directory / "color.0001.png").is_file() ) self.assertTrue( (result_directory / "color.0002.png").is_file() ) self.assertEqual(len(inference_options), 2) for _image, options in inference_options: self.assertEqual(options, {}) def test_extract_frames_copies_selected_frames(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_directory = root / "frames" input_directory.mkdir() for index in range(1, 6): (input_directory / f"frame-{index:04d}.jpg").touch() context = FakeJobContext(root / "result", "extract-frames") self.frames_job.run( context, { "input": str(input_directory), "video_frames": 3, "start_frame": 2, }, ) result_directory = context.directory self.assertEqual( sorted( path.name for path in result_directory.iterdir() if path.is_file() ), ["color.0001.jpg", "color.0002.jpg", "color.0003.jpg"], ) self.assertEqual(context.status["progress"], 1.0) self.assertIn("3 frame", context.status["message"]) def test_da3_debug_saves_only_direct_api_prediction(self): prediction = SimpleNamespace( depth=np.full((1, 2, 2), 2.0, dtype=np.float32), conf=np.ones((1, 2, 2), dtype=np.float32), extrinsics=np.eye(4, dtype=np.float32)[None, :3], intrinsics=np.eye(3, dtype=np.float32)[None], processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), is_metric=0, sky=None, scale_factor=None, ) fake_da3 = ModuleType("server.models.onnx_da3") calls = [] def fake_infer(_session, images, **options): calls.append((images, options)) return prediction fake_da3.infer = fake_infer fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" input_path.touch() with patch.dict( sys.modules, { "server.models.onnx_da3": fake_da3, "server.models.onnx_runtime": fake_runtime, }, ): context = FakeJobContext(root / "result", "debug-da3") self.debug_da3_job.run( context, { "input": str(input_path), "model": "DA3", "video_frames": 1, "device": "cpu", "process_res": 518, "process_res_method": "lower_bound_resize", "ref_view_strategy": "middle", }, ) with np.load(context.directory / "prediction.npz") as result: self.assertEqual( set(result.files), { "depth", "conf", "extrinsics", "intrinsics", "processed_images", "is_metric", }, ) np.testing.assert_array_equal(result["conf"], prediction.conf) self.assertFalse(any(context.directory.glob("*.glb"))) self.assertFalse(any(context.directory.glob("*.png"))) self.assertEqual(calls[0][0], [str(input_path)]) self.assertEqual(calls[0][1]["process_res"], 518) def test_moge2_debug_saves_each_direct_api_dictionary(self): prediction = { "points": np.ones((2, 2, 3), dtype=np.float32), "depth": np.ones((2, 2), dtype=np.float32), "normal": np.ones((2, 2, 3), dtype=np.float32), "mask": np.ones((2, 2), dtype=bool), "intrinsics": np.eye(3, dtype=np.float32), } with tempfile.TemporaryDirectory() as directory: root = Path(directory) frames = root / "frames" frames.mkdir() (frames / "a.png").touch() (frames / "b.png").touch() context = FakeJobContext(root / "result", "debug-moge2") with patch.object( self.debug_moge2_job, "infer_one", return_value=prediction, ) as infer_one: self.debug_moge2_job.run( context, { "input": str(frames), "model": "MOGE2_VITS_NORMAL", "video_frames": 2, "device": "cpu", "resolution_level": 5, }, ) self.assertEqual(infer_one.call_count, 2) archives = sorted(context.directory.glob("prediction.*.npz")) self.assertEqual(len(archives), 2) with np.load(archives[0]) as result: self.assertEqual(set(result.files), set(prediction)) self.assertFalse(any(context.directory.glob("*.glb"))) def test_ben2_debug_saves_direct_api_foreground(self): class Foreground: def save(self, path): Path(path).write_bytes(b"foreground") with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" input_path.touch() context = FakeJobContext(root / "result", "debug-ben2") with patch.object( self.debug_ben2_job, "infer_foreground", return_value=(Foreground(), 1.0, 2.0), ) as infer_foreground: self.debug_ben2_job.run( context, { "input": str(input_path), "model": "BEN2_BASE", "video_frames": 1, "device": "cpu", }, ) infer_foreground.assert_called_once() self.assertEqual( (context.directory / "foreground.png").read_bytes(), b"foreground", ) def test_da3_removes_background_after_depth_using_original_input(self): order = [] inference_inputs = [] fake_da3 = ModuleType("server.models.onnx_da3") def fake_infer(_session, images, **_options): order.append("depth") inference_inputs.append(images) return SimpleNamespace( depth=np.full((1, 2, 2), 2.0, dtype=np.float32), conf=np.ones((1, 2, 2), dtype=np.float32), extrinsics=np.eye(4, dtype=np.float32)[None, :3], intrinsics=np.eye(3, dtype=np.float32)[None], processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), ) fake_da3.infer = fake_infer fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" class FakeForeground: def convert(self, _mode): return self def save(self, path, **_options): Path(path).touch() def fake_foreground(image, *_args, **_kwargs): order.append("background") self.assertEqual(image, input_path) return FakeForeground(), 1.0, 2.0 with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" input_path.touch() with patch.dict( sys.modules, { "server.models.onnx_da3": fake_da3, "server.models.onnx_runtime": fake_runtime, }, ), patch.object( self.foreground_job, "infer_foreground", side_effect=fake_foreground, ): context = FakeJobContext( root / "result", "generate-depth", ) self.depth_job.run( context, { "input": str(input_path), "model": "DA3", "remove_background_model": "BEN2", "video_frames": 1, "export_format": "mini_npz", "device": "cpu", "process_res": 518, "process_res_method": "upper_bound_resize", "ref_view_strategy": "first", "inference_mode": "per_frame", }, ) self.assertEqual(order, ["depth", "background"]) self.assertEqual(inference_inputs, [[str(input_path)]]) def test_reconstruct_copies_frames_for_multi_frame_input(self): inference_inputs = [] fake_da3 = ModuleType("server.models.onnx_da3") fake_da3.create_session = lambda _model_dir, _device: object() def fake_infer(_session, images, **_options): inference_inputs.append(images) return SimpleNamespace( depth=np.full((1, 2, 2), 2.0, dtype=np.float32), conf=np.ones((1, 2, 2), dtype=np.float32), extrinsics=np.eye(4, dtype=np.float32)[None, :3], intrinsics=np.eye(3, dtype=np.float32)[None], processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), ) fake_da3.infer = fake_infer fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_directory = root / "frames" input_directory.mkdir() (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() with patch.dict( sys.modules, { "server.models.onnx_da3": fake_da3, "server.models.onnx_runtime": fake_runtime, }, ): context = FakeJobContext( root / "result", "generate-depth", ) self.depth_job.run( context, { "input": str(input_directory), "model": "DA3", "video_frames": 16, "export_format": "mini_npz", "device": "cpu", "process_res": 518, "process_res_method": "upper_bound_resize", "ref_view_strategy": "first", "conf_thresh_percentile": 40.0, "num_max_points": 1000, "inference_mode": "per_frame", }, ) result_directory = context.directory self.assertTrue( (result_directory / "color.0001.jpg").is_file() ) self.assertTrue( (result_directory / "color.0002.jpg").is_file() ) self.assertEqual(len(inference_inputs), 2) self.assertTrue(all(len(paths) == 1 for paths in inference_inputs)) with np.load(result_directory / "results.npz") as result: self.assertEqual(result["depth"].shape, (2, 2, 2)) self.assertTrue( (result_directory / "results.npz").is_file() ) self.assertFalse( (result_directory / "exports").exists() ) self.assertEqual(context.status["progress"], 1.0) def test_mesh_confidence_uses_percentile(self): depth = np.ones((1, 2, 2), dtype=np.float32) confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]]) threshold = self.geometry.confidence_percentile_threshold( depth, confidence, 50.0, ) self.assertEqual(threshold, 2.5) def test_zero_mesh_confidence_disables_confidence_filter(self): depth = np.ones((2, 2), dtype=np.float32) confidence = np.full((2, 2), np.nan, dtype=np.float32) x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0]) points = np.stack((x, y, depth), axis=-1) threshold = self.geometry.confidence_percentile_threshold( depth, confidence, 0.0, ) faces = self.geometry.mesh_triangle_faces( points, depth, confidence, threshold, filter_edges=False, max_edge_length=0.0, ) self.assertIsNone(threshold) self.assertEqual(len(faces), 2) def test_mesh_filters_triangles_independently(self): depth = np.ones((2, 2), dtype=np.float32) confidence = np.ones((2, 2), dtype=np.float32) confidence[1, 0] = 0.0 x, y = np.meshgrid( np.arange(2, dtype=np.float32), np.arange(2, dtype=np.float32), ) points = np.stack((x, y, depth), axis=-1) faces = self.geometry.mesh_triangle_faces( points, depth, confidence, confidence_threshold=0.5, filter_edges=False, max_edge_length=0.0, ) np.testing.assert_array_equal(faces, [[0, 3, 1]]) def test_mesh_background_mask_combines_with_confidence(self): depth = np.ones((2, 3), dtype=np.float32) confidence = np.ones((2, 3), dtype=np.float32) x, y = np.meshgrid( np.arange(3, dtype=np.float32), np.arange(2, dtype=np.float32), ) points = np.stack((x, y, depth), axis=-1) foreground = np.ones((2, 3), dtype=bool) foreground[:, 0] = False faces = self.geometry.mesh_triangle_faces( points, depth, confidence, confidence_threshold=None, filter_edges=False, max_edge_length=0.0, vertex_mask=foreground, ) np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]]) def test_depth_edge_filter_ignores_smooth_slope(self): smooth = np.tile( np.linspace(2.0, 2.04, 5, dtype=np.float32), (3, 1), ) discontinuous = smooth.copy() discontinuous[:, 3:] += 1.0 self.assertFalse(self.geometry.depth_edge_mask(smooth).any()) self.assertTrue(self.geometry.depth_edge_mask(discontinuous)[:, 2:4].all()) def test_higher_edge_culling_allows_less_stretch(self): limits = [ self.geometry.edge_culling_stretch_limit(value) for value in (0.0, 50.0, 100.0) ] self.assertEqual(limits, [10000.0, 2.0, 1.0]) def test_stretched_edge_culling_uses_camera_space_edge_ratio(self): flat_points = np.asarray( [ [[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]], [[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]], ], dtype=np.float32, ) stretched_points = flat_points.copy() stretched_points[0, 1] = (2.0, 0.0, 2.0) triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32) flat_mask = self.geometry.triangle_edge_stretch_mask( flat_points, triangles, 1.0, ) loose_mask = self.geometry.triangle_edge_stretch_mask( stretched_points, triangles, 2.0, ) strict_mask = self.geometry.triangle_edge_stretch_mask( stretched_points, triangles, 1.0, ) self.assertTrue(flat_mask.all()) self.assertTrue(loose_mask.all()) self.assertFalse(strict_mask.all()) def test_mesh_faces_apply_confidence_threshold_and_face_direction(self): points = np.array( [ [[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [2.0, 0.0, 1.0]], [[0.0, 1.0, 1.0], [1.0, 1.0, 1.0], [2.0, 1.0, 1.0]], ], dtype=np.float32, ) depth = np.ones((2, 3), dtype=np.float32) confidence = np.ones((2, 3), dtype=np.float32) confidence[:, 0] = 0.0 faces = self.geometry.mesh_triangle_faces( points, depth, confidence, confidence_threshold=0.5, filter_edges=False, max_edge_length=0.0, ) np.testing.assert_array_equal(faces, [[1, 4, 5], [1, 5, 2]]) vertices = points.reshape(-1, 3)[faces[0]] normal = np.cross( vertices[1] - vertices[0], vertices[2] - vertices[0], ) self.assertLess(normal[2], 0.0) def test_mesh_percentile_reduces_face_count(self): depth = np.ones((3, 3), dtype=np.float32) confidence = np.arange(1, 10, dtype=np.float32).reshape(3, 3) x, y = np.meshgrid( np.arange(3, dtype=np.float32), np.arange(3, dtype=np.float32), ) points = np.stack((x, y, depth), axis=-1) face_counts = [] for percentile in (0.0, 50.0, 100.0): threshold = self.geometry.confidence_percentile_threshold( depth, confidence, percentile, ) faces = self.geometry.mesh_triangle_faces( points, depth, confidence, confidence_threshold=threshold, filter_edges=False, max_edge_length=0.0, ) face_counts.append(len(faces)) self.assertEqual(face_counts, [8, 2, 0]) def test_mesh_compaction_removes_unreferenced_vertices(self): vertices = np.arange(18, dtype=np.float32).reshape(6, 3) faces = np.asarray([[1, 3, 5]], dtype=np.int32) compact_vertices, compact_faces = self.geometry.compact_mesh( vertices, faces, ) np.testing.assert_array_equal( compact_vertices, vertices[[1, 3, 5]], ) np.testing.assert_array_equal(compact_faces, [[0, 1, 2]]) def test_surface_background_requires_every_vertex_to_be_foreground(self): faces = np.asarray( ((0, 1, 4), (0, 4, 3), (1, 2, 5), (1, 5, 4)), dtype=np.int32, ) foreground = np.asarray( ((True, True, False), (False, True, False)), dtype=bool, ) filtered = self.geometry.remove_background_faces(faces, foreground) np.testing.assert_array_equal(filtered, faces[[0]]) def test_surface_background_mask_keeps_any_nonzero_alpha(self): from PIL import Image with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) for alpha, expected in ((0, False), (1, True), (127, True)): path = root / f"alpha-{alpha}.png" Image.new("RGBA", (2, 2), (10, 20, 30, alpha)).save(path) mask = self.geometry.foreground_mesh_mask(path, (2, 2)) self.assertEqual(mask.dtype, np.bool_) self.assertTrue(np.all(mask == expected)) def test_image_mesh_data_keeps_camera_geometry_and_uv(self): depth = np.full((3, 3), 2.0, dtype=np.float32) confidence = np.ones((3, 3), dtype=np.float32) intrinsics = np.asarray( [[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32, ) frame = self.surface.GeometryFrame( depth=depth, points=self.geometry.unproject_camera_depth(depth, intrinsics), intrinsics=intrinsics, valid_mask=np.ones(depth.shape, dtype=bool), confidence=confidence, ) mesh = self.surface.build_surface(frame) with tempfile.TemporaryDirectory() as directory: result_path = self.surface_outputs.write_surface_archive( mesh, Path(directory) / "mesh.npz", ) with np.load(result_path) as result: self.assertEqual(result["vertices"].shape, (9, 3)) self.assertEqual(result["faces"].shape, (8, 3)) self.assertEqual(result["uv"].shape, (9, 2)) self.assertEqual(float(result["reference_depth"]), 2.0) self.assertNotIn("postprocess_method", result.files) self.assertEqual(result["point_depth"].shape, (9,)) self.assertEqual( result["face_confidence"].shape, (8,), ) self.assertEqual( set(result.files), { "vertices", "faces", "uv", "intrinsics", "image_size", "reference_depth", "point_depth", "face_confidence", "face_depth_discontinuity", }, ) np.testing.assert_allclose(result["uv"][0], [0.0, 1.0]) np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0]) self.assertNotIn("normal_map", result.files) def test_moge_surface_uses_the_same_archive_without_normal_data(self): depth = np.full((3, 3), 2.0, dtype=np.float32) intrinsics = np.asarray( [[2.0, 0.0, 1.0], [0.0, 2.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32, ) frame = self.surface.GeometryFrame( depth=depth, points=self.geometry.unproject_camera_depth(depth, intrinsics), intrinsics=intrinsics, valid_mask=np.ones(depth.shape, dtype=bool), model_mask=np.asarray( ( (False, False, False), (True, True, True), (True, True, True), ), dtype=bool, ), ) mesh = self.surface.build_surface(frame) with tempfile.TemporaryDirectory() as directory: path = self.surface_outputs.write_surface_archive( mesh, Path(directory) / "mesh.npz", ) with np.load(path) as result: self.assertIn("face_mask", result.files) self.assertNotIn("face_confidence", result.files) self.assertNotIn("normal", result.files) self.assertNotIn("normal_map", result.files) def test_moge_surface_normal_is_exported_as_a_texture(self): from PIL import Image normal = np.asarray( [[[0.0, 0.0, -1.0]]], dtype=np.float32, ) with tempfile.TemporaryDirectory() as directory: path = self.surface_outputs.write_normal_texture( normal, Path(directory) / "normal.png", ) with Image.open(path) as image: self.assertEqual(image.mode, "RGB") self.assertEqual(image.size, (1, 1)) red, green, blue = image.getpixel((0, 0)) self.assertIn(red, {127, 128}) self.assertIn(green, {127, 128}) self.assertEqual(blue, 255) def test_da3_debug_writer_keeps_api_values_without_derived_attributes(self): prediction = SimpleNamespace( depth=np.ones((1, 2, 3), dtype=np.float32), conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3), intrinsics=np.eye(3, dtype=np.float32)[None], extrinsics=np.eye(4, dtype=np.float32)[None], ) with tempfile.TemporaryDirectory() as directory: output_path = self.debug_da3_outputs.write_prediction( prediction, Path(directory) / "prediction.npz", ) with np.load(output_path) as attributes: self.assertEqual(attributes["depth"].shape, (1, 2, 3)) self.assertEqual(attributes["conf"].shape, (1, 2, 3)) self.assertNotIn("confidence_percentile", attributes.files) if __name__ == "__main__": unittest.main()