Merge pull request #57 from OlafenwaMoses/moses/vizion3d-playground

fix failing bundle missing file

Author: OlafenwaMoses

pyproject.toml

@@ -157,6 +157,15 @@ dev = [
     "ruff>=0.9.0",
     "twine>=6.0.0",
 ]
+torchmcubes = [
+    "torchmcubes",
+]
+
+[tool.uv.sources]
+torchmcubes = { git = "https://github.com/tatsy/torchmcubes.git" }
+
+[tool.uv.extra-build-dependencies]
+torchmcubes = ["torch"]
 
 # ---------------------------------------------------------------------------
 # Hatch build — wheel & sdist

tests/integration/test_reconstruction_direct.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import importlib.util
 import json
 import os
 import time
@@ -18,12 +19,8 @@
 import pytest
 from PIL import Image
 
-pytest.importorskip("open3d", reason="open3d required")
-pytest.importorskip("trimesh", reason="trimesh required")
-pytest.importorskip("rembg", reason="rembg required")
-pytest.importorskip("mcubes", reason="PyMCubes required")
-pytest.importorskip("omegaconf", reason="omegaconf required")
-pytest.importorskip("einops", reason="einops required")
+for module_name in ("open3d", "trimesh", "rembg", "mcubes", "omegaconf", "einops"):
+    assert importlib.util.find_spec(module_name), f"{module_name} is required"
 
 from vizion3d.lifting.utils import create_ply_binary  # noqa: E402
 from vizion3d.reconstruction import (  # noqa: E402
@@ -34,6 +31,7 @@
     SceneComponents3DReconstructionCommand,
     SceneComponents3DReconstructionConfig,
 )
+from vizion3d.reconstruction.defaults import resolve_model_bundle  # noqa: E402
 from vizion3d.reconstruction.handlers import Object3DReconstructionHandler  # noqa: E402
 
 OBJECT_LIMIT = float(os.environ.get("VIZION3D_TEST_RECON_OBJECT_LIMIT", "240.0"))
@@ -44,17 +42,13 @@
 
 @pytest.fixture(scope="session")
 def reconstruction_model_bundle() -> str:
-    bundle = Path(__file__).resolve().parents[2] / "scene-components-3d-models.zip"
-    if not bundle.is_file():
-        pytest.skip(f"Reconstruction model bundle not found: {bundle}")
-    return str(bundle)
+    return str(resolve_model_bundle())
 
 
 @pytest.fixture(scope="session")
 def reconstruction_image_bytes() -> bytes:
     path = Path(__file__).parent.parent / "assets" / RECONSTRUCTION_IMAGE
-    if not path.is_file():
-        pytest.skip(f"Reconstruction integration image not found: {path}")
+    assert path.is_file(), f"Reconstruction integration image not found: {path}"
     image = Image.open(path)
     assert max(image.size) <= TEST_IMAGE_MAX_DIMENSION
     return path.read_bytes()

tests/integration/test_reconstruction_grpc.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import importlib.util
 import json
 import os
 import time
@@ -15,14 +16,11 @@
 import pytest
 from PIL import Image
 
-pytest.importorskip("open3d", reason="open3d required")
-pytest.importorskip("trimesh", reason="trimesh required")
-pytest.importorskip("rembg", reason="rembg required")
-pytest.importorskip("mcubes", reason="PyMCubes required")
-pytest.importorskip("omegaconf", reason="omegaconf required")
-pytest.importorskip("einops", reason="einops required")
+for module_name in ("open3d", "trimesh", "rembg", "mcubes", "omegaconf", "einops"):
+    assert importlib.util.find_spec(module_name), f"{module_name} is required"
 
 from vizion3d.proto import lifting_pb2  # noqa: E402
+from vizion3d.reconstruction.defaults import resolve_model_bundle  # noqa: E402
 from vizion3d.reconstruction.handlers import Object3DReconstructionHandler  # noqa: E402
 
 OBJECT_LIMIT = float(os.environ.get("VIZION3D_TEST_RECON_GRPC_OBJECT_LIMIT", "260.0"))
@@ -33,17 +31,13 @@
 
 @pytest.fixture(scope="session")
 def reconstruction_model_bundle() -> str:
-    bundle = Path(__file__).resolve().parents[2] / "scene-components-3d-models.zip"
-    if not bundle.is_file():
-        pytest.skip(f"Reconstruction model bundle not found: {bundle}")
-    return str(bundle)
+    return str(resolve_model_bundle())
 
 
 @pytest.fixture(scope="session")
 def reconstruction_image_bytes() -> bytes:
     path = Path(__file__).parent.parent / "assets" / RECONSTRUCTION_IMAGE
-    if not path.is_file():
-        pytest.skip(f"Reconstruction integration image not found: {path}")
+    assert path.is_file(), f"Reconstruction integration image not found: {path}"
     image = Image.open(path)
     assert max(image.size) <= TEST_IMAGE_MAX_DIMENSION
     return path.read_bytes()

tests/integration/test_reconstruction_rest.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import base64
+import importlib.util
 import json
 import os
 import time
@@ -17,13 +18,10 @@
 from fastapi.testclient import TestClient
 from PIL import Image
 
-pytest.importorskip("open3d", reason="open3d required")
-pytest.importorskip("trimesh", reason="trimesh required")
-pytest.importorskip("rembg", reason="rembg required")
-pytest.importorskip("mcubes", reason="PyMCubes required")
-pytest.importorskip("omegaconf", reason="omegaconf required")
-pytest.importorskip("einops", reason="einops required")
+for module_name in ("open3d", "trimesh", "rembg", "mcubes", "omegaconf", "einops"):
+    assert importlib.util.find_spec(module_name), f"{module_name} is required"
 
+from vizion3d.reconstruction.defaults import resolve_model_bundle  # noqa: E402
 from vizion3d.reconstruction.handlers import Object3DReconstructionHandler  # noqa: E402
 from vizion3d.server.rest.app import app  # noqa: E402
 
@@ -37,17 +35,13 @@
 
 @pytest.fixture(scope="session")
 def reconstruction_model_bundle() -> str:
-    bundle = Path(__file__).resolve().parents[2] / "scene-components-3d-models.zip"
-    if not bundle.is_file():
-        pytest.skip(f"Reconstruction model bundle not found: {bundle}")
-    return str(bundle)
+    return str(resolve_model_bundle())
 
 
 @pytest.fixture(scope="session")
 def reconstruction_image_bytes() -> bytes:
     path = Path(__file__).parent.parent / "assets" / RECONSTRUCTION_IMAGE
-    if not path.is_file():
-        pytest.skip(f"Reconstruction integration image not found: {path}")
+    assert path.is_file(), f"Reconstruction integration image not found: {path}"
     image = Image.open(path)
     assert max(image.size) <= TEST_IMAGE_MAX_DIMENSION
     return path.read_bytes()

uv.lock

@@ -8,6 +8,7 @@
 import os
 import sys
 import threading
+import types
 from pathlib import Path
 from types import ModuleType
 
@@ -100,6 +101,35 @@ def _module_file(module: ModuleType | None) -> Path | None:
     return Path(module.__file__).resolve()
 
 
+def _ensure_torchmcubes_compat() -> None:
+    try:
+        import torchmcubes  # noqa: F401
+
+        return
+    except ImportError:
+        pass
+
+    try:
+        import mcubes
+        import torch
+    except ImportError as exc:
+        raise ImportError(
+            "3D reconstruction requires torchmcubes or PyMCubes for marching cubes."
+        ) from exc
+
+    module = types.ModuleType("torchmcubes")
+
+    def marching_cubes(volume, threshold):
+        vertices, faces = mcubes.marching_cubes(volume.detach().cpu().numpy(), threshold)
+        return (
+            torch.as_tensor(vertices, dtype=volume.dtype, device=volume.device),
+            torch.as_tensor(faces, dtype=torch.long, device=volume.device),
+        )
+
+    module.marching_cubes = marching_cubes
+    sys.modules["torchmcubes"] = module
+
+
 def _import_tsr_from_bundle(root: Path):
     source = (root / "TripoSR").resolve()
     system_py = source / "tsr" / "system.py"
@@ -117,6 +147,7 @@ def _import_tsr_from_bundle(root: Path):
     if source_str in sys.path:
         sys.path.remove(source_str)
     sys.path.insert(0, source_str)
+    _ensure_torchmcubes_compat()
 
     try:
         from tsr.system import TSR
@@ -172,7 +203,8 @@ def _load_model(self, model_bundle: str | None, requested_device: str):
                 import trimesh
             except ImportError as exc:
                 raise ImportError(
-                    "3D reconstruction requires torch, trimesh, einops, omegaconf, and PyMCubes."
+                    "3D reconstruction requires torch, trimesh, einops, omegaconf, "
+                    "and PyMCubes or torchmcubes."
                 ) from exc
 
             triposr = root / "TripoSR"