# neuropoly/totalspineseg architecture

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

- Interactive diagram: https://gitdiagram.com/neuropoly/totalspineseg
- Repository: https://github.com/neuropoly/totalspineseg

## Overview

TotalSpineSeg is a Python medical-imaging tool for segmenting and labeling spinal structures in MRI, with a two-model nnU-Net inference workflow and optional localizer-guided labeling. The documented workflow predicts semantic regions and landmarks in step 1, derives odd-disc guidance, predicts again in step 2, then assigns individual vertebra and disc labels. It also supports single-model labeling, training, dataset preparation, weight installation, and preview generation. Inference and utility excerpts are partial, so uncertain internal calls are omitted; documented but unsampled training components use tree-supported paths.

Main entry points: [`totalspineseg/inference.py`](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/inference.py) (Inference entry), [`totalspineseg/init_inference.py`](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/init_inference.py) (Weight setup), [`scripts/prepare_datasets.sh`](https://github.com/neuropoly/totalspineseg/blob/main/scripts/prepare_datasets.sh) (Dataset preparation), [`totalspineseg/utils/iterative_label.py`](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/iterative_label.py) (Iterative labeling)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_inference["MRI inference"]
  node_inference_cli["Inference entry<br/>[inference.py]"]
  node_weights_init["Weight setup<br/>[init_inference.py]"]
end

subgraph group_labeling["Labeling and outputs"]
  node_step1_cleanup["Component cleanup"]
  node_iter_label["Iterative labeling<br/>[iterative_label.py]"]
  node_alternate["Odd disc extraction"]
  node_final_seg["Final segmentation<br/>[inference.py]"]
end

subgraph group_models["Models and labels"]
  node_weight_store[("Model weights")]
  node_step1_model["Step 1 model<br/>[dataset_step1.json]"]
  node_step2_model["Step 2 model<br/>[dataset_step2.json]"]
  node_label_maps["Label mappings"]
end

subgraph group_training["Training workflows"]
  node_training_data["Training datasets<br/>[training_data.json]"]
  node_prepare_data["Dataset preparation"]
  node_train_script["Training launcher<br/>[train.sh]"]
  node_trainer["nnU-Net trainer"]
end

subgraph group_support["Image utilities"]
  node_preview["Preview images<br/>[preview_jpg.py]"]
  node_augment["Image augmentation<br/>[augment.py]"]
  node_image_tools["Image transforms<br/>[image.py]"]
  node_average4d["4D averaging<br/>[average4d.py]"]
end

node_user(("MRI user"))
node_mri["MRI input"]
node_release_host(("Model releases"))
node_localizer["Localizer input"]

node_user -->|"runs"| node_inference_cli
node_user -->|"provides"| node_mri
node_mri -->|"reads"| node_inference_cli
node_inference_cli -->|"initializes"| node_weights_init
node_weights_init -->|"downloads"| node_release_host
node_weights_init -->|"installs"| node_weight_store
node_inference_cli -->|"predicts"| node_step1_model
node_step1_model -->|"produces masks"| node_step1_cleanup
node_inference_cli -->|"cleans"| node_step1_cleanup
node_inference_cli -->|"labels discs"| node_iter_label
node_inference_cli -->|"extracts discs"| node_alternate
node_inference_cli -->|"predicts"| node_step2_model
node_inference_cli -->|"assigns labels"| node_iter_label
node_localizer -.->|"optional guidance"| node_inference_cli
node_label_maps -->|"defines labels"| node_iter_label
node_inference_cli -->|"writes"| node_final_seg
node_inference_cli -->|"generates previews"| node_preview
node_training_data -->|"supplies datasets"| node_prepare_data
node_prepare_data -->|"prepares inputs"| node_train_script
node_train_script -->|"runs training"| node_trainer
node_trainer -->|"produces weights"| node_weight_store
node_prepare_data -.->|"uses augmentation"| node_augment
node_user -->|"runs utility"| node_average4d
node_user -->|"uses image tools"| node_image_tools

click node_inference_cli "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/inference.py"
click node_weights_init "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/init_inference.py"
click node_weight_store "https://github.com/neuropoly/totalspineseg/tree/main/totalspineseg/models"
click node_step1_model "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/resources/datasets/dataset_step1.json"
click node_step1_cleanup "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/largest_component.py"
click node_iter_label "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/iterative_label.py"
click node_alternate "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/extract_alternate.py"
click node_step2_model "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/resources/datasets/dataset_step2.json"
click node_label_maps "https://github.com/neuropoly/totalspineseg/tree/main/totalspineseg/resources/labels_maps"
click node_final_seg "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/inference.py"
click node_preview "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/preview_jpg.py"
click node_training_data "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/resources/data/training_data.json"
click node_prepare_data "https://github.com/neuropoly/totalspineseg/blob/main/scripts/prepare_datasets.sh"
click node_train_script "https://github.com/neuropoly/totalspineseg/blob/main/scripts/train.sh"
click node_trainer "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/models/nnUNetTrainerDAExt.py"
click node_augment "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py"
click node_image_tools "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/image.py"
click node_average4d "https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/average4d.py"

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_inference_cli,node_weights_init,node_user toneBlue
class node_step1_cleanup,node_iter_label,node_alternate,node_final_seg toneAmber
class node_weight_store,node_step1_model,node_step2_model,node_label_maps toneMint
class node_training_data,node_prepare_data,node_train_script,node_trainer toneRose
class node_preview,node_augment,node_image_tools,node_average4d,node_mri,node_release_host,node_localizer toneIndigo
```

## Components

### MRI inference

- [Inference entry](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/inference.py): `totalspineseg/inference.py`
- [Weight setup](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/init_inference.py): `totalspineseg/init_inference.py`

### Labeling and outputs

- [Component cleanup](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/largest_component.py): `totalspineseg/utils/largest_component.py`
- [Iterative labeling](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/iterative_label.py): `totalspineseg/utils/iterative_label.py`
- [Odd disc extraction](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/extract_alternate.py): `totalspineseg/utils/extract_alternate.py`
- [Final segmentation](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/inference.py): `totalspineseg/inference.py`

### Models and labels

- [Model weights](https://github.com/neuropoly/totalspineseg/tree/main/totalspineseg/models): `totalspineseg/models`
- [Step 1 model](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/resources/datasets/dataset_step1.json): `totalspineseg/resources/datasets/dataset_step1.json`
- [Step 2 model](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/resources/datasets/dataset_step2.json): `totalspineseg/resources/datasets/dataset_step2.json`
- [Label mappings](https://github.com/neuropoly/totalspineseg/tree/main/totalspineseg/resources/labels_maps): `totalspineseg/resources/labels_maps`

### Training workflows

- [Training datasets](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/resources/data/training_data.json): `totalspineseg/resources/data/training_data.json`
- [Dataset preparation](https://github.com/neuropoly/totalspineseg/blob/main/scripts/prepare_datasets.sh): `scripts/prepare_datasets.sh`
- [Training launcher](https://github.com/neuropoly/totalspineseg/blob/main/scripts/train.sh): `scripts/train.sh`
- [nnU-Net trainer](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/models/nnUNetTrainerDAExt.py): `totalspineseg/models/nnUNetTrainerDAExt.py`

### Image utilities

- [Preview images](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/preview_jpg.py): `totalspineseg/utils/preview_jpg.py`
- [Image augmentation](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py): `totalspineseg/utils/augment.py`
- [Image transforms](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/image.py): `totalspineseg/utils/image.py`
- [4D averaging](https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/average4d.py): `totalspineseg/utils/average4d.py`

### Other components

- **MRI user**
- **MRI input**
- **Model releases**
- **Localizer input**

## Connections

- MRI user → Inference entry: runs
- MRI user → MRI input: provides
- MRI input → Inference entry: reads
- Inference entry → Weight setup: initializes
- Weight setup → Model releases: downloads
- Weight setup → Model weights: installs
- Inference entry → Step 1 model: predicts
- Step 1 model → Component cleanup: produces masks
- Inference entry → Component cleanup: cleans
- Inference entry → Iterative labeling: labels discs
- Inference entry → Odd disc extraction: extracts discs
- Inference entry → Step 2 model: predicts
- Inference entry → Iterative labeling: assigns labels
- Localizer input → Inference entry: optional guidance
- Label mappings → Iterative labeling: defines labels
- Inference entry → Final segmentation: writes
- Inference entry → Preview images: generates previews
- Training datasets → Dataset preparation: supplies datasets
- Dataset preparation → Training launcher: prepares inputs
- Training launcher → nnU-Net trainer: runs training
- nnU-Net trainer → Model weights: produces weights
- Dataset preparation → Image augmentation: uses augmentation
- MRI user → 4D averaging: runs utility
- MRI user → Image transforms: uses image tools