# horizonrobotics/uni3r architecture

> Architecture diagram of the GitHub repository horizonrobotics/uni3r, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-05-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/horizonrobotics/uni3r
- Repository: https://github.com/horizonrobotics/uni3r

## Overview

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/uni3r.py`](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/uni3r.py) (Uni3R core), [`uni3r/datasets`](https://github.com/horizonrobotics/uni3r/tree/main/uni3r/datasets) (Datasets), [`train.py`](https://github.com/horizonrobotics/uni3r/blob/main/train.py) (train.py), [`test.py`](https://github.com/horizonrobotics/uni3r/blob/main/test.py) (test.py)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_group_ingest["Data flow"]
  node_data_process["Data process<br/>preprocess"]
  node_scannet_export["ScanNet export<br/>export pipeline<br/>[export.py]"]
  node_scannet_proc["ScanNet proc<br/>processor"]
  node_scannetpp_preproc["ScanNet++ proc<br/>processor"]
  node_datasets["Datasets<br/>loaders"]
  node_scannet_ds["ScanNet ds<br/>dataset loader<br/>[scannet.py]"]
  node_scannetpp_ds["ScanNet++ ds<br/>dataset loader<br/>[scannetpp.py]"]
  node_testdata_ds["Test data<br/>dataset loader<br/>[testdata.py]"]
end

subgraph group_group_model["Uni3R model"]
  node_train_py["train.py<br/>entrypoint<br/>[train.py]"]
  node_test_py["test.py<br/>entrypoint<br/>[test.py]"]
  node_scripts["Scripts<br/>shell wrappers"]
  node_default_cfg["Default cfg<br/>config<br/>[default.yaml]"]
  node_uni3r_core["Uni3R core<br/>model composition<br/>[uni3r.py]"]
  node_dust3r_feat["DUSt3R feat<br/>geometry backbone"]
  node_depth_head["Depth head<br/>prediction head<br/>[dpt_head.py]"]
  node_ptv3["Point aggregator<br/>point encoder<br/>[ptv3.py]"]
  node_vggt["VGGT branch<br/>feature branch<br/>[vggt_w_intr_emb.py]"]
  node_lseg["LSeg<br/>semantic branch<br/>[lseg.py]"]
  node_loss["Losses<br/>objective<br/>[loss.py]"]
end

subgraph group_group_render["Gaussian rendering"]
  node_gaussian_head["Gaussian head<br/>scene head<br/>[gaussian_head.py]"]
  node_gaussian_model["Gauss model<br/>scene model<br/>[gaussian_model.py]"]
  node_cuda_splat["CUDA splat<br/>renderer utils<br/>[cuda_splatting.py]"]
  node_camera_utils["Camera utils<br/>pose utils<br/>[camera_utils.py]"]
  node_umeyama["Alignment<br/>[umeyama.py]"]
  node_viz_utils["Viz utils<br/>visualization"]
end

subgraph group_group_external["External deps"]
  node_dust3r_sub["DUSt3R<br/>external submodule"]
  node_diff_gauss["Diff Gaussian<br/>external submodule"]
  node_lang_seg_sub["LangSeg<br/>external submodule"]
  node_vggt_sub["VGGT<br/>external submodule"]
end

node_scripts -->|"launches"| node_train_py
node_scripts -->|"launches"| node_test_py
node_default_cfg -->|"configures"| node_train_py
node_default_cfg -->|"configures"| node_test_py
node_data_process -->|"processes"| node_scannet_proc
node_data_process -->|"processes"| node_scannetpp_preproc
node_scannet_proc -->|"exports"| node_scannet_export
node_datasets -->|"loads"| node_scannet_ds
node_datasets -->|"loads"| node_scannetpp_ds
node_datasets -->|"loads"| node_testdata_ds
node_train_py -->|"consumes"| node_datasets
node_test_py -->|"consumes"| node_datasets
node_train_py -->|"builds"| node_uni3r_core
node_test_py -->|"runs"| node_uni3r_core
node_uni3r_core -->|"uses"| node_dust3r_feat
node_uni3r_core -->|"uses"| node_depth_head
node_uni3r_core -->|"uses"| node_ptv3
node_uni3r_core -->|"uses"| node_vggt
node_uni3r_core -->|"uses"| node_lseg
node_uni3r_core -->|"produces"| node_gaussian_head
node_gaussian_head -->|"updates"| node_gaussian_model
node_gaussian_model -->|"renders with"| node_cuda_splat
node_gaussian_model -->|"backed by"| node_diff_gauss
node_uni3r_core -->|"optimized by"| node_loss
node_loss -->|"uses"| node_camera_utils
node_loss -->|"uses"| node_umeyama
node_viz_utils -->|"inspects"| node_gaussian_model
node_dust3r_feat -->|"depends on"| node_dust3r_sub
node_lseg -->|"depends on"| node_lang_seg_sub
node_vggt -->|"depends on"| node_vggt_sub

click node_train_py "https://github.com/horizonrobotics/uni3r/blob/main/train.py"
click node_test_py "https://github.com/horizonrobotics/uni3r/blob/main/test.py"
click node_scripts "https://github.com/horizonrobotics/uni3r/tree/main/scripts"
click node_default_cfg "https://github.com/horizonrobotics/uni3r/blob/main/configs/default.yaml"
click node_data_process "https://github.com/horizonrobotics/uni3r/tree/main/data_process"
click node_scannet_export "https://github.com/horizonrobotics/uni3r/blob/main/data_process/scannet/export_data/export.py"
click node_scannet_proc "https://github.com/horizonrobotics/uni3r/blob/main/data_process/scannet/scannet_processor.py"
click node_scannetpp_preproc "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets_preprocess/scannetpp_preprocess.py"
click node_datasets "https://github.com/horizonrobotics/uni3r/tree/main/uni3r/datasets"
click node_scannet_ds "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets/scannet.py"
click node_scannetpp_ds "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets/scannetpp.py"
click node_testdata_ds "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets/testdata.py"
click node_uni3r_core "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/uni3r.py"
click node_dust3r_feat "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/dust3r_with_feature.py"
click node_depth_head "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/dpt_head.py"
click node_ptv3 "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/ptv3.py"
click node_vggt "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/vggt_w_intr_emb.py"
click node_lseg "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/lseg.py"
click node_gaussian_head "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/gaussian_head.py"
click node_gaussian_model "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/gaussian_model.py"
click node_cuda_splat "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/cuda_splatting.py"
click node_loss "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/loss.py"
click node_camera_utils "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/camera_utils.py"
click node_umeyama "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/umeyama.py"
click node_viz_utils "https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/visualization_utils.py"
click node_dust3r_sub "https://github.com/horizonrobotics/uni3r/tree/main/submodules/dust3r"
click node_diff_gauss "https://github.com/horizonrobotics/uni3r/tree/main/submodules/diff-gaussian-rasterization"
click node_lang_seg_sub "https://github.com/horizonrobotics/uni3r/tree/main/submodules/lang_seg"
click node_vggt_sub "https://github.com/horizonrobotics/uni3r/tree/main/submodules/vggt"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_data_process,node_scannet_export,node_scannet_proc,node_scannetpp_preproc,node_datasets,node_scannet_ds,node_scannetpp_ds,node_testdata_ds toneBlue
class node_train_py,node_test_py,node_scripts,node_default_cfg,node_uni3r_core,node_dust3r_feat,node_depth_head,node_ptv3,node_vggt,node_lseg,node_loss toneAmber
class node_gaussian_head,node_gaussian_model,node_cuda_splat,node_camera_utils,node_umeyama,node_viz_utils toneMint
class node_dust3r_sub,node_diff_gauss,node_lang_seg_sub,node_vggt_sub toneRose
```

