Add reconstruction tasks

Author: OlafenwaMoses

docs/api/reconstruction.md

@@ -0,0 +1,3 @@
+# Reconstruction API
+
+::: vizion3d.reconstruction

docs/hardware_acceleration.md

@@ -124,7 +124,7 @@ uv add "vizion3d[cuda]"
 
 Because `vizion3d[cuda]` declares no torch dependency, pip will not touch the CUDA wheel installed in step 1.
 
-vizion3d detects CUDA via `torch.cuda.is_available()` at runtime and moves models and tensors to the GPU automatically.
+vizion3d detects CUDA via `torch.cuda.is_available()` at runtime and moves models and tensors to the GPU automatically. For reconstruction workloads, combine the extras as `vizion3d[cuda,reconstruction]`; the CUDA extra also installs `onnxruntime-gpu` so mandatory `rembg` background removal can use ONNX Runtime's CUDA provider when available.
 
 ---
 

docs/reconstruction/object_3d_reconstruction.md

@@ -0,0 +1,43 @@
+# Object 3D Reconstruction
+
+`Object3DReconstruction` takes a close-range image of one object and produces:
+
+- a cleaned, uniformly gray PLY mesh;
+- a uniformly gray point cloud sampled from the mesh surface.
+
+Install the runtime dependencies with `pip install "vizion3d[reconstruction]"`.
+The task resolves `scene-components-3d-models.zip` from the repository root,
+`~/.cache/vizion3d/models`, or `VIZION3D_RECONSTRUCTION_MODEL_BUNDLE`.
+Set `VIZION3D_TRIPOSR_SOURCE` when the TripoSR Python source is outside the
+repository's `research/3D_Object-Reconstruction/TripoSR` directory.
+
+```python
+from vizion3d.reconstruction import (
+    Object3DReconstruction,
+    Object3DReconstructionCommand,
+)
+
+result = Object3DReconstruction().run(
+    Object3DReconstructionCommand(image_input="object.png")
+)
+mesh = result.mesh
+point_cloud = result.point_cloud
+```
+
+Inputs always use the bundled `rembg/u2net.onnx` model for background removal.
+
+## Device
+
+`Object3DReconstructionConfig(device="auto")` propagates the selected device to
+TripoSR and to `rembg` where the installed ONNX Runtime providers support it.
+`auto` prefers CUDA, then Apple/CoreML-compatible acceleration for `rembg`, then
+CPU. If an accelerated TripoSR or `rembg` run fails, the task retries that stage
+on CPU.
+
+## Input Resolution
+
+The task limits the longest input-image dimension to 1080 pixels before
+background removal. The resize preserves aspect ratio. This avoids spending
+memory and inference time on source pixels that cannot pass through TripoSR's
+final 512 by 512 conditioning input. The config may lower this limit, but
+values above 1080 are rejected.

docs/reconstruction/scene_components_3d_reconstruction.md

@@ -0,0 +1,40 @@
+# Scene Components 3D Reconstruction
+
+`SceneComponents3DReconstruction` accepts one scene image, estimates depth,
+detects and segments objects, maps each mask back to the original-resolution
+image, enhances each object crop with Real-ESRGAN, then runs
+`Object3DReconstruction` for each selected component.
+
+Each selected crop always goes through `rembg` background removal inside
+`Object3DReconstruction`; there is no scene or object option to skip it.
+
+Each component contains a uniformly gray mesh and point cloud, together with
+its label, confidence, source bounding box, and geometry counts.
+
+```python
+from vizion3d.reconstruction import (
+    SceneComponents3DReconstruction,
+    SceneComponents3DReconstructionCommand,
+)
+
+result = SceneComponents3DReconstruction().run(
+    SceneComponents3DReconstructionCommand(image_input="scene.jpg")
+)
+```
+
+## Device
+
+The nested object config's `device` setting is propagated through the scene
+pipeline. TripoSR, `rembg`, and scene Real-ESRGAN use the requested accelerator
+when the installed runtime supports it, and retry on CPU if the accelerated
+stage fails. Mesh cleanup and point sampling remain CPU operations because they
+are handled by `trimesh`.
+
+## Input Resolution
+
+The scene-level `max_input_dimension=1080` applies to depth and segmentation
+analysis. Object crops are still taken from the original image, enhanced with
+Real-ESRGAN, then independently capped at 1080 pixels by
+`Object3DReconstruction` before foreground processing. TripoSR ultimately
+conditions every crop at its required 512 by 512 input size. Set the
+scene-level limit to `0` to disable only the depth and segmentation resize.

docs/tasks/index.md

@@ -22,3 +22,10 @@ has a direct Python facade and, where available, REST and gRPC adapters.
 | Task | Python import | REST | gRPC |
 |---|---|---|---|
 | [Scale Observation](../observation/scale_observation.md) | `vizion3d.observation.ScaleObservation` | `/observation/scale-observation` | `RunScaleObservation` |
