How horizonrobotics/uni3r works
1. What this repo is
Uni3R is a 3D reconstruction + semantic understanding system for unposed multi-view images. The core output is a Gaussian-splatting-based 3D scene representation enriched with semantics, trained and evaluated mainly on ScanNet and ScanNet++.
2. Top-level execution path
The main runtime entry points are train.py and test.py, with shell wrappers in scripts/ for different view counts: 2-view, 8-view, and 16-view training/evaluation. configs/default.yaml and script arguments control dataset paths, output directories, and experiment settings.
3. Data preparation boundary
Raw dataset conversion is separated from model training. The data_process/ tree contains dataset-specific processors, especially for ScanNet and ScanNet++. data_process/scannet/export_data/ is an adapted export pipeline for ScanNet sensor data, turning raw scene assets into training-ready structures.
4. Dataset layer
The repo’s own dataset interfaces live under uni3r/datasets/, with loaders for:
- scannet.py
- scannetpp.py
- testdata.py
These loaders feed multi-view samples into training/inference, likely bundling images with camera metadata and scene-level grouping.
5. Model orchestration
uni3r/uni3r.py is the central model composition layer. It ties together geometry extraction, feature fusion, semantic prediction, and Gaussian-splatting reconstruction. uni3r/configs.py and uni3r/utils/weight_modify.py support experiment configuration and checkpoint adaptation.
6. Geometry and feature backbone
The architecture reuses DUSt3R from submodules/dust3r/ as a major geometric prior and multi-view feature backbone. In the Uni3R wrapper, uni3r/dust3r_with_feature.py, uni3r/dpt_head.py, uni3r/ptv3.py, and uni3r/vggt_w_intr_emb.py suggest separate branches for dense visual features, depth/geometry heads, and point/feature aggregation.
7. Semantic understanding branch
uni3r/lseg.py integrates LSeg for language-driven semantic segmentation. This is the semantic side of the system, used to attach class/semantic information to the reconstructed 3D representation rather than producing geometry alone.
8. Gaussian splatting reconstruction
uni3r/gaussian_head.py and uni3r/utils/gaussian_model.py convert fused scene features into a Gaussian scene model. Rendering and optimization rely on CUDA-enabled splatting utilities such as uni3r/utils/cuda_splatting.py and the external submodules/diff-gaussian-rasterization package. The repo also depends on gsplat and pytorch3d as geometry/rendering infrastructure.
9. Training objective and optimization
uni3r/loss.py defines the training losses, including the geometry loss mentioned in the README. The training loop is organized to optimize reconstruction quality, geometric consistency, and semantic alignment over multi-view samples. Supporting utilities like uni3r/utils/umeyama.py and uni3r/utils/camera_utils.py indicate explicit pose/alignment handling in the pipeline.
10. Visualization and scene utilities
uni3r/utils/visualization_utils.py, graphics_utils.py, points_process.py, and sh_utils.py support point-cloud handling, visualization, and rendering-related math. These are not primary model components, but they sit on the critical path for debugging, evaluation, and producing outputs.
11. External submodule dependencies
The repository is heavily built around vendored research code:
- submodules/dust3r/: multi-view reconstruction backbone, training/inference, datasets, and CroCo-based encoder components
- submodules/diff-gaussian-rasterization: differentiable Gaussian rasterizer
- Additional placeholders for lang_seg and vggt
These submodules are not just libraries; they are core architectural dependencies for geometry and rendering.
12. End-to-end data flow
The main flow is: raw dataset images + camera data → dataset preprocessing (data_process/, uni3r/datasets_preprocess/) → training dataset loaders (uni3r/datasets/) → geometry/feature backbone (dust3r-based modules) → semantic branch (LSeg) + Gaussian head → differentiable Gaussian rendering → losses/optimization → checkpoints and evaluation outputs.
13. Operational outputs
Trained models are saved under the configured output directory, while pretrained checkpoints are expected in checkpoints/pretrained_models/. The README also points to HuggingFace-hosted Uni3R checkpoints and a separate demo checkpoint for LSeg-related components.
14. What to draw in an architecture diagram
A good diagram should show three major subsystems feeding one another:
- Data ingestion/preprocessing
- Multi-view reconstruction + semantic model
- Differentiable Gaussian rendering + losses
Around those, show the external dependencies: DUSt3R, diff-gaussian-rasterization, LSeg, Pytorch3D, and gsplat.
Main entry points: uni3r.py Uni3R core · datasets Datasets · train.py train.py · test.py test.py
Components 29
Data flow
Raw scene data preparation and multi-view dataset loading.
- Data process · Dataset conversion and export boundary before training.
data_process - ScanNet export · Raw ScanNet sensor-data export into training-ready structure.
data_process/scannet/export_data/export.py - ScanNet proc · ScanNet-specific preprocessing orchestration.
data_process/scannet/scannet_processor.py - ScanNet++ proc · ScanNet++ preprocessing support.
uni3r/datasets_preprocess/scannetpp_preprocess.py - Datasets · Repo-owned multi-view dataset interfaces.
uni3r/datasets - ScanNet ds · Loads ScanNet multi-view training and inference samples.
uni3r/datasets/scannet.py - ScanNet++ ds · Loads ScanNet++ multi-view samples.
uni3r/datasets/scannetpp.py - Test data · Lightweight evaluation/demo sample loader.
uni3r/datasets/testdata.py
Uni3R model
Core reconstruction and semantic prediction modules.
- train.py · Main training entry point.
train.py - test.py · Main evaluation entry point.
test.py - Scripts · Convenience launchers for train/test across view counts.
scripts - Default cfg · Top-level experiment settings and paths.
configs/default.yaml - Uni3R core · Central orchestration of geometry, semantics, and 3D reconstruction.
uni3r/uni3r.py - DUSt3R feat · Multi-view geometric prior and feature backbone wrapper.
uni3r/dust3r_with_feature.py - Depth head · Dense geometric prediction head.
uni3r/dpt_head.py - Point aggregator · Point/feature aggregation branch.
uni3r/ptv3.py - VGGT branch · Alternate feature/intrinsic embedding branch.
uni3r/vggt_w_intr_emb.py - LSeg · Language-guided semantic segmentation integration.
uni3r/lseg.py - Losses · Training objectives for geometry, reconstruction, and semantics.
uni3r/loss.py
Gaussian rendering
Gaussian scene representation, splatting, and optimization utilities.
- Gaussian head · Converts fused scene features into Gaussian representation.
uni3r/gaussian_head.py - Gauss model · Gaussian scene storage and update utilities.
uni3r/utils/gaussian_model.py - CUDA splat · CUDA-enabled splatting utilities.
uni3r/utils/cuda_splatting.py - Camera utils · Camera and pose handling helpers.
uni3r/utils/camera_utils.py - Alignment · Rigid alignment utility used in pose/scene normalization.
uni3r/utils/umeyama.py - Viz utils · Point-cloud and scene visualization utilities.
uni3r/utils/visualization_utils.py
External deps
Vendored research submodules used as core dependencies.
- DUSt3R · Vendored multi-view reconstruction backbone and training stack.
submodules/dust3r - Diff Gaussian · Differentiable Gaussian rasterization backend.
submodules/diff-gaussian-rasterization - LangSeg · Vendored language/semantic segmentation dependency.
submodules/lang_seg - VGGT · Vendored feature/intrinsic embedding dependency.
submodules/vggt
Connections 30
- Scripts → train.py: launches · Shell wrappers invoke training with view-count presets.
- Scripts → test.py: launches · Shell wrappers invoke evaluation with view-count presets.
- Default cfg → train.py: configures · Default config supplies experiment and dataset settings.
- Default cfg → test.py: configures · Default config supplies evaluation settings.
- Data process → ScanNet proc: processes · Repository preprocessing boundary for ScanNet.
- Data process → ScanNet++ proc: processes · Repository preprocessing boundary for ScanNet++.
- ScanNet proc → ScanNet export: exports · ScanNet processor uses export pipeline to build training data.
- Datasets → ScanNet ds: loads · Dataset package exposes ScanNet loader.
- Datasets → ScanNet++ ds: loads · Dataset package exposes ScanNet++ loader.
- Datasets → Test data: loads · Dataset package exposes test/demo loader.
- train.py → Datasets: consumes · Training consumes multi-view samples from dataset loaders.
- test.py → Datasets: consumes · Evaluation consumes dataset samples.
- train.py → Uni3R core: builds · Training instantiates the Uni3R model.
- test.py → Uni3R core: runs · Evaluation runs the Uni3R model for inference.
- Uni3R core → DUSt3R feat: uses · Core model wraps the DUSt3R-based geometric backbone.
- Uni3R core → Depth head: uses · Core model attaches dense geometry head.
- Uni3R core → Point aggregator: uses · Core model includes point aggregation branch.
- Uni3R core → VGGT branch: uses · Core model includes VGGT/intrinsic embedding branch.
- Uni3R core → LSeg: uses · Core model integrates semantic segmentation branch.
- Uni3R core → Gaussian head: produces · Core model feeds fused features into Gaussian scene head.
- Gaussian head → Gauss model: updates · Head parameterizes the Gaussian scene model.
- Gauss model → CUDA splat: renders with · Gaussian model relies on CUDA splatting utilities.
- Gauss model → Diff Gaussian: backed by · Differentiable rasterization backend for training/rendering.
- Uni3R core → Losses: optimized by · Training objectives supervise the full reconstruction pipeline.
- Losses → Camera utils: uses · Losses and supervision rely on camera/pose handling.
- Losses → Alignment: uses · Losses use alignment for pose/scene consistency.
- Viz utils → Gauss model: inspects · Visualization utilities operate on Gaussian scene outputs.
- DUSt3R feat → DUSt3R: depends on · Geometry backbone is built on vendored DUSt3R code.
- LSeg → LangSeg: depends on · Semantic branch is backed by language segmentation submodule.
- VGGT branch → VGGT: depends on · Feature branch is backed by vendored VGGT code.