import importlib.util import inspect 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_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_reconstruct_reports_status_before_loading_torch(self): source = inspect.getsource(self.worker.reconstruct) self.assertLess( source.index('"Starting AI runtime"'), source.index("import torch"), ) 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) self.assertIn( "cancel_check=cancel_check", inspect.getsource(self.worker.reconstruct), ) self.assertIn( "cancel_check=cancel_check", inspect.getsource(self.worker.remove_background), ) 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_optional_background_removal_runs_before_da3_model_load(self): source = inspect.getsource(self.worker.reconstruct) self.assertLess( source.index("infer_foreground("), source.index("DepthAnything3.from_pretrained"), ) 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_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_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_download_passes_optional_file_patterns(self): calls = [] fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face.snapshot_download = lambda **options: calls.append( options ) args = SimpleNamespace( repository="PramaLLC/BEN2", model_dir="model", status="status.json", allow_pattern=["config.json", "model.safetensors"], ) 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.object( self.worker, "resolve_hf_endpoint", return_value=(self.worker.HF_ENDPOINT, False), ): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") self.worker.download_model(args) self.assertEqual( calls[0]["allow_patterns"], ["config.json", "model.safetensors"], ) def test_download_model_switches_to_mirror_before_hf_import(self): calls = [] fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face.snapshot_download = lambda **options: calls.append( options ) 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.object( self.worker, "resolve_hf_endpoint", return_value=("https://hf-mirror.com", True), ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( self.worker, "write_status", ) as mocked_write: args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") self.worker.download_model(args) self.assertEqual( os.environ["HF_ENDPOINT"], "https://hf-mirror.com", ) messages = [call.args[2] for call in mocked_write.call_args_list] self.assertTrue(any("(mirror)" in message for message in messages)) def test_ben2_background_removal_writes_fast_transparent_png(self): inference_options = [] save_options = [] 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): 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_uses_official_da3_export_parameters_for_all_four_formats(self): inference_options = [] export_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() 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__ = [] 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, }, ): for export_format in self.worker.EXPORT_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, ) 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"]) self.assertEqual(len(inference_options), 4) self.assertEqual(len(export_calls), 4) 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: 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, {}) infer_gs_values = [options["infer_gs"] for _, options in inference_options] self.assertEqual(infer_gs_values, [False, True, False, False]) if __name__ == "__main__": unittest.main()