icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe

Browse Source 
- Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing.
- Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations.
- Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Added model synchronization script `sync_models.py` for downloading and verifying ONNX models.
- Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality.
- Updated `vendor.lock.toml` to remove unused artifacts.
This commit is contained in:
马南
2026-08-14 10:27:09 +08:00
parent 13fab9ce88
commit 2ad30bf4a3
41 changed files with 1569 additions and 2109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
tests/test_addon.py
View File

@@ -280,6 +280,10 @@ class AddonTest(unittest.TestCase):
annotations["snapform_process_res_factor"]["subtype"],
"PERCENTAGE",
)
self.assertEqual(
annotations["snapform_model"]["default"],
0,
)
def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self):
image_mesh = self.addon.image_mesh
@@ -644,8 +648,11 @@ class AddonTest(unittest.TestCase):
def test_conversion_dialog_exposes_requested_settings(self):
operator = self.addon.ConvertImageToDepthPlane
self.assertEqual(operator.__annotations__["snapform_model"]["default"], 1)
self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"])
self.assertEqual(
operator.__annotations__["snapform_model"]["default"],
0,
)
self.assertFalse(operator.__annotations__["snapform_use_ray_pose"]["default"])
self.assertEqual(
operator.__annotations__["snapform_process_res_factor"]["default"],
100.0,
@@ -1094,7 +1101,6 @@ class AddonTest(unittest.TestCase):
for snapform_model, snapform_ray_pose in (
("SMALL", True),
("BASE", False),
("LARGE_1_1", True),
("GIANT_1_1", False),
("NESTED_GIANT_LARGE_1_1", True),
):
@@ -1947,8 +1953,8 @@ class AddonTest(unittest.TestCase):
self.addon.DebugPanel.bl_options,
)
panels = (
self.addon.Da3OutputDebugPanel,
self.addon.Moge2ApiDebugPanel,
self.addon.Da3OutputDebugPanel,
self.addon.Ben2RemovalDebugPanel,
)
for panel in panels:
@@ -1956,7 +1962,7 @@ class AddonTest(unittest.TestCase):
self.assertIn("DEFAULT_CLOSED", panel.bl_options)
self.assertEqual(
[panel.bl_label for panel in panels],
["DA3", "MoGe-2", "BEN2"],
["MoGe-2", "DA3", "BEN2"],
)
def test_service_panel_exposes_shared_service_state_and_controls(self):
@@ -2483,7 +2489,7 @@ class AddonTest(unittest.TestCase):
for index, value in enumerate(command)
if value == "--allow-pattern"
]
self.assertEqual(patterns, ["config.json", "model.safetensors"])
self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
label_index = command.index("--label")
self.assertEqual(command[label_index + 1], "BEN2 Base")
@@ -2512,15 +2518,12 @@ class AddonTest(unittest.TestCase):
self.assertEqual(command[0], sys.executable)
self.assertIn("--uv", command)
self.assertIn("C:/uv/uv.exe", command)
self.assertEqual(
command[command.index("--torch-variant") + 1],
"cu128",
)
self.assertNotIn("--torch-variant", command)
self.assertTrue(
any("install_env.py" in str(value) for value in command)
)
def test_install_env_command_auto_selects_torch_variant(self):
def test_install_env_command_does_not_select_legacy_runtime(self):
client = importlib.import_module("snapform.client")
with patch.object(
client,
@@ -2542,10 +2545,7 @@ class AddonTest(unittest.TestCase):
command = client.install_env_command(Path("status.json"))
self.assertNotIn("--pypi-index", command)
self.assertEqual(
command[command.index("--torch-variant") + 1],
"cpu",
)
self.assertNotIn("--torch-variant", command)
self.assertIn("C:/uv/uv.exe", command)
def test_completed_export_opens_job_directory(self):

53
tests/test_ben2_onnx.py Normal file
View File

@@ -0,0 +1,53 @@
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
import numpy as np
from PIL import Image
SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform"
sys.path.insert(0, str(SOURCE_DIR))
import ben2_onnx
class FakeSession:
def get_inputs(self):
return [SimpleNamespace(name="input")]
def get_outputs(self):
return [SimpleNamespace(name="alpha")]
def run(self, output_names, feeds):
self.output_names = output_names
self.feeds = feeds
gradient = np.linspace(0.0, 1.0, 1024, dtype=np.float32)
return [np.tile(gradient, (1024, 1))[None, None]]
class Ben2OnnxTest(unittest.TestCase):
def test_preprocess_matches_model_layout(self):
image = Image.new("RGB", (40, 20), (64, 128, 192))
result = ben2_onnx.preprocess(image)
self.assertEqual(result.shape, (1, 3, 1024, 1024))
self.assertEqual(result.dtype, np.float32)
self.assertTrue(result.flags.c_contiguous)
def test_infer_restores_original_size_and_alpha(self):
session = FakeSession()
image = Image.new("RGB", (40, 20), (10, 20, 30))
result = ben2_onnx.infer(session, image)
self.assertEqual(result.mode, "RGBA")
self.assertEqual(result.size, image.size)
self.assertEqual(session.feeds["input"].shape, (1, 3, 1024, 1024))
alpha = np.asarray(result.getchannel("A"))
self.assertLess(int(alpha[:, 0].max()), int(alpha[:, -1].min()))
if __name__ == "__main__":
unittest.main()

61
tests/test_da3_onnx.py Normal file
View File

@@ -0,0 +1,61 @@
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
import numpy as np
from PIL import Image
SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform"
sys.path.insert(0, str(SOURCE_DIR))
import da3_onnx
class FakeSession:
def get_inputs(self):
return [SimpleNamespace(name="images"), SimpleNamespace(name="strategy")]
def run(self, output_names, feeds):
self.feeds = feeds
_, views, _, height, width = feeds["images"].shape
depth = np.ones((1, views, height, width), dtype=np.float32)
confidence = np.full_like(depth, 0.75)
pose = np.zeros((1, views, 9), dtype=np.float32)
pose[..., 6] = 1.0
pose[..., 7:] = np.pi / 2.0
return depth, confidence, pose
class Da3OnnxTest(unittest.TestCase):
def test_target_size_uses_patch_multiple(self):
self.assertEqual(
da3_onnx.target_size((800, 600), 504, "upper_bound_resize"),
(504, 378),
)
def test_infer_returns_camera_and_image_arrays(self):
with tempfile.TemporaryDirectory() as directory:
image_path = Path(directory) / "input.png"
Image.new("RGB", (80, 60), (10, 20, 30)).save(image_path)
session = FakeSession()
result = da3_onnx.infer(
session,
[image_path],
process_res=56,
ref_view_strategy="middle",
)
self.assertEqual(result.depth.shape, (1, 42, 56))
self.assertEqual(result.conf.shape, result.depth.shape)
self.assertEqual(result.processed_images.shape, (1, 42, 56, 3))
self.assertEqual(result.extrinsics.shape, (1, 3, 4))
self.assertEqual(result.intrinsics.shape, (1, 3, 3))
self.assertEqual(int(session.feeds["strategy"][0]), 1)
np.testing.assert_allclose(result.extrinsics[0, :3, :3], np.eye(3))
if __name__ == "__main__":
unittest.main()

260
tests/test_install_env.py
View File

@@ -36,7 +36,7 @@ class InstallEnvTest(unittest.TestCase):
(vendor / "moge.tar.gz").write_bytes(b"moge")
return vendor
def test_installs_da3_from_archives_without_gsplat(self):
def test_installs_onnx_runtime_without_training_frameworks(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
@@ -71,77 +71,31 @@ class InstallEnvTest(unittest.TestCase):
commands += [
call.args[4] for call in mocked_run_install.call_args_list
]
da3_command = next(
command
for command in commands
if any("depth-anything-3 @" in str(value) for value in command)
)
self.assertIn(
f"addict=={bootstrap.ADDICT_VERSION}",
da3_command,
)
self.assertTrue(
any(
"depth-anything-3 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in da3_command
)
)
self.assertFalse(
any("git+" in str(value) for value in da3_command)
)
self.assertFalse(
any(
"gsplat" in str(value)
for command in commands
for value in command
)
)
self.assertFalse(
any(
"--no-build-isolation" in str(value)
for command in commands
for value in command
)
)
ben2_command = next(
command
for command in commands
if any("ben2 @" in str(value) for value in command)
)
self.assertTrue(
any(
"ben2 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in ben2_command
)
)
self.assertFalse(
any("git+" in str(value) for value in ben2_command)
)
service_command = next(
command
for command in commands
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
)
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command)
self.assertIn("setuptools", service_command)
self.assertNotIn("ninja", service_command)
self.assertIn(
f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}",
service_command,
)
for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"):
self.assertFalse(
any(forbidden in str(value) for value in service_command),
forbidden,
)
check_command = bootstrap.verify_command(
bootstrap.venv_python(venv_dir)
)
self.assertTrue(
any("import torch;" in str(value) for value in check_command)
)
self.assertTrue(
any(
"from ben2 import BEN_Base" in str(value)
for value in check_command
)
)
self.assertFalse(any("torch" in str(value) for value in check_command))
self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command)
)
self.assertTrue(
any("onnxruntime" in str(value) for value in check_command)
)
self.assertFalse(
any("import gsplat" in str(value) for value in check_command)
)
@@ -157,24 +111,24 @@ class InstallEnvTest(unittest.TestCase):
if any(
repository in command
for repository in (
"Ruicheng/moge-2-vitb-normal",
"PramaLLC/BEN2",
"Ruicheng/moge-2-vits-normal-onnx",
"onnx-community/BEN2-ONNX",
)
)
]
self.assertEqual(len(default_downloads), 2)
self.assertTrue(any("Ruicheng/moge-2-vitb-normal" in map(str, command) for command in default_downloads))
self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads))
ben2_download = next(
command
for command in default_downloads
if "PramaLLC/BEN2" in command
if "onnx-community/BEN2-ONNX" in command
)
patterns = [
ben2_download[index + 1]
for index, value in enumerate(ben2_download)
if value == "--allow-pattern"
]
self.assertEqual(patterns, ["config.json", "model.safetensors"])
self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
self.assertNotIn("--stage", ben2_download)
self.assertNotIn("--stage-label", ben2_download)
self.assertEqual(
@@ -182,7 +136,7 @@ class InstallEnvTest(unittest.TestCase):
"BEN2 Base",
)
def test_darwin_install_skips_cuda_index_and_xformers(self):
def test_darwin_install_uses_default_onnx_runtime(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
@@ -224,15 +178,15 @@ class InstallEnvTest(unittest.TestCase):
commands = [
call.args[4] for call in mocked_run_install.call_args_list
]
torch_command = next(
runtime_command = next(
command
for command in commands
if any(
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value)
for value in command
)
)
self.assertNotIn("--index-url", torch_command)
self.assertNotIn("--index-url", runtime_command)
self.assertFalse(
any(
"xformers" in str(value)
@@ -241,30 +195,27 @@ class InstallEnvTest(unittest.TestCase):
)
)
def test_intel_mac_is_rejected_with_clear_error(self):
def test_intel_mac_uses_portable_onnx_runtime(self):
bootstrap = load_install_env()
args = SimpleNamespace(
uv="uv",
venv_dir=".venv",
models_dir="models",
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status="status.json",
)
with patch.object(
bootstrap.sys,
"platform",
"darwin",
), patch.object(
bootstrap.platform,
"machine",
return_value="x86_64",
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
bootstrap.install(args)
mocked_rmtree.assert_not_called()
):
stages = bootstrap.install_stages(
SimpleNamespace(
models_dir="models",
status="status.json",
),
".venv",
"python",
"uv",
)
runtime = stages[1][1]
self.assertIn(
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}",
runtime,
)
def test_run_stage_writes_stage_progress(self):
bootstrap = load_install_env()
@@ -414,10 +365,10 @@ class InstallEnvTest(unittest.TestCase):
)
self.assertEqual(ends[-1], 1.0)
self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0])
runtime_ranges = stage_ranges[:4]
runtime_ranges = stage_ranges[:2]
self.assertEqual(
[item[0] for item in runtime_ranges],
[1, 2, 3, 3],
[1, 2],
)
self.assertTrue(
all(item[0] is None for item in stage_ranges[4:])
@@ -606,19 +557,19 @@ class InstallEnvTest(unittest.TestCase):
self.assertEqual(bootstrap.format_elapsed(12), "12s")
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
def test_ai_component_estimate_continues_across_internal_steps(self):
def test_onnx_runtime_estimate_decreases_with_progress(self):
bootstrap = load_install_env()
first = bootstrap.estimated_remaining(
"AI components",
bootstrap.XFORMERS_SLICE[0],
)
second = bootstrap.estimated_remaining(
"AI components",
"ONNX Runtime",
bootstrap.AI_COMPONENTS_SLICE[0],
)
second = bootstrap.estimated_remaining(
"ONNX Runtime",
sum(bootstrap.AI_COMPONENTS_SLICE) / 2,
)
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["AI components"])
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["ONNX Runtime"])
self.assertGreater(second, 0)
self.assertLess(second, first)
@@ -727,7 +678,7 @@ class InstallEnvTest(unittest.TestCase):
ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
self.assertEqual(ready["devices"], ["cuda", "cpu"])
self.assertEqual(ready["devices"], ["directml", "cpu"])
self.assertEqual(
ready["manifest_revision"],
bootstrap.MANIFEST_REVISION,
@@ -814,7 +765,7 @@ class InstallEnvTest(unittest.TestCase):
def fake_run(command, environment=None):
commands.append(command)
if any("moge-2-vitb-normal" in str(value) for value in command):
if any("moge-2-vits-normal" in str(value) for value in command):
raise subprocess.CalledProcessError(1, command)
with tempfile.TemporaryDirectory() as directory:
@@ -852,7 +803,7 @@ class InstallEnvTest(unittest.TestCase):
self.assertIn("Environment is installed", status["message"])
self.assertTrue(
any(
"PramaLLC/BEN2" in map(str, command)
"onnx-community/BEN2-ONNX" in map(str, command)
for command in commands
)
)
@@ -983,25 +934,13 @@ class InstallEnvTest(unittest.TestCase):
for value in command
)
)
ai_command = next(
runtime_command = next(
command
for command in install_commands
if any("depth-anything-3 @" in str(value) for value in command)
)
self.assertTrue(
any(
"depth-anything-3 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in ai_command
)
)
self.assertTrue(
any(
"ben2 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in ai_command
)
if any("onnxruntime" in str(value) for value in command)
)
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
self.assertFalse(any("torch" in str(value) for value in runtime_command))
def test_install_reuses_existing_uv(self):
bootstrap = load_install_env()
@@ -1047,81 +986,12 @@ class InstallEnvTest(unittest.TestCase):
self.assertEqual(uv_path.read_bytes(), b"binary")
def test_copy_vendor_inputs_copies_missing_archives(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
bundled = root / "bundled"
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
vendor = root / "vendor"
bootstrap.copy_vendor_inputs(bundled, vendor)
self.assertEqual(
(vendor / "depth-anything-3.tar.gz").read_bytes(),
b"da3",
)
self.assertEqual(
(vendor / "ben2.tar.gz").read_bytes(),
b"ben2",
)
def test_copy_vendor_inputs_replaces_same_size_stale_archive(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
bundled = root / "bundled"
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"new")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"old")
bootstrap.copy_vendor_inputs(bundled, vendor)
self.assertEqual(
(vendor / "depth-anything-3.tar.gz").read_bytes(),
b"new",
)
def test_remove_legacy_vendor_removes_only_install_inputs(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(root / "depth-anything-3.tar.gz").write_bytes(b"da3")
(root / "ben2.tar.gz").write_bytes(b"ben2")
models = root / "models"
models.mkdir()
bootstrap.remove_legacy_vendor(root)
self.assertFalse(vendor.exists())
self.assertFalse((root / "depth-anything-3.tar.gz").exists())
self.assertFalse((root / "ben2.tar.gz").exists())
self.assertTrue(models.is_dir())
def test_install_stages_uses_local_source_archives(self):
def test_install_stages_use_only_runtime_packages(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
@@ -1141,16 +1011,14 @@ class InstallEnvTest(unittest.TestCase):
uv_path,
)
ai_command = next(
runtime_command = next(
command
for _label, command, _environment, kind in stages
if kind == "install"
and any("depth-anything-3" in str(value) for value in command)
and any("onnxruntime" in str(value) for value in command)
)
da3_local = (vendor / "depth-anything-3.tar.gz").as_uri()
ben2_local = (vendor / "ben2.tar.gz").as_uri()
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
self.assertFalse(any("torch" in str(value) for value in runtime_command))
def test_install_stages_download_only_default_models(self):
bootstrap = load_install_env()
@@ -1189,13 +1057,13 @@ class InstallEnvTest(unittest.TestCase):
if kind == "download"
}
self.assertIn(
"moge2_vitb_normal",
" ".join(map(str, downloads["MoGe-2 ViT-B Normal"])),
"moge-2-vits-normal-onnx",
" ".join(map(str, downloads["MoGe-2 ViT-S Normal"])),
)
self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"])))
self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"])))
self.assertEqual(
set(downloads),
{"MoGe-2 ViT-B Normal", "BEN2 Base"},
{"MoGe-2 ViT-S Normal", "BEN2 Base"},
)
def test_install_does_not_verify_env(self):

64
tests/test_models.py
View File

@@ -30,12 +30,6 @@ class ModelTest(unittest.TestCase):
self.assertEqual(model.repository, "depth-anything/DA3-BASE")
self.assertEqual(model.license_name, "Apache 2.0")
def test_large_model_exposes_noncommercial_license(self):
model = self.models.get_model("LARGE_1_1")
self.assertEqual(model.license_name, "CC BY-NC 4.0")
self.assertFalse(model.enabled)
def test_giant_models_expose_gaussian_head(self):
for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"):
with self.subTest(model_key=model_key):
@@ -54,19 +48,20 @@ class ModelTest(unittest.TestCase):
def test_ben2_download_only_includes_required_hugging_face_files(self):
model = self.models.get_downloadable_model("BEN2_BASE")
self.assertEqual(model.repository, "PramaLLC/BEN2")
self.assertEqual(model.license_name, "MIT")
self.assertEqual(
model.download_patterns,
("config.json", "model.safetensors"),
)
self.assertEqual(model.repository, "onnx-community/BEN2-ONNX")
self.assertEqual(model.license_name, "Apache 2.0")
self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",))
self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",))
self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx")
self.assertNotIn("BEN2_BASE", self.models.MODELS)
def test_da3_model_directories_share_the_da3_prefix(self):
def test_da3_model_directories_use_repository_names(self):
for model_key, model in self.models.MODELS.items():
if not model.enabled:
continue
with self.subTest(model_key=model_key):
self.assertTrue(
model.directory_name.startswith("da3_"),
model.directory_name.startswith("DA3-"),
model.directory_name,
)
self.assertEqual(
@@ -74,20 +69,37 @@ class ModelTest(unittest.TestCase):
model.directory_name,
)
def test_ben2_model_directory_uses_ben2_prefix(self):
def test_ben2_model_directory_uses_repository_name(self):
self.assertEqual(
self.models.model_directory_name("BEN2_BASE"),
"ben2_base",
"BEN2-ONNX",
)
def test_moge2_uses_single_checkpoint_and_mit_license(self):
model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL")
def test_enabled_da3_models_use_onnx_and_r2_checksums(self):
for model in self.models.MODELS.values():
if not model.enabled:
continue
with self.subTest(model=model.key):
self.assertTrue(model.download_patterns[0].endswith(".onnx"))
self.assertEqual(model.ready_patterns, ("model.onnx",))
self.assertEqual(model.r2_files[0][0], "model.onnx")
self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal")
self.assertEqual(model.license_name, "MIT")
self.assertEqual(model.download_patterns, ("model.pt",))
self.assertEqual(model.ready_patterns, ("model.pt",))
self.assertFalse(model.supports_gs)
def test_moge2_models_use_official_onnx_models(self):
for key, variant in (("MOGE2_VITS_NORMAL", "vits"),
("MOGE2_VITB_NORMAL", "vitb"),
("MOGE2_VITL_NORMAL", "vitl")):
with self.subTest(key=key):
model = self.models.get_downloadable_model(key)
self.assertEqual(
model.repository,
f"Ruicheng/moge-2-{variant}-normal-onnx",
)
self.assertEqual(model.license_name, "MIT")
self.assertEqual(model.download_patterns, ("model.onnx",))
self.assertEqual(model.ready_patterns, ("model.onnx",))
self.assertEqual(model.r2_directory, model.repository.split("/")[-1])
self.assertEqual(model.r2_files[0][0], "model.onnx")
self.assertFalse(model.supports_gs)
def test_moge2_models_match_uploaded_r2_variants(self):
self.assertEqual(
@@ -96,5 +108,9 @@ class ModelTest(unittest.TestCase):
)
self.assertEqual(
[model.r2_directory for model in self.models.MOGE2_MODELS.values()],
["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"],
[
"moge-2-vits-normal-onnx",
"moge-2-vitb-normal-onnx",
"moge-2-vitl-normal-onnx",
],
)

130
tests/test_moge_onnx.py Normal file
View File

@@ -0,0 +1,130 @@
import importlib.util
import sys
import unittest
from pathlib import Path
from unittest.mock import patch
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_module():
specification = importlib.util.spec_from_file_location(
"snapform_moge_onnx",
PROJECT_ROOT / "src" / "snapform" / "moge_onnx.py",
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class FakeInput:
def __init__(self, name):
self.name = name
class FakeSession:
def __init__(self, outputs):
self.outputs = outputs
self.calls = []
def run(self, names, inputs):
self.calls.append((names, inputs))
return [self.outputs[name] for name in names]
def get_outputs(self):
return [FakeInput(name) for name in self.outputs]
class MogeOnnxTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.moge = load_module()
def test_auto_provider_prefers_accelerated_runtime(self):
providers = self.moge.select_providers(
"auto",
["CPUExecutionProvider", "DmlExecutionProvider"],
)
self.assertEqual(
providers,
["DmlExecutionProvider", "CPUExecutionProvider"],
)
def test_explicit_unavailable_provider_falls_back_to_cpu(self):
providers = self.moge.select_providers(
"cuda",
["CPUExecutionProvider"],
)
self.assertEqual(providers, ["CPUExecutionProvider"])
def test_infer_uses_dynamic_token_input_and_numpy_postprocess(self):
height, width = 8, 12
raw = {
"points": np.zeros((1, height, width, 3), dtype=np.float32),
"normal": np.zeros((1, height, width, 3), dtype=np.float32),
"mask": np.ones((1, height, width), dtype=np.float32),
"scale": np.ones((1,), dtype=np.float32),
}
raw["points"][..., 2] = 1.0
raw["normal"][..., 2] = 1.0
session = FakeSession(raw)
expected = {
"points": np.zeros((height, width, 3), dtype=np.float32),
"depth": np.ones((height, width), dtype=np.float32),
"normal": raw["normal"][0],
"mask": np.ones((height, width), dtype=bool),
"intrinsics": np.eye(3, dtype=np.float32),
}
with patch.object(self.moge, "postprocess", return_value=expected):
result = self.moge.infer(
session,
np.zeros((height, width, 3), dtype=np.float32),
resolution_level=9,
)
self.assertIs(result, expected)
names, inputs = session.calls[0]
self.assertEqual(names, ["points", "normal", "mask", "scale"])
self.assertEqual(inputs["image"].shape, (1, 3, height, width))
self.assertEqual(inputs["num_tokens"].dtype, np.int64)
self.assertEqual(inputs["num_tokens"], 3600)
def test_postprocess_recovers_camera_projection_and_metric_scale(self):
height, width = 64, 80
focal = np.float32(0.9)
shift = np.float32(0.3)
metric_scale = np.float32(2.0)
uv = self.moge.normalized_view_plane_uv(width, height, np.float32)
depth = np.linspace(1.0, 2.4, height * width, dtype=np.float32).reshape(
height,
width,
)
raw_points = np.concatenate(
(uv * depth[..., None] / focal, (depth - shift)[..., None]),
axis=-1,
)
raw = {
"points": raw_points[None],
"normal": np.broadcast_to(
np.array((0.0, 0.0, 1.0), dtype=np.float32),
(1, height, width, 3),
),
"mask": np.ones((1, height, width), dtype=np.float32),
"metric_scale": np.array((metric_scale,), dtype=np.float32),
}
result = self.moge.postprocess(raw)
np.testing.assert_allclose(result["depth"], depth * metric_scale, rtol=1e-3)
np.testing.assert_allclose(result["points"][..., 2], result["depth"], rtol=1e-5)
self.assertTrue(result["mask"].all())
if __name__ == "__main__":
unittest.main()

95
tests/test_pack.py
View File

@@ -1,7 +1,4 @@
import ast
import importlib.util
import io
import tarfile
import tempfile
import tomllib
import unittest
@@ -64,8 +61,6 @@ class PackTest(unittest.TestCase):
self.assertIn("service.py", names)
self.assertNotIn("snapform/__init__.py", names)
self.assertNotIn("importer.py", names)
self.assertIn("vendor/depth-anything-3.tar.gz", names)
self.assertIn("vendor/ben2.tar.gz", names)
self.assertIn(
f"vendor/{pack.PLATFORMS[platform]}",
names,
@@ -108,95 +103,5 @@ class PackTest(unittest.TestCase):
{"windows-x64", "macos-arm64", "linux-x64"},
)
def test_bundled_depth_anything_3_has_no_xformers_dependency(self):
pack = load_pack()
with tarfile.open(
fileobj=io.BytesIO(
pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
),
mode="r:gz",
) as archive:
raw_pyproject = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"pyproject.toml"
).read()
pyproject = raw_pyproject.decode("utf-8")
self.assertNotIn("xformers", pyproject)
self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
def test_bundled_depth_anything_3_exports_are_lazy(self):
pack = load_pack()
with tarfile.open(
fileobj=io.BytesIO(
pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
),
mode="r:gz",
) as archive:
source = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"src/depth_anything_3/utils/export/__init__.py"
).read().decode("utf-8")
tree = ast.parse(source)
module_imports = [
ast.get_source_segment(source, node)
for node in tree.body
if isinstance(node, (ast.Import, ast.ImportFrom))
]
self.assertEqual(
module_imports,
["from depth_anything_3.specs import Prediction"],
)
def test_official_depth_anything_3_baseline_is_pristine(self):
pack = load_pack()
archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz"
with tarfile.open(archive_path, "r:gz") as archive:
raw_pyproject = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"pyproject.toml"
).read()
export_source = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"src/depth_anything_3/utils/export/__init__.py"
).read().decode("utf-8")
pyproject = raw_pyproject.decode("utf-8")
self.assertIn("xformers", pyproject)
self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
self.assertIn(
"from depth_anything_3.utils.export.gs import",
export_source,
)
def test_bundled_depth_anything_3_qr_falls_back_on_mps(self):
pack = load_pack()
with tarfile.open(
fileobj=io.BytesIO(
pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
),
mode="r:gz",
) as archive:
source = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"src/depth_anything_3/utils/ray_utils.py"
).read().decode("utf-8")
self.assertIn('use_cpu = device.type == "mps"', source)
self.assertIn("A_tilde = A_tilde.cpu()", source)
self.assertIn("Q_tilde = Q_tilde.to(device=device)", source)
def test_official_depth_anything_3_ray_utils_is_pristine(self):
pack = load_pack()
archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz"
with tarfile.open(archive_path, "r:gz") as archive:
source = archive.extractfile(
"Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
"src/depth_anything_3/utils/ray_utils.py"
).read().decode("utf-8")
self.assertIn(
"Q_tilde, R_tilde = torch.linalg.qr(A_tilde)",
source,
)
self.assertNotIn("use_cpu", source)
if __name__ == "__main__":
unittest.main()