+
+## Reconstruction
+
+| Task | Python import | REST | gRPC |
+|---|---|---|---|
+| [Object 3D Reconstruction](../reconstruction/object_3d_reconstruction.md) | `vizion3d.reconstruction.Object3DReconstruction` | `/reconstruction/object-3d-reconstruction` | `RunObject3DReconstruction` |
+| [Scene Components 3D Reconstruction](../reconstruction/scene_components_3d_reconstruction.md) | `vizion3d.reconstruction.SceneComponents3DReconstruction` | `/reconstruction/scene-components-3d-reconstruction` | `RunSceneComponents3DReconstruction` |

mkdocs.yml

@@ -40,8 +40,12 @@ nav:
           - YOLOE-26L Prompt-Free Classes: annotation/yoloe_26l_prompt_free_classes.md
       - Observation:
           - Scale Observation: observation/scale_observation.md
+      - Reconstruction:
+          - Object 3D Reconstruction: reconstruction/object_3d_reconstruction.md
+          - Scene Components 3D Reconstruction: reconstruction/scene_components_3d_reconstruction.md
   - Concepts:
       - Camera Intrinsics Matrix: concepts/camera_intrinsics.md
   - API Reference:
       - Lifting (2D → 3D): api/lifting.md
       - Observation: api/observation.md
+      - Reconstruction: api/reconstruction.md

pyproject.toml

