import importlib.util import json 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_worker(): specification = importlib.util.spec_from_file_location( "snapform_worker", PROJECT_ROOT / "src" / "snapform" / "worker.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) return module class WorkerTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.worker = load_worker() def test_foreground_preserves_source_rgb_and_ben2_alpha(self): alpha = SimpleNamespace(size=(2, 1)) source_rgba = SimpleNamespace(size=(2, 1), applied_alpha=None) source_rgba.putalpha = lambda value: setattr( source_rgba, "applied_alpha", value, ) source = SimpleNamespace(convert=lambda _mode: source_rgba) foreground_rgba = SimpleNamespace( getchannel=lambda _name: alpha, ) foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba) result = self.worker.preserve_source_rgb(source, foreground) self.assertIs(result, source_rgba) self.assertIs(result.applied_alpha, alpha) 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.worker.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.worker.TaskCancelled): self.worker.check_cancelled(lambda: True) self.worker.check_cancelled(lambda: False) def test_select_device_supports_mps_and_auto_fallback(self): mps_torch = SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: False), backends=SimpleNamespace( mps=SimpleNamespace(is_available=lambda: True), ), ) mps_torch.device = lambda device_type: SimpleNamespace( type=device_type ) self.assertEqual( self.worker.select_device("mps", mps_torch).type, "mps", ) self.assertEqual( self.worker.select_device("auto", mps_torch).type, "mps", ) self.assertEqual( self.worker.select_device("cpu", mps_torch).type, "cpu", ) cuda_torch = SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: True), ) cuda_torch.device = lambda device_type: SimpleNamespace( type=device_type ) self.assertEqual( self.worker.select_device("auto", cuda_torch).type, "cuda", ) cpu_only_torch = SimpleNamespace( cuda=SimpleNamespace(is_available=lambda: False), backends=SimpleNamespace( mps=SimpleNamespace(is_available=lambda: False), ), ) cpu_only_torch.device = lambda device_type: SimpleNamespace( type=device_type ) self.assertEqual( self.worker.select_device("auto", cpu_only_torch).type, "cpu", ) with self.assertRaisesRegex(RuntimeError, "MPS was requested"): self.worker.select_device("mps", cpu_only_torch) def test_reconstruct_parser_accepts_background_model_directory(self): 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" image_path.touch() images, temporary_directory = self.worker.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.worker.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.worker.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.worker.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_extract_video_frames_reads_frames_sequentially(self): class FakeCapture: def __init__(self): self.position = 0 def isOpened(self): return True def read(self): if self.position >= 100: return False, None frame = self.position self.position += 1 return True, frame def release(self): return None class FakeCv2: def __init__(self): self.captures = [] def VideoCapture(self, _path): capture = FakeCapture() self.captures.append(capture) return capture def imwrite(self, path, frame): Path(path).write_text(str(frame), encoding="utf-8") return True with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_cv2 = FakeCv2() image_paths = self.worker.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_cv2, ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) self.assertEqual(len(fake_cv2.captures), 1) self.assertEqual(fake_cv2.captures[0].position, 10) def test_extract_video_frames_starts_after_offset(self): class FakeCapture: def __init__(self): self.position = 0 def isOpened(self): return True def read(self): if self.position >= 100: return False, None frame = self.position self.position += 1 return True, frame def release(self): return None class FakeCv2: def __init__(self): self.captures = [] def VideoCapture(self, _path): capture = FakeCapture() self.captures.append(capture) return capture def imwrite(self, path, frame): Path(path).write_text(str(frame), encoding="utf-8") return True with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_cv2 = FakeCv2() image_paths = self.worker.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_cv2, 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_cv2.captures[0].position, 80) def test_extract_video_frames_tolerates_short_video_tail(self): class ShortCapture: def __init__(self): self.position = 0 def isOpened(self): return True def read(self): if self.position >= 99: return False, None frame = self.position self.position += 1 return True, frame def release(self): return None class FakeCv2: def VideoCapture(self, _path): return ShortCapture() def imwrite(self, path, frame): Path(path).write_text(str(frame), encoding="utf-8") return True with tempfile.TemporaryDirectory() as directory: root = Path(directory) image_paths = self.worker.extract_video_frames( root / "clip.mkv", root / "out", 10, FakeCv2(), ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) def test_npz_is_complete_requires_readable_archive(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) valid_path = root / "valid.npz" np.savez_compressed( valid_path, depth=np.zeros((4, 4)), ) self.assertTrue(self.worker.npz_is_complete(valid_path)) empty_path = root / "empty.npz" empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18) self.assertFalse(self.worker.npz_is_complete(empty_path)) def test_wait_for_exported_paths_waits_for_complete_npz(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) npz_path = root / "exports" / "mini_npz" / "results.npz" npz_path.parent.mkdir(parents=True) npz_path.touch() with patch.object( self.worker, "npz_is_complete", side_effect=[False, True], ): result = self.worker.wait_for_exported_paths( root, "mini_npz", timeout=2.0, ) self.assertEqual(result, [npz_path]) def test_move_exported_file_retries_windows_file_lock(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source = root / "source.npz" destination = root / "destination.npz" source.touch() path_type = type(source) original_replace = path_type.replace calls = 0 def locked_once(path, target): nonlocal calls calls += 1 if calls == 1: error = PermissionError("file is in use") error.winerror = 32 raise error return original_replace(path, target) with patch.object(path_type, "replace", new=locked_once), patch.object( self.worker.time, "sleep", ) as sleep: result = self.worker.move_exported_file( source, destination, ) self.assertEqual(result, destination) self.assertTrue(destination.is_file()) sleep.assert_called_once_with(0.1) 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.worker.prepare_input(input_path, 16, directory) def test_moge_runtime_disables_unsupported_xformers_attention(self): with patch.dict(os.environ, {}, clear=True): self.worker.prepare_moge_runtime() self.assertEqual(os.environ["XFORMERS_DISABLED"], "1") 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", status="status.json", allow_pattern=[], ) 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, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( self.worker, "_verify_model_complete", return_value=None, ): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") self.worker.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", status="status.json", allow_pattern=[], ) 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, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") with self.assertRaisesRegex( RuntimeError, "All download sources failed", ): self.worker.download_model(args) def test_download_model_reports_all_failed_sources(self): fake_hugging_face = ModuleType("huggingface_hub") def fake_snapshot_download(**options): raise OSError("direct down") fake_hugging_face.snapshot_download = fake_snapshot_download args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", status="status.json", allow_pattern=[], ) 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, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( self.worker, "download_from_r2", side_effect=OSError("mirror down"), ): args.r2_directory = "DA3Small" args.r2_file = ["config.json|100|abc"] args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"): self.worker.download_model(args) def test_download_from_r2_downloads_declared_files(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", status="status.json", allow_pattern=["config.json"], r2_directory="DA3Small", r2_file=["config.json|100|abc"], ) downloaded = [] with tempfile.TemporaryDirectory() as directory, patch.object( self.worker, "_download_r2_file", side_effect=lambda *args, **kwargs: downloaded.append(args[1]), ), patch.object(self.worker, "write_status"): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") self.worker.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("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_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("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_directory = root / "frames" input_directory.mkdir() (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() args = SimpleNamespace( input=str(input_directory), model_dir="model", output_dir=str(root / "result"), status=str(root / "status.json"), device="cpu", video_frames=8, ) self.worker.remove_background(args) result_directory = Path(args.output_dir) self.assertTrue( (result_directory / "foreground.0001.png").is_file() ) self.assertTrue( (result_directory / "foreground.0002.png").is_file() ) self.assertTrue( (result_directory / "frame.0001.jpg").is_file() ) self.assertTrue( (result_directory / "frame.0002.jpg").is_file() ) self.assertEqual(len(inference_options), 2) for _image, options in inference_options: self.assertEqual(options, {"refine_foreground": False}) def test_remove_background_parser_accepts_video_frames(self): args = self.worker.parser().parse_args( [ "remove-background", "--input", "input.mp4", "--model-dir", "ben2", "--output-dir", "result", "--status", "status.json", "--device", "auto", "--video-frames", "24", ] ) self.assertEqual(args.command, "remove-background") self.assertEqual(args.video_frames, 24) def test_extract_frames_parser_accepts_frame_range(self): args = self.worker.parser().parse_args( [ "extract-frames", "--input", "frames", "--output-dir", "result", "--status", "status.json", "--video-frames", "24", "--start-frame", "70", ] ) self.assertEqual(args.command, "extract-frames") self.assertEqual(args.video_frames, 24) self.assertEqual(args.start_frame, 70) 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() args = SimpleNamespace( input=str(input_directory), output_dir=str(root / "result"), status=str(root / "status.json"), video_frames=3, start_frame=2, ) self.worker.extract_frames(args) result_directory = Path(args.output_dir) self.assertEqual( sorted( path.name for path in result_directory.iterdir() if path.is_file() ), ["frame.0001.jpg", "frame.0002.jpg", "frame.0003.jpg"], ) status = json.loads( Path(args.status).read_text(encoding="utf-8") ) self.assertEqual(status["progress"], 1.0) self.assertIn("3 frame", status["message"]) def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self): inference_options = [] export_calls = [] mesh_calls = [] class FakeModel: @classmethod def from_pretrained(cls, model_dir): if model_dir != "model": raise AssertionError(f"Unexpected model directory: {model_dir}") return cls() def to(self, device): if device.type != "cpu": raise AssertionError(f"Unexpected device: {device.type}") return self def eval(self): return None def inference(self, images, **options): inference_options.append((images, options)) return SimpleNamespace( depth=np.full((1, 2, 2), 2.0, dtype=np.float32), ) def fake_export(prediction, export_format, export_dir, **options): export_calls.append((prediction, export_format, export_dir, options)) output_directory = Path(export_dir) if export_format == "glb": (output_directory / "scene.glb").touch() elif export_format == "gs_ply": ply_directory = output_directory / "gs_ply" ply_directory.mkdir() (ply_directory / "0000.ply").touch() elif export_format == "depth_vis": depth_directory = output_directory / "depth_vis" depth_directory.mkdir() (depth_directory / "0000.jpg").touch() else: mini_directory = output_directory / "exports" / "mini_npz" mini_directory.mkdir(parents=True) np.savez_compressed( mini_directory / "results.npz", depth=np.zeros((14, 14)), ) fake_torch = ModuleType("torch") fake_torch.device = lambda device_type: SimpleNamespace( type=device_type, ) fake_torch.inference_mode = nullcontext fake_api = ModuleType("depth_anything_3.api") fake_api.DepthAnything3 = FakeModel fake_export_module = ModuleType("depth_anything_3.utils.export") fake_export_module.export = fake_export fake_utils = ModuleType("depth_anything_3.utils") fake_utils.__path__ = [] fake_package = ModuleType("depth_anything_3") fake_package.__path__ = [] def fake_mesh_export( prediction, images, output_directory, confidence_percentile, filter_edges, max_edge_length, ): mesh_calls.append( ( prediction, images, confidence_percentile, filter_edges, max_edge_length, ) ) output_path = Path(output_directory) / "mesh.glb" output_path.touch() return output_path with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" input_path.touch() with patch.dict( sys.modules, { "torch": fake_torch, "depth_anything_3": fake_package, "depth_anything_3.api": fake_api, "depth_anything_3.utils": fake_utils, "depth_anything_3.utils.export": fake_export_module, }, ), patch.object( self.worker, "export_textured_mesh_glb", side_effect=fake_mesh_export, ): debug_formats = ( "glb", "gs_ply", "depth_vis", "mini_npz", "mesh", ) for export_format in debug_formats: args = SimpleNamespace( input=str(input_path), video_frames=16, output_dir=str(root / export_format), export_format=export_format, status=str(root / f"{export_format}.json"), device="cpu", model_dir="model", process_res=518, process_res_method="lower_bound_resize", use_ray_pose=1, ref_view_strategy="middle", conf_thresh_percentile=35.0, num_max_points=12345, show_cameras=0, gs_views_interval=2, mesh_filter_edges=True, mesh_max_edge_length=0.2, ) self.worker.reconstruct(args) status = json.loads( Path(args.status).read_text(encoding="utf-8") ) self.assertEqual(status["progress"], 1.0) self.assertIn(export_format, status["message"]) result_directory = Path(args.output_dir) if export_format == "gs_ply": self.assertTrue((result_directory / "0000.ply").is_file()) self.assertFalse((result_directory / "gs_ply").exists()) if export_format == "depth_vis": self.assertTrue((result_directory / "0000.jpg").is_file()) self.assertFalse((result_directory / "depth_vis").exists()) self.assertEqual(len(inference_options), 5) self.assertEqual(len(export_calls), 4) self.assertEqual(len(mesh_calls), 1) for images, options in inference_options: self.assertEqual(images, [str(input_path)]) self.assertEqual(options["process_res"], 518) self.assertEqual( options["process_res_method"], "lower_bound_resize", ) self.assertTrue(options["use_ray_pose"]) self.assertEqual(options["ref_view_strategy"], "middle") for ( _prediction, export_format, _export_dir, options, ) in export_calls: np.testing.assert_array_equal( _prediction.depth, np.full((1, 2, 2), 2.0, dtype=np.float32), ) if export_format == "glb": self.assertEqual( options, { "glb": { "conf_thresh_percentile": 35.0, "num_max_points": 12345, "show_cameras": False, "export_depth_vis": False, } }, ) elif export_format == "gs_ply": self.assertEqual( options, {"gs_ply": {"gs_views_interval": 2}}, ) else: self.assertEqual(options, {}) self.assertEqual(mesh_calls[0][1], [input_path]) self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2)) infer_gs_values = [options["infer_gs"] for _, options in inference_options] self.assertEqual( infer_gs_values, [False, True, False, False, False], ) def test_reconstruct_copies_frames_for_multi_frame_input(self): class FakeModel: @classmethod def from_pretrained(cls, model_dir): return cls() def to(self, device): return self def eval(self): return None def inference(self, images, **options): return SimpleNamespace( depth=np.full((2, 2, 2), 2.0, dtype=np.float32), ) def fake_export(prediction, export_format, export_dir, **options): mini_directory = Path(export_dir) / "exports" / "mini_npz" mini_directory.mkdir(parents=True) np.savez_compressed( mini_directory / "results.npz", depth=np.zeros((2, 14, 14)), ) fake_torch = ModuleType("torch") fake_torch.device = lambda device_type: SimpleNamespace( type=device_type, ) fake_torch.inference_mode = nullcontext fake_api = ModuleType("depth_anything_3.api") fake_api.DepthAnything3 = FakeModel fake_export_module = ModuleType("depth_anything_3.utils.export") fake_export_module.export = fake_export fake_utils = ModuleType("depth_anything_3.utils") fake_utils.__path__ = [] fake_package = ModuleType("depth_anything_3") fake_package.__path__ = [] 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, { "torch": fake_torch, "depth_anything_3": fake_package, "depth_anything_3.api": fake_api, "depth_anything_3.utils": fake_utils, "depth_anything_3.utils.export": fake_export_module, }, ): args = SimpleNamespace( input=str(input_directory), video_frames=16, output_dir=str(root / "result"), export_format="mini_npz", status=str(root / "status.json"), device="cpu", model_dir="model", process_res=518, process_res_method="upper_bound_resize", use_ray_pose=1, ref_view_strategy="first", conf_thresh_percentile=40.0, num_max_points=1000, show_cameras=0, gs_views_interval=1, remove_background_model_dir="", ) self.worker.reconstruct(args) result_directory = Path(args.output_dir) self.assertTrue( (result_directory / "frame.0001.jpg").is_file() ) self.assertTrue( (result_directory / "frame.0002.jpg").is_file() ) self.assertTrue( (result_directory / "results.npz").is_file() ) self.assertFalse( (result_directory / "exports").exists() ) status = json.loads(Path(args.status).read_text(encoding="utf-8")) self.assertEqual(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.worker.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.worker.confidence_percentile_threshold( depth, confidence, 0.0, ) faces = self.worker.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.worker.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.worker.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.worker.depth_edge_mask(smooth).any()) self.assertTrue(self.worker.depth_edge_mask(discontinuous)[:, 2:4].all()) def test_higher_edge_culling_allows_less_stretch(self): limits = [ self.worker.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.worker.triangle_edge_stretch_mask( flat_points, triangles, 1.0, ) loose_mask = self.worker.triangle_edge_stretch_mask( stretched_points, triangles, 2.0, ) strict_mask = self.worker.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.worker.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.worker.confidence_percentile_threshold( depth, confidence, percentile, ) faces = self.worker.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.worker.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_faces_are_removed_before_export(self): faces = np.asarray( ((0, 1, 4), (0, 4, 3), (1, 2, 5), (1, 5, 4)), dtype=np.int32, ) foreground = np.asarray( ((False, True, True), (False, True, True)), dtype=bool, ) filtered = self.worker.remove_background_faces(faces, foreground) np.testing.assert_array_equal(filtered, faces[2:]) def test_image_mesh_data_keeps_camera_geometry_and_uv(self): 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): Path(path).write_bytes(b"texture") fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) depth = np.full((1, 3, 3), 2.0, dtype=np.float32) confidence = np.ones((1, 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, ) prediction = SimpleNamespace( depth=depth, conf=confidence, intrinsics=intrinsics, ) progress = [] with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, {"PIL": fake_pil}, ), patch.object( self.worker, "compact_mesh", side_effect=lambda vertices, faces: ( vertices, faces, ), ): result_path, texture_path = self.worker.export_image_mesh_data( prediction, Path(directory) / "input.png", Path(directory) / "result", progress_callback=lambda step, total, label: progress.append( (step, total, label) ), ) 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["point_confidence_percentile"].shape, (9,), ) self.assertEqual( result["face_edge_stretch"].shape, (8,), ) self.assertEqual(result["face_background"].shape, (8,)) np.testing.assert_allclose(result["uv"][0], [0.0, 1.0]) np.testing.assert_allclose(result["uv"][-1], [1.0, 0.0]) self.assertTrue(texture_path.is_file()) self.assertEqual([item[0] for item in progress], list(range(1, 7))) self.assertTrue(all(item[1] == 6 for item in progress)) def test_prediction_attributes_keep_debug_structure(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.worker.export_prediction_attributes( prediction, directory, ) with np.load(output_path) as attributes: self.assertEqual(attributes["depth"].shape, (1, 2, 3)) self.assertEqual(attributes["conf"].shape, (1, 2, 3)) self.assertEqual( attributes["confidence_percentile"].shape, (1, 2, 3), ) if __name__ == "__main__": unittest.main()