60
tests/test_service.py
View File

@@ -40,7 +40,6 @@ class ServiceTest(unittest.TestCase):
def create_cache(self):
cache = self.service.ModelCache.__new__(self.service.ModelCache)
cache.torch = SimpleNamespace(cuda=FakeCuda())
cache.ben2_model = None
cache.ben2_key = None
cache.da3_model = None
@@ -52,65 +51,41 @@ class ServiceTest(unittest.TestCase):
def test_ben2_model_is_reused_for_matching_runtime_key(self):
loads = []
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
loads.append(model_dir)
return cls()
def to(self, device):
self.device = device
return self
def eval(self):
return None
ben2 = ModuleType("ben2")
ben2.BEN_Base = FakeModel
ben2 = ModuleType("ben2_onnx")
ben2.create_session = lambda model_dir, device: (
loads.append((model_dir, device)) or object()
)
cache = self.create_cache()
device = SimpleNamespace(type="cuda")
with patch.dict(sys.modules, {"ben2": ben2}):
with patch.dict(sys.modules, {"ben2_onnx": ben2}):
first, _load_ms = cache.get_ben2("model", device)
second, cached_load_ms = cache.get_ben2("model", device)
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(loads, [("model", "cuda")])
self.assertEqual(cache.status()["ben2"], "model")
self.assertEqual(cached_load_ms, 0.0)
def test_da3_model_is_reused_for_matching_runtime_key(self):
loads = []
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
loads.append(model_dir)
return cls()
def to(self, device):
self.device = device
return self
def eval(self):
return None
package = ModuleType("depth_anything_3")
api = ModuleType("depth_anything_3.api")
api.DepthAnything3 = FakeModel
da3 = ModuleType("da3_onnx")
da3.create_session = lambda model_dir, device: (
loads.append((model_dir, device)) or object()
)
cache = self.create_cache()
device = SimpleNamespace(type="cuda")
with patch.dict(
sys.modules,
{
"depth_anything_3": package,
"depth_anything_3.api": api,
"da3_onnx": da3,
},
):
first = cache.get_da3("model", device)
second = cache.get_da3("model", device)
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(loads, [("model", "cuda")])
self.assertEqual(cache.status()["da3"], "model")
def test_depth_model_cache_keeps_only_da3_or_moge(self):
@@ -135,13 +110,10 @@ class ServiceTest(unittest.TestCase):
moge = object()
def load_moge(_model_dir, _device, key):
cache.moge_model = moge
cache.moge_key = key
return moge
with patch.object(cache, "_load_moge", side_effect=load_moge):
self.assertIs(cache.get_moge("moge", device), moge)
fake_moge = ModuleType("moge_onnx")
fake_moge.create_session = lambda _model_dir, _device: moge
with patch.dict(sys.modules, {"moge_onnx": fake_moge}):
self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge)
self.assertIsNone(cache.da3_model)
self.assertIsNotNone(cache.ben2_model)