@@ -30,8 +30,8 @@ dependencies = [
     "fastapi>=0.115.0",
     "clean-ioc>=1.3.0",
     "pydantic>=2.0.0",
-    "grpcio>=1.60.0",
-    "grpcio-tools>=1.60.0",
+    "grpcio>=1.80.0",
+    "grpcio-tools>=1.80.0",
     "uvicorn>=0.23.0",
     "python-multipart>=0.0.6",
     "transformers>=5.6.2",
@@ -89,8 +89,21 @@ dependencies = [
 [project.optional-dependencies]
 cpu  = ["torch>=2.4.0", "torchvision>=0.19.0"]
 mps  = ["torch>=2.4.0", "torchvision>=0.19.0"]
-cuda = []
+cuda = ["onnxruntime-gpu>=1.20.0"]
 amd  = []
+reconstruction = [
+    "torch>=2.4.0",
+    "torchvision>=0.19.0",
+    "trimesh>=4.0.0",
+    "einops>=0.7.0",
+    "omegaconf>=2.3.0",
+    "PyMCubes>=0.1.6",
+    "rembg>=2.0.67",
+    "onnxruntime>=1.20.0",
+    "opencv-python>=4.8.0",
+    "basicsr>=1.4.2",
+    "realesrgan>=0.3.0",
+]
 
 [project.urls]
 Homepage = "https://github.com/OlafenwaMoses/vizion3D"

tests/assets/reconstruction_scene_480.jpg

@@ -0,0 +1,219 @@
+"""
+Integration tests — direct Python entry points for 3D reconstruction.
+
+These tests run real image bytes through the new reconstruction tasks. They use
+low-cost mesh settings so the path is exercised without the full production
+point-count cost.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+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")
+
+from vizion3d.lifting.utils import create_ply_binary  # noqa: E402
+from vizion3d.reconstruction import (  # noqa: E402
+    Object3DReconstruction,
+    Object3DReconstructionCommand,
+    Object3DReconstructionConfig,
+    SceneComponents3DReconstruction,
+    SceneComponents3DReconstructionCommand,
+    SceneComponents3DReconstructionConfig,
+)
+from vizion3d.reconstruction.handlers import Object3DReconstructionHandler  # noqa: E402
+
+OBJECT_LIMIT = float(os.environ.get("VIZION3D_TEST_RECON_OBJECT_LIMIT", "240.0"))
+SCENE_LIMIT = float(os.environ.get("VIZION3D_TEST_RECON_SCENE_LIMIT", "420.0"))
+TEST_IMAGE_MAX_DIMENSION = 480
+RECONSTRUCTION_IMAGE = "reconstruction_scene_480.jpg"
+
+
+@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)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def triposr_source_available():
+    source = (
+        Path(__file__).resolve().parents[2]
+        / "research"
+        / "3D_Object-Reconstruction"
+        / "TripoSR"
+    )
+    if not (source / "tsr" / "system.py").is_file():
+        pytest.skip(f"TripoSR source not found: {source}")
+    os.environ["VIZION3D_TRIPOSR_SOURCE"] = str(source)
+
+
+@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}")
+    image = Image.open(path)
+    assert max(image.size) <= TEST_IMAGE_MAX_DIMENSION
+    return path.read_bytes()
+
+
+def _save_mesh_and_cloud(result, run_dir: Path, stem: str) -> None:
+    run_dir.mkdir(parents=True, exist_ok=True)
+    mesh_path = run_dir / f"{stem}_mesh.ply"
+    cloud_path = run_dir / f"{stem}_point_cloud.ply"
+    mesh_path.write_bytes(result.mesh.export(file_type="ply", encoding="binary_little_endian"))
+    points = np.asarray(result.point_cloud.points).astype(np.float32)
+    colors = (np.asarray(result.point_cloud.colors) * 255).astype(np.uint8)
+    cloud_path.write_bytes(create_ply_binary(points, colors))
+
+
+def _assert_object_result(result) -> None:
+    assert result.vertex_count > 0
+    assert result.face_count > 0
+    assert result.point_count > 0
+    assert result.mesh.vertices.shape[0] == result.vertex_count
+    assert result.mesh.faces.shape[0] == result.face_count
+    assert result.point_cloud.has_points()
+    assert len(np.asarray(result.point_cloud.points)) == result.point_count
+    assert np.all(np.asarray(result.mesh.visual.vertex_colors)[:, :3] == 211)
+    assert np.allclose(np.asarray(result.point_cloud.colors), 211 / 255)
+
+
+def test_object_3d_reconstruction_runs_real_image(
+    reconstruction_image_bytes,
+    reconstruction_model_bundle,
+    tmp_path,
+    timing_collector,
+):
+    image = Image.open(BytesIO(reconstruction_image_bytes))
+    assert max(image.size) <= TEST_IMAGE_MAX_DIMENSION
+    Object3DReconstructionHandler._models.clear()
+    config = Object3DReconstructionConfig(
+        max_input_dimension=512,
+        marching_cubes_resolution=64,
+        point_count=2_048,
+        smoothing_iterations=0,
+        device=os.environ.get("VIZION3D_TEST_RECON_DEVICE", "cpu"),
+    )
+
+    t0 = time.perf_counter()
+    result = Object3DReconstruction().run(
+        Object3DReconstructionCommand(
+            image_input=reconstruction_image_bytes,
+            model_bundle=reconstruction_model_bundle,
+            advanced_config=config,
+        )
+    )
+    elapsed = time.perf_counter() - t0
+
+    _assert_object_result(result)
+    _save_mesh_and_cloud(result, tmp_path / "object_3d_reconstruction", "object")
+    timing_collector.add(
+        "Direct",
+        "TripoSR object",
+        1,
+        elapsed,
+        str(tmp_path / "object_3d_reconstruction"),
+        task="Object 3D Reconstruction",
+        model="TripoSR",
+        device=config.device,
+    )
+    assert elapsed < OBJECT_LIMIT
+
+
+def test_scene_components_3d_reconstruction_runs_real_image(
+    reconstruction_image_bytes,
+    reconstruction_model_bundle,
+    local_model_path,
+    local_annotation_model_path,
+    tmp_path,
+    timing_collector,
+):
+    image = Image.open(BytesIO(reconstruction_image_bytes))
+    assert max(image.size) <= TEST_IMAGE_MAX_DIMENSION
+    Object3DReconstructionHandler._models.clear()
+    object_config = Object3DReconstructionConfig(
+        max_input_dimension=512,
+        marching_cubes_resolution=64,
+        point_count=1_024,
+        smoothing_iterations=0,
+        device=os.environ.get("VIZION3D_TEST_RECON_DEVICE", "cpu"),
+    )
+    config = SceneComponents3DReconstructionConfig(
+        max_input_dimension=640,
+        max_objects=1,
+        confidence_threshold=0.05,
+        padding_ratio=0.1,
+        object_config=object_config,
+    )
+
+    t0 = time.perf_counter()
+    result = SceneComponents3DReconstruction().run(
+        SceneComponents3DReconstructionCommand(
+            image_input=reconstruction_image_bytes,
+            model_bundle=reconstruction_model_bundle,
+            depth_model_backend=local_model_path,
+            annotation_model_backend=local_annotation_model_path,
+            advanced_config=config,
+        )
+    )
+    elapsed = time.perf_counter() - t0
+
+    assert result.source_image_size[0] > 0 and result.source_image_size[1] > 0
+    assert max(result.source_image_size) <= TEST_IMAGE_MAX_DIMENSION
+    assert max(result.analysis_image_size) <= config.max_input_dimension
+    assert result.depth_backend_used
+    assert result.annotation_backend_used
+    assert result.reconstruction_backend_used
+    assert len(result.components) >= 1
+
+    run_dir = tmp_path / "scene_components_3d_reconstruction"
+    summary = {
+        "source_image_size": result.source_image_size,
+        "analysis_image_size": result.analysis_image_size,
+        "components": [
+            {
+                "label": component.label,
+                "class_id": component.class_id,
+                "confidence": component.confidence,
+                "vertex_count": component.vertex_count,
+                "face_count": component.face_count,
+                "point_count": component.point_count,
+            }
+            for component in result.components
+        ],
+    }
+    run_dir.mkdir(parents=True, exist_ok=True)
+    (run_dir / "summary.json").write_text(json.dumps(summary, indent=2))
+
+    for index, component in enumerate(result.components, start=1):
+        _assert_object_result(component)
+        _save_mesh_and_cloud(component, run_dir, f"component_{index:02d}")
+
+    timing_collector.add(
+        "Direct",
+        "Scene components",
+        1,
+        elapsed,
+        str(run_dir),
+        task="Scene Components 3D Reconstruction",
+        model="Depth + YOLO + RealESRGAN + TripoSR",
+        device=object_config.device,
+    )
+    assert elapsed < SCENE_LIMIT