SpatialHub

A spatial AI and perception library providing unified ONNX Runtime inference adapters for computer vision and 3D spatial computing models.

Motivation & Architecture

3D spatial vision systems (such as 6D object pose estimation, Visual SLAM, and 3D reconstruction) are multi-stage pipelines composed of representation learning, geometric matching, depth prediction, proposal segmentation, and pose optimization.

In practice, integrating multiple vision models into a unified pipeline presents practical challenges:

• Interface Variations: Models differ in coordinate conventions, dictionary formats, and output tensor dimensions.
• Environment Conflicts: Combining distinct model implementations often introduces dependency conflicts and heavy runtime footprints.

SpatialHub addresses this by providing standardized dataclass returns over a unified ONNX Runtime execution layer:

Sensor & Asset Inputs                   Perception Adapters (ONNX)          Standardized Dataclasses
─────────────────────────────────      ─────────────────────────────       ─────────────────────────
Single RGB Image                  ───►  DINOv2 (Feature Extraction)   ───►  FeatureExtractionResult
Image Pair                        ───►  EfficientLoFTR (Matching)     ───►  MatchResult
RGB Images + Intrinsics (K)       ───►  Depth Anything 3 (Depth)      ───►  DepthPredictionResult
RGB Image + CAD Mesh (.ply)       ───►  FastSAM / SAM / CNOS (Masks)  ───►  SegmentationResult
RGB-D + Intrinsics (K) + CAD Mesh ───►  FoundationPose (6D Pose)      ───►  PoseEstimationResult
                                                                                    │
                                                                                    ▼
Downstream Systems (Visual SLAM, 3D Reconstruction, Robotics Manipulation)

Downstream spatial algorithms operate on these return types, allowing models to be interchanged without modifying downstream pipeline logic.

Key Principles

• Standardized Return Contracts: Dataclass structures across model families (MatchResult, DepthPredictionResult, FeatureExtractionResult, SegmentationResult, PoseEstimationResult).
• ONNX Runtime Execution: Inference executes via ONNX Runtime with NumPy and OpenCV vector operations.
• Model Weight Resolution: Automatically fetches and caches .onnx weight files from Hugging Face Hub.
• Hardware Acceleration: Supports CPU, CUDA, and TensorRT execution providers.
• Offscreen Rendering Utilities: Headless G-buffer and batched atlas renderer for CAD mesh template matching and pose estimation.

Supported Models

Installation

Requires Python 3.12+. Choose the installation for your hardware setup:

Standard Installation

For CPU inference:

pip install "spatialhub[cpu]"

For NVIDIA GPU acceleration (CUDA / TensorRT):

pip install "spatialhub[gpu]"

3D CAD & Rendering Installation

For workflows requiring 3D CAD mesh loading and template rendering (e.g. FoundationPose, CNOS):

For CPU with rendering:

pip install "spatialhub[cpu,render]"

For GPU with rendering:

pip install "spatialhub[gpu,render]"

Quickstart

from spatialhub import FoundationPose, EfficientLoFTR, DepthAnything3, DINOv2, FastSAM, SAM, CNOS

# 6D Object Pose Estimation (FoundationPose)
est = FoundationPose(
    model_path="mesh.obj",
    model_unit="mm",
    scorer_weights="scorer.onnx",
    refiner_weights="refiner.onnx",
)
pose_res = est.estimate(rgb=rgb_img, depth=depth_img, K=cam_K, mask=obj_mask)
pose_res.visualize(draw_bbox=True, draw_axes=True, save_path="pose.png")

# Feature Matching (EfficientLoFTR)
matcher = EfficientLoFTR()
match_res = matcher.match("img1.jpg", "img2.jpg", max_dim=1024)
match_res.visualize(top_k=50, save_path="matches.png")

# Depth Estimation (Depth Anything 3)
estimator = DepthAnything3(model_name="da3_base")
depth_res = estimator.estimate_depth(images=["view1.png", "view2.png"])
depth_viz = estimator.visualize(depth_res.depth[0])

# Feature Extraction (DINOv2)
dino = DINOv2(model_variant="vitl14")
feat_res = dino.extract_features("image.png", l2_normalize=True)

# Proposal Segmentation (FastSAM)
fastsam = FastSAM(model_variant="x")
seg_res = fastsam.generate_masks("scene.png", conf_threshold=0.3)
seg_res.visualize_mask(save_path="fastsam_masks.png")

ONNX Export

Export scripts for generating .onnx models from source repositories are located under tools/export/:

uv run tools/export/export_efficient_loftr.py --checkpoint weights/eloftr_outdoor.ckpt --output-folder onnx_weight

See the ONNX Export Guide for options and details.

License

Core SpatialHub code is released under the Apache 2.0 License. Pretrained model weights and submodule architectures maintain their respective original licenses.