278
tests/test_worker.py
View File

@@ -66,57 +66,6 @@ class WorkerTest(unittest.TestCase):
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(
[
@@ -436,12 +385,6 @@ class WorkerTest(unittest.TestCase):
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")
@@ -605,8 +548,11 @@ class WorkerTest(unittest.TestCase):
)
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_ben2 = ModuleType("ben2_onnx")
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())
@@ -614,7 +560,7 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{
"torch": fake_torch,
"ben2": fake_ben2,
"ben2_onnx": fake_ben2,
"PIL": fake_pil,
},
):
@@ -635,7 +581,7 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(
inference_options[0][1],
{"refine_foreground": False},
{},
)
self.assertEqual(
save_options[0],
@@ -690,8 +636,11 @@ class WorkerTest(unittest.TestCase):
)
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_ben2 = ModuleType("ben2_onnx")
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())
@@ -699,7 +648,7 @@ class WorkerTest(unittest.TestCase):
sys.modules,
{
"torch": fake_torch,
"ben2": fake_ben2,
"ben2_onnx": fake_ben2,
"PIL": fake_pil,
},
):
@@ -734,7 +683,7 @@ class WorkerTest(unittest.TestCase):
self.assertEqual(len(inference_options), 2)
for _image, options in inference_options:
self.assertEqual(options, {"refine_foreground": False})
self.assertEqual(options, {})
def test_remove_background_parser_accepts_video_frames(self):
args = self.worker.parser().parse_args(
@@ -813,64 +762,25 @@ class WorkerTest(unittest.TestCase):
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,
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),
)
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__ = []
fake_da3 = ModuleType("da3_onnx")
fake_da3.create_session = lambda _model_dir, _device: object()
def fake_infer(_session, images, **options):
inference_options.append((images, options))
return prediction
fake_da3.infer = fake_infer
fake_runtime = ModuleType("onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
def fake_mesh_export(
prediction,
@@ -893,6 +803,20 @@ class WorkerTest(unittest.TestCase):
output_path.touch()
return output_path
def fake_point_cloud(_prediction, output_directory, *_args):
(Path(output_directory) / "scene.glb").touch()
def fake_depth_preview(_prediction, output_directory):
target = Path(output_directory) / "depth_vis"
target.mkdir()
(target / "0000.jpg").touch()
def fake_mini_npz(_prediction, output_directory):
np.savez_compressed(
Path(output_directory) / "results.npz",
depth=np.zeros((1, 2, 2)),
)
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
@@ -900,20 +824,28 @@ class WorkerTest(unittest.TestCase):
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,
"da3_onnx": fake_da3,
"onnx_runtime": fake_runtime,
},
), patch.object(
self.worker,
"export_textured_mesh_glb",
side_effect=fake_mesh_export,
), patch.object(
self.worker,
"export_da3_point_cloud",
side_effect=fake_point_cloud,
), patch.object(
self.worker,
"export_da3_depth_preview",
side_effect=fake_depth_preview,
), patch.object(
self.worker,
"export_da3_mini_npz",
side_effect=fake_mini_npz,
):
debug_formats = (
"glb",
"gs_ply",
"depth_vis",
"mini_npz",
"mesh",
@@ -945,15 +877,11 @@ class WorkerTest(unittest.TestCase):
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(inference_options), 4)
self.assertEqual(len(mesh_calls), 1)
for images, options in inference_options:
self.assertEqual(images, [str(input_path)])
@@ -964,82 +892,21 @@ class WorkerTest(unittest.TestCase):
)
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_da3 = ModuleType("da3_onnx")
fake_da3.create_session = lambda _model_dir, _device: object()
fake_da3.infer = lambda *_args, **_kwargs: SimpleNamespace(
depth=np.full((2, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((2, 2, 2), dtype=np.float32),
extrinsics=np.repeat(np.eye(4, dtype=np.float32)[None, :3], 2, axis=0),
intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0),
processed_images=np.zeros((2, 2, 2, 3), dtype=np.uint8),
)
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__ = []
fake_runtime = ModuleType("onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
@@ -1050,11 +917,8 @@ class WorkerTest(unittest.TestCase):
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,
"da3_onnx": fake_da3,
"onnx_runtime": fake_runtime,
},
):
args = SimpleNamespace(