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

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)