## Components

### Data flow

Raw scene data preparation and multi-view dataset loading.

- [Data process](https://github.com/horizonrobotics/uni3r/tree/main/data_process): preprocess · Dataset conversion and export boundary before training. · `data_process`
- [ScanNet export](https://github.com/horizonrobotics/uni3r/blob/main/data_process/scannet/export_data/export.py): export pipeline · Raw ScanNet sensor-data export into training-ready structure. · `data_process/scannet/export_data/export.py`
- [ScanNet proc](https://github.com/horizonrobotics/uni3r/blob/main/data_process/scannet/scannet_processor.py): processor · ScanNet-specific preprocessing orchestration. · `data_process/scannet/scannet_processor.py`
- [ScanNet++ proc](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets_preprocess/scannetpp_preprocess.py): processor · ScanNet++ preprocessing support. · `uni3r/datasets_preprocess/scannetpp_preprocess.py`
- [Datasets](https://github.com/horizonrobotics/uni3r/tree/main/uni3r/datasets): loaders · Repo-owned multi-view dataset interfaces. · `uni3r/datasets`
- [ScanNet ds](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets/scannet.py): dataset loader · Loads ScanNet multi-view training and inference samples. · `uni3r/datasets/scannet.py`
- [ScanNet++ ds](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets/scannetpp.py): dataset loader · Loads ScanNet++ multi-view samples. · `uni3r/datasets/scannetpp.py`
- [Test data](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/datasets/testdata.py): dataset loader · Lightweight evaluation/demo sample loader. · `uni3r/datasets/testdata.py`

### Uni3R model

Core reconstruction and semantic prediction modules.

- [train.py](https://github.com/horizonrobotics/uni3r/blob/main/train.py): entrypoint · Main training entry point. · `train.py`
- [test.py](https://github.com/horizonrobotics/uni3r/blob/main/test.py): entrypoint · Main evaluation entry point. · `test.py`
- [Scripts](https://github.com/horizonrobotics/uni3r/tree/main/scripts): shell wrappers · Convenience launchers for train/test across view counts. · `scripts`
- [Default cfg](https://github.com/horizonrobotics/uni3r/blob/main/configs/default.yaml): config · Top-level experiment settings and paths. · `configs/default.yaml`
- [Uni3R core](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/uni3r.py): model composition · Central orchestration of geometry, semantics, and 3D reconstruction. · `uni3r/uni3r.py`
- [DUSt3R feat](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/dust3r_with_feature.py): geometry backbone · Multi-view geometric prior and feature backbone wrapper. · `uni3r/dust3r_with_feature.py`
- [Depth head](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/dpt_head.py): prediction head · Dense geometric prediction head. · `uni3r/dpt_head.py`
- [Point aggregator](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/ptv3.py): point encoder · Point/feature aggregation branch. · `uni3r/ptv3.py`
- [VGGT branch](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/vggt_w_intr_emb.py): feature branch · Alternate feature/intrinsic embedding branch. · `uni3r/vggt_w_intr_emb.py`
- [LSeg](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/lseg.py): semantic branch · Language-guided semantic segmentation integration. · `uni3r/lseg.py`
- [Losses](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/loss.py): objective · Training objectives for geometry, reconstruction, and semantics. · `uni3r/loss.py`

### Gaussian rendering

Gaussian scene representation, splatting, and optimization utilities.

- [Gaussian head](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/gaussian_head.py): scene head · Converts fused scene features into Gaussian representation. · `uni3r/gaussian_head.py`
- [Gauss model](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/gaussian_model.py): scene model · Gaussian scene storage and update utilities. · `uni3r/utils/gaussian_model.py`
- [CUDA splat](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/cuda_splatting.py): renderer utils · CUDA-enabled splatting utilities. · `uni3r/utils/cuda_splatting.py`
- [Camera utils](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/camera_utils.py): pose utils · Camera and pose handling helpers. · `uni3r/utils/camera_utils.py`
- [Alignment](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/umeyama.py): pose alignment · Rigid alignment utility used in pose/scene normalization. · `uni3r/utils/umeyama.py`
- [Viz utils](https://github.com/horizonrobotics/uni3r/blob/main/uni3r/utils/visualization_utils.py): visualization · Point-cloud and scene visualization utilities. · `uni3r/utils/visualization_utils.py`

### External deps

Vendored research submodules used as core dependencies.

- [DUSt3R](https://github.com/horizonrobotics/uni3r/tree/main/submodules/dust3r): external submodule · Vendored multi-view reconstruction backbone and training stack. · `submodules/dust3r`
- [Diff Gaussian](https://github.com/horizonrobotics/uni3r/tree/main/submodules/diff-gaussian-rasterization): external submodule · Differentiable Gaussian rasterization backend. · `submodules/diff-gaussian-rasterization`
- [LangSeg](https://github.com/horizonrobotics/uni3r/tree/main/submodules/lang_seg): external submodule · Vendored language/semantic segmentation dependency. · `submodules/lang_seg`
- [VGGT](https://github.com/horizonrobotics/uni3r/tree/main/submodules/vggt): external submodule · Vendored feature/intrinsic embedding dependency. · `submodules/vggt`

## Connections

- 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.