cuhk-aim-group/endogaussian

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.