# cuhk-aim-group/endogaussian architecture

> Architecture diagram of the GitHub repository cuhk-aim-group/endogaussian, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-19). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/cuhk-aim-group/endogaussian
- Repository: https://github.com/cuhk-aim-group/endogaussian

## Overview

1. **What this repo is**
   This is a PyTorch-based 4D Gaussian splatting system for dynamic endoscopic scene reconstruction. The core idea is: ingest a time-ordered endoscopic dataset, optimize a deformable Gaussian scene representation, then render novel views/frames and evaluate image quality.

2. **Main execution entry points**
   The root-level workflow is centered on `train.py`, `render.py`, `metrics.py`, and `full_eval.py`. These are the primary orchestration scripts for training, rendering checkpoints, computing metrics, and running larger evaluation sweeps. The `journal/` tree mirrors the same structure for the later TMI version.

3. **Configuration is dataset-scene driven**
   The `arguments/` package contains scene-specific config modules for EndoNeRF and SCARED, and `journal/arguments/` adds Hamlyn. These files encode per-sequence hyperparameters and dataset paths, so the architecture is not “one model fits all”; it is configured per clip/sequence from Python config modules.

4. **Data ingestion boundary**
   Raw datasets are expected under `data/...` in a preprocessed layout. The `scene/` package is the dataset layer: `dataset.py`, `dataset_readers.py`, `endo_loader.py`, and `colmap_loader.py` turn endoscopic or COLMAP-style assets into camera objects, image streams, and scene metadata. This is the boundary between filesystem datasets and the learning system.

5. **Core scene representation**
   The central model lives in `scene/gaussian_model.py`. Supporting modules such as `deformation.py`, `hexplane.py`, and `regulation.py` define the temporal deformation field and regularization terms that make the scene dynamic rather than static. In architecture terms, this is the “world model” the optimizer updates.

6. **Camera and geometry utilities**
   `scene/cameras.py` plus the `utils/` helpers provide camera parameter handling, image/graphics math, spherical-harmonics utilities, scene normalization, and general experiment utilities. These modules are shared infrastructure used by both training and rendering.

7. **Rendering pipeline**
   Rendering is handled through `gaussian_renderer/`, which wraps the differentiable splatting backend and includes `network_gui.py` for live/interactive visualization support. The renderer consumes the optimized Gaussian scene plus camera parameters and outputs images, depth, and other render products used by training and evaluation.

8. **Optimization and loss computation**
   Training is not just reconstruction loss; it also uses utility losses from `utils/loss_utils.py` and regularization from `scene/regulation.py`. The workflow is typical: load a batch of cameras/frames, render through the Gaussian renderer, compare against ground truth, and backpropagate into the Gaussian/deformation parameters.

9. **Evaluation and reconstruction outputs**
   `metrics.py` computes quantitative results, and in the journal version `journal/metrics.py` is paired with `journal/lpipsPyTorch/` for perceptual LPIPS evaluation. `render.py` and `journal/render.py` support offline rendering of train/test/video splits, and the journal version also mentions optional point-cloud reconstruction via an extra flag.

10. **Custom CUDA extensions are part of the architecture**
   Two compiled submodules are essential: `submodules/depth-diff-gaussian-rasterization` provides the differentiable Gaussian rasterizer with depth support, and `submodules/simple-knn` provides fast KNN/spatial queries. These are not optional helpers; they are the performance-critical backend that makes the model train and render in real time.

11. **Repository structure has two code generations**
   The root-level code appears to be the EndoGaussian implementation tied to the arXiv release, while `journal/` is a fuller fork for the TMI paper. The journal tree adds Hamlyn support, LPIPS, extraction utilities, and mirrored scene/renderer/training code, but the high-level architecture stays the same.

12. **External datasets and preprocessing**
   The system depends on external medical datasets rather than a built-in synthetic generator: EndoNeRF, SCARED, and Hamlyn. The repository assumes those datasets are preprocessed into the expected directory layout before training, so data preparation is a separate upstream stage in the pipeline.

13. **Outputs and experiment organization**
   Trained models and rendered artifacts are written under per-experiment output directories such as `output/endonerf/pulling`. This makes the boundary between “experiment definition” and “experiment results” very explicit, and it is what the evaluation and rendering scripts consume later.

14. **Where to draw the architecture diagram**
   A good diagram would show: dataset inputs on the left; config modules selecting a scene; the scene/data loaders creating cameras and frames; the Gaussian/deformation model in the center; the custom CUDA rasterizer underneath; and training/rendering/evaluation scripts on the right consuming the same model. The `journal/` tree can be shown as a parallel, expanded implementation of the same pipeline.

Main entry points: [`scene`](https://github.com/cuhk-aim-group/endogaussian/tree/master/scene) (Scene layer), [`train.py`](https://github.com/cuhk-aim-group/endogaussian/blob/master/train.py) (Train), [`render.py`](https://github.com/cuhk-aim-group/endogaussian/blob/master/render.py) (Render), [`journal/train.py`](https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/train.py) (Train)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_root["Root pipeline"]
  node_root_args["Configs"]
  node_root_train["Train<br/>orchestrator<br/>[train.py]"]
  node_root_render["Render<br/>orchestrator<br/>[render.py]"]
  node_root_metrics["Metrics<br/>evaluation<br/>[metrics.py]"]
  node_root_eval["Full eval<br/>sweep runner<br/>[full_eval.py]"]
  node_root_scene["Scene layer<br/>data+scene"]
  node_root_renderer["Renderer"]
  node_root_utils["Utils"]
end

subgraph group_journal["Journal pipeline"]
  node_jr_args["Configs"]
  node_jr_train["Train<br/>orchestrator<br/>[train.py]"]
  node_jr_render["Render<br/>orchestrator<br/>[render.py]"]
  node_jr_metrics["Metrics<br/>evaluation<br/>[metrics.py]"]
  node_jr_eval["Full eval<br/>sweep runner<br/>[full_eval.py]"]
  node_jr_scene["Scene layer<br/>data+scene"]
  node_jr_renderer["Renderer"]
  node_jr_utils["Utils"]
  node_jr_lpips["LPIPS<br/>perceptual metric"]
  node_jr_extract["Extract<br/>[extract.py]"]
end

subgraph group_backend["Native backend"]
  node_raster["Rasterizer<br/>cuda extension"]
  node_knn["KNN<br/>cuda extension"]
end

subgraph group_data["External data"]
  node_root_data[("Datasets<br/>external inputs")]
  node_root_out[("Outputs<br/>experiment artifacts")]
end

node_root_args -->|"selects"| node_root_train
node_root_args -->|"selects"| node_root_render
node_root_data -->|"loads"| node_root_scene
node_root_scene -->|"feeds"| node_root_train
node_root_scene -->|"feeds"| node_root_render
node_root_scene -->|"models"| node_root_renderer
node_root_renderer -->|"uses"| node_raster
node_root_scene -->|"uses"| node_knn
node_root_utils -->|"supports"| node_root_train
node_root_utils -->|"supports"| node_root_render
node_root_renderer -->|"renders"| node_root_train
node_root_train -->|"writes"| node_root_out
node_root_render -->|"writes"| node_root_out
node_root_metrics -->|"reads"| node_root_out
node_root_eval -->|"drives"| node_root_metrics
node_jr_args -->|"selects"| node_jr_train
node_jr_scene -->|"feeds"| node_jr_train
node_jr_scene -->|"models"| node_jr_renderer
node_jr_renderer -->|"uses"| node_raster
node_jr_utils -->|"supports"| node_jr_train
node_jr_renderer -->|"renders"| node_jr_train
node_jr_lpips -->|"extends"| node_jr_metrics
node_jr_render -->|"writes"| node_root_out
node_jr_extract -->|"exports"| node_root_out

click node_root_args "https://github.com/cuhk-aim-group/endogaussian/tree/master/arguments"
click node_root_train "https://github.com/cuhk-aim-group/endogaussian/blob/master/train.py"
click node_root_render "https://github.com/cuhk-aim-group/endogaussian/blob/master/render.py"
click node_root_metrics "https://github.com/cuhk-aim-group/endogaussian/blob/master/metrics.py"
click node_root_eval "https://github.com/cuhk-aim-group/endogaussian/blob/master/full_eval.py"
click node_root_scene "https://github.com/cuhk-aim-group/endogaussian/tree/master/scene"
click node_root_renderer "https://github.com/cuhk-aim-group/endogaussian/tree/master/gaussian_renderer"
click node_root_utils "https://github.com/cuhk-aim-group/endogaussian/tree/master/utils"
click node_jr_args "https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/arguments"
click node_jr_train "https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/train.py"
click node_jr_render "https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/render.py"
click node_jr_metrics "https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/metrics.py"
click node_jr_eval "https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/full_eval.py"
click node_jr_scene "https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/scene"
click node_jr_renderer "https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/gaussian_renderer"
click node_jr_utils "https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/utils"
click node_jr_lpips "https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/lpipsPyTorch"
click node_jr_extract "https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/extract.py"
click node_raster "https://github.com/cuhk-aim-group/endogaussian/tree/master/submodules/depth-diff-gaussian-rasterization"
click node_knn "https://github.com/cuhk-aim-group/endogaussian/tree/master/submodules/simple-knn"

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_root_args,node_root_train,node_root_render,node_root_metrics,node_root_eval,node_root_scene,node_root_renderer,node_root_utils toneBlue
class node_jr_args,node_jr_train,node_jr_render,node_jr_metrics,node_jr_eval,node_jr_scene,node_jr_renderer,node_jr_utils,node_jr_lpips,node_jr_extract toneAmber
class node_raster,node_knn toneMint
class node_root_data,node_root_out toneRose
```

## Components

### Root pipeline

Primary EndoGaussian implementation

- [Configs](https://github.com/cuhk-aim-group/endogaussian/tree/master/arguments): scene configs · Scene-specific Python argument modules for EndoNeRF and SCARED · `arguments`
- [Train](https://github.com/cuhk-aim-group/endogaussian/blob/master/train.py): orchestrator · Main optimization entry point · `train.py`
- [Render](https://github.com/cuhk-aim-group/endogaussian/blob/master/render.py): orchestrator · Offline rendering entry point · `render.py`
- [Metrics](https://github.com/cuhk-aim-group/endogaussian/blob/master/metrics.py): evaluation · Quantitative evaluation entry point · `metrics.py`
- [Full eval](https://github.com/cuhk-aim-group/endogaussian/blob/master/full_eval.py): sweep runner · Batch evaluation over scenes and checkpoints · `full_eval.py`
- [Scene layer](https://github.com/cuhk-aim-group/endogaussian/tree/master/scene): data+scene · Dataset loading, cameras, deformation, Gaussian model, and regularization · `scene`
- [Renderer](https://github.com/cuhk-aim-group/endogaussian/tree/master/gaussian_renderer): differentiable renderer · Gaussian splatting renderer and live GUI support · `gaussian_renderer`
- [Utils](https://github.com/cuhk-aim-group/endogaussian/tree/master/utils): math+training utils · Shared camera, geometry, loss, SH, scene, and system helpers · `utils`

### Journal pipeline

Expanded TMI fork with mirrored architecture

- [Configs](https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/arguments): scene configs · Journal scene configs including EndoNeRF, SCARED, and Hamlyn · `journal/arguments`
- [Train](https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/train.py): orchestrator · Journal training entry point · `journal/train.py`
- [Render](https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/render.py): orchestrator · Journal offline rendering entry point · `journal/render.py`
- [Metrics](https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/metrics.py): evaluation · Journal quantitative evaluation with LPIPS support · `journal/metrics.py`
- [Full eval](https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/full_eval.py): sweep runner · Journal batch evaluation entry point · `journal/full_eval.py`
- [Scene layer](https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/scene): data+scene · Journal copy of dataset, camera, deformation, and Gaussian model stack · `journal/scene`
- [Renderer](https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/gaussian_renderer): differentiable renderer · Journal renderer and live GUI support · `journal/gaussian_renderer`
- [Utils](https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/utils): math+training utils · Journal helpers for cameras, graphics, losses, SH, params, and timing · `journal/utils`
- [LPIPS](https://github.com/cuhk-aim-group/endogaussian/tree/master/journal/lpipsPyTorch): perceptual metric · LPIPS evaluation module used in journal metrics · `journal/lpipsPyTorch`
- [Extract](https://github.com/cuhk-aim-group/endogaussian/blob/master/journal/extract.py): utility · Journal extraction/export helper · `journal/extract.py`

### Native backend

Performance-critical CUDA extensions

- [Rasterizer](https://github.com/cuhk-aim-group/endogaussian/tree/master/submodules/depth-diff-gaussian-rasterization): cuda extension · Differentiable depth-aware Gaussian rasterization backend · `submodules/depth-diff-gaussian-rasterization`
- [KNN](https://github.com/cuhk-aim-group/endogaussian/tree/master/submodules/simple-knn): cuda extension · Fast spatial nearest-neighbor queries · `submodules/simple-knn`

### External data

Preprocessed medical datasets and experiment outputs

- **Datasets**: external inputs · Preprocessed EndoNeRF, SCARED, and related experiment data
- **Outputs**: experiment artifacts · Checkpoints, renders, and evaluation outputs

## Connections

- Configs → Train: selects · Scene configs parameterize training runs
- Configs → Render: selects · Configs define render-time scene settings
- Datasets → Scene layer: loads · Preprocessed datasets are ingested by the scene layer
- Scene layer → Train: feeds · Scenes provide cameras, frames, and metadata to training
- Scene layer → Render: feeds · Scene objects drive offline rendering
- Scene layer → Renderer: models · Gaussian model and deformation fields are consumed by the renderer
- Renderer → Rasterizer: uses · Renderer delegates splatting to the CUDA rasterizer
- Scene layer → KNN: uses · Spatial queries support scene and model operations
- Utils → Train: supports · Losses, geometry, and system helpers support optimization
- Utils → Render: supports · Shared utilities support rendering and output handling
- Renderer → Train: renders · Training optimizes through rendered outputs
- Train → Outputs: writes · Checkpoints and experiment artifacts are saved per run
- Render → Outputs: writes · Rendered frames and videos are exported
- Metrics → Outputs: reads · Metrics consume rendered outputs and checkpoints
- Full eval → Metrics: drives · Full evaluation orchestrates metric computation
- Configs → Train: selects · Journal configs parameterize training
- Scene layer → Train: feeds · Journal scene layer supplies data and models
- Scene layer → Renderer: models · Journal renderer consumes the same deformable Gaussian world model
- Renderer → Rasterizer: uses · Journal renderer uses the shared CUDA backend
- Utils → Train: supports · Journal training uses shared helper utilities
- Renderer → Train: renders · Journal optimization is driven by rendered outputs
- LPIPS → Metrics: extends · LPIPS is part of journal metric computation
- Render → Outputs: writes · Journal render outputs are saved alongside experiments
- Extract → Outputs: exports · Extraction utility produces additional artifacts