nan
2ad30bf4a3
- 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.
131 lines
4.1 KiB
Python
131 lines
4.1 KiB
Python
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()
|