Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 6 additions & 12 deletions tests/integration/test_reconstruction_direct.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

from __future__ import annotations

import importlib.util
import json
import os
import time
Expand All @@ -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
Expand All @@ -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"))
Expand All @@ -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()
Expand Down
18 changes: 6 additions & 12 deletions tests/integration/test_reconstruction_grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from __future__ import annotations

import importlib.util
import json
import os
import time
Expand All @@ -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"))
Expand All @@ -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()
Expand Down
18 changes: 6 additions & 12 deletions tests/integration/test_reconstruction_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from __future__ import annotations

import base64
import importlib.util
import json
import os
import time
Expand All @@ -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

Expand All @@ -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()
Expand Down
13 changes: 13 additions & 0 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

34 changes: 33 additions & 1 deletion vizion3d/reconstruction/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import os
import sys
import threading
import types
from pathlib import Path
from types import ModuleType

Expand Down Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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"
Expand Down
Loading