# katherlab/braintrain architecture

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

- Interactive diagram: https://gitdiagram.com/katherlab/braintrain
- Repository: https://github.com/katherlab/braintrain

## Overview

1. **Overall shape**
   BrainTrain is a Python/PyTorch-style MRI analysis pipeline organized around a few top-level entry-point scripts. The repo’s main flow is: raw brain MRI data is preprocessed into standardized NumPy volumes, those volumes feed supervised or self-supervised training, and the trained models then drive evaluation, feature extraction, and explainability maps.

2. **Configuration boundary**
   `config.py` is the central runtime boundary for paths, cohort names, input CSV locations, tensor directories, GPU/device selection, and likely training hyperparameters. The README indicates that the repo expects separate train/test cohort folders and output folders for models, logs, scores, evaluations, and explainability artifacts.

3. **Data ingestion and preprocessing**
   `preprocess.py` is the ingestion stage for raw MRI data. The README describes a pipeline that converts DICOM to NIfTI, applies N4 bias correction, registers to MNI, skull-strips, and resamples to fixed 3D tensor shapes such as `96×96×96`. The output of this stage is standardized `.npy` volumes named by subject ID (`eid`), which become the canonical input format for the rest of the system.

4. **Dataset and label handling**
   `utils/dataloaders/dataloader.py` and `utils/dataloaders/dataloader_ssl.py` are the main data access layer. The supervised loader pairs CSV labels with `.npy` tensors, while the SSL loader likely returns unlabeled or differently augmented volumes for pretraining. `utils/label_mapping.py` probably normalizes task-specific labels into training targets, keeping task semantics out of the training scripts.

5. **Model architecture layer**
   `utils/architectures/` is the model boundary. The presence of `sfcn_cls.py`, `sfcn_ssl2.py`, `head.py`, and `lora_layers.py` suggests a shared backbone with multiple heads and adaptation strategies:
   - `sfcn_cls.py` for supervised classification/regression
   - `sfcn_ssl2.py` for self-supervised representation learning
   - `head.py` for task-specific output layers
   - `lora_layers.py` for parameter-efficient fine-tuning
   `utils/models.py` likely provides model construction and switching between backbones such as SFCN, DenseNet, SwinTransformer, and LoRA-enabled variants as described in the README.

6. **Training modes**
   `train.py` is the supervised/transfer-learning entry point. The README shows multiple modes: training from scratch, linear probing, full fine-tuning of an SSL backbone, and LoRA adaptation. Architecturally, this implies a common training loop with a pluggable backbone, selectable frozen/unfrozen parameters, and a task head. The training script consumes the standardized `.npy` volumes and CSV labels, then writes checkpoints and logs to the configured output directories.

7. **Self-supervised learning path**
   `ssl_train.py` is a separate pretraining pipeline for unlabeled MRI data. It depends on `utils/augmentations/ssl_transforms.py`, which is the likely augmentation boundary for SSL-specific views of the same volume. This branch produces a pretrained backbone that is later reused by `train.py` for linear probing, full fine-tuning, or LoRA-based adaptation.

8. **Losses, optimization, and distributed execution**
   `utils/criterions.py` contains the task losses used across training modes, while `utils/distribution.py` suggests support for multi-GPU or distributed training/evaluation. This separation indicates that the training scripts stay thin and delegate optimization details to shared utilities, making the same core model usable in both single-device and distributed setups.

9. **Evaluation and metrics**
   `test.py` is the evaluation entry point. According to the README, it generates ROC and PR curves, confusion matrices, classification metrics, and bootstrap confidence intervals. The test stage consumes trained checkpoints and the same CSV/tensor inputs, then emits numeric scores and plots into the evaluation output tree.

10. **Explainability and atlas-based regional analysis**
   `heatmap.py` is the explainability pipeline. It produces GradCAM, saliency, and gradient-based overlays on MRI volumes, and it also performs regional attribution analysis using an AAL atlas. `prepare_atlas.py` likely prepares the atlas and region masks used to map voxel-level attributions to anatomically meaningful brain regions. This creates a second output branch from model predictions: not just scores, but spatial explanations.

11. **Feature extraction and visualization**
   `features.py` is a representation-analysis utility. The README says it extracts features from the SSL-pretrained model and visualizes them with UMAP or t-SNE. This is a downstream analysis path that sits between the trained backbone and interpretability/cluster analysis, distinct from the classifier evaluation flow.

12. **What to draw in an architecture diagram**
   The key boundaries are: raw MRI data → preprocessing → standardized `.npy` tensors; tensors + CSV metadata → dataloaders; dataloaders → model architectures; models → training/checkpointing or evaluation/explainability. The main branching point is whether the pipeline uses supervised training, SSL pretraining, feature extraction, or attribution analysis, but all branches share the same standardized volume format and central configuration.

Main entry points: [`train.py`](https://github.com/katherlab/braintrain/blob/main/train.py) (Train), [`utils/models.py`](https://github.com/katherlab/braintrain/blob/main/utils/models.py) (Models), [`config.py`](https://github.com/katherlab/braintrain/blob/main/config.py) (Config), [`ssl_train.py`](https://github.com/katherlab/braintrain/blob/main/ssl_train.py) (SSL train)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g1["Data prep"]
  node_n2["Preprocess<br/>MRI pipeline<br/>[preprocess.py]"]
end

subgraph group_g2["Core training"]
  node_n1["Config<br/>[config.py]"]
  node_n3["Train set<br/>supervised data<br/>[dataloader.py]"]
  node_n4["SSL set<br/>self-supervised data<br/>[dataloader_ssl.py]"]
  node_n5["Labels<br/>target mapping<br/>[label_mapping.py]"]
  node_n6["Models<br/>model factory<br/>[models.py]"]
  node_n11["Losses<br/>criterion<br/>[criterions.py]"]
  node_n12["Dist utils<br/>distributed helper<br/>[distribution.py]"]
  node_n13["Train<br/>[train.py]"]
  node_n14["SSL train<br/>pretraining<br/>[ssl_train.py]"]
  node_n19["SSL augments<br/>augmentation<br/>[ssl_transforms.py]"]
end

subgraph group_g3["Downstream analysis"]
  node_n15["Test<br/>evaluation<br/>[test.py]"]
  node_n16["Features<br/>representation analysis<br/>[features.py]"]
  node_n17["Heatmaps<br/>explainability<br/>[heatmap.py]"]
  node_n18["Atlas prep<br/>atlas utils<br/>[prepare_atlas.py]"]
end

subgraph group_g4["Model internals"]
  node_n7["SFCN cls<br/>backbone<br/>[sfcn_cls.py]"]
  node_n8["SFCN SSL<br/>backbone<br/>[sfcn_ssl2.py]"]
  node_n9["Task head<br/>output head<br/>[head.py]"]
  node_n10["LoRA layers<br/>adaptation<br/>[lora_layers.py]"]
end

node_n2 -->|"standardized volumes"| node_n3
node_n2 -->|"standardized volumes"| node_n4
node_n5 -->|"targets"| node_n3
node_n1 -->|"paths"| node_n2
node_n1 -->|"run config"| node_n13
node_n1 -->|"run config"| node_n14
node_n1 -->|"eval config"| node_n15
node_n1 -->|"viz config"| node_n17
node_n13 -->|"builds"| node_n6
node_n14 -->|"builds"| node_n6
node_n6 -->|"backbone"| node_n7
node_n6 -->|"backbone"| node_n8
node_n6 -->|"head"| node_n9
node_n6 -->|"adaptation"| node_n10
node_n13 -->|"objective"| node_n11
node_n14 -->|"augments"| node_n19
node_n13 -.->|"optional"| node_n12
node_n14 -.->|"optional"| node_n12
node_n13 -->|"checkpoints"| node_n15
node_n13 -->|"embeddings"| node_n16
node_n13 -->|"predictions"| node_n17
node_n17 -->|"regions"| node_n18

click node_n1 "https://github.com/katherlab/braintrain/blob/main/config.py"
click node_n2 "https://github.com/katherlab/braintrain/blob/main/preprocess.py"
click node_n3 "https://github.com/katherlab/braintrain/blob/main/utils/dataloaders/dataloader.py"
click node_n4 "https://github.com/katherlab/braintrain/blob/main/utils/dataloaders/dataloader_ssl.py"
click node_n5 "https://github.com/katherlab/braintrain/blob/main/utils/label_mapping.py"
click node_n6 "https://github.com/katherlab/braintrain/blob/main/utils/models.py"
click node_n7 "https://github.com/katherlab/braintrain/blob/main/utils/architectures/sfcn_cls.py"
click node_n8 "https://github.com/katherlab/braintrain/blob/main/utils/architectures/sfcn_ssl2.py"
click node_n9 "https://github.com/katherlab/braintrain/blob/main/utils/architectures/head.py"
click node_n10 "https://github.com/katherlab/braintrain/blob/main/utils/architectures/lora_layers.py"
click node_n11 "https://github.com/katherlab/braintrain/blob/main/utils/criterions.py"
click node_n12 "https://github.com/katherlab/braintrain/blob/main/utils/distribution.py"
click node_n13 "https://github.com/katherlab/braintrain/blob/main/train.py"
click node_n14 "https://github.com/katherlab/braintrain/blob/main/ssl_train.py"
click node_n15 "https://github.com/katherlab/braintrain/blob/main/test.py"
click node_n16 "https://github.com/katherlab/braintrain/blob/main/features.py"
click node_n17 "https://github.com/katherlab/braintrain/blob/main/heatmap.py"
click node_n18 "https://github.com/katherlab/braintrain/blob/main/prepare_atlas.py"
click node_n19 "https://github.com/katherlab/braintrain/blob/main/utils/augmentations/ssl_transforms.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_n2 toneBlue
class node_n1,node_n3,node_n4,node_n5,node_n6,node_n11,node_n12,node_n13,node_n14,node_n19 toneAmber
class node_n15,node_n16,node_n17,node_n18 toneMint
class node_n7,node_n8,node_n9,node_n10 toneRose
```

## Components

### Data prep

Raw MRI ingestion and standardization

- [Preprocess](https://github.com/katherlab/braintrain/blob/main/preprocess.py): MRI pipeline · Converts raw scans into standardized 3D NumPy volumes. · `preprocess.py`

### Core training

Shared configuration, data access, models, and training logic

- [Config](https://github.com/katherlab/braintrain/blob/main/config.py): runtime config · Central paths, cohorts, devices, and hyperparameters. · `config.py`
- [Train set](https://github.com/katherlab/braintrain/blob/main/utils/dataloaders/dataloader.py): supervised data · CSV-driven dataset loader for labeled NumPy volumes. · `utils/dataloaders/dataloader.py`
- [SSL set](https://github.com/katherlab/braintrain/blob/main/utils/dataloaders/dataloader_ssl.py): self-supervised data · Dataset loader for unlabeled or SSL training volumes. · `utils/dataloaders/dataloader_ssl.py`
- [Labels](https://github.com/katherlab/braintrain/blob/main/utils/label_mapping.py): target mapping · Maps task-specific cohort labels into training targets. · `utils/label_mapping.py`
- [Models](https://github.com/katherlab/braintrain/blob/main/utils/models.py): model factory · Builds and switches between supported model variants. · `utils/models.py`
- [Losses](https://github.com/katherlab/braintrain/blob/main/utils/criterions.py): criterion · Shared training losses across modes. · `utils/criterions.py`
- [Dist utils](https://github.com/katherlab/braintrain/blob/main/utils/distribution.py): distributed helper · Multi-GPU and distributed execution support. · `utils/distribution.py`
- [Train](https://github.com/katherlab/braintrain/blob/main/train.py): supervised training · From-scratch, probing, fine-tuning, and LoRA training entry point. · `train.py`
- [SSL train](https://github.com/katherlab/braintrain/blob/main/ssl_train.py): pretraining · Self-supervised pretraining entry point. · `ssl_train.py`
- [SSL augments](https://github.com/katherlab/braintrain/blob/main/utils/augmentations/ssl_transforms.py): augmentation · Augmentation pipeline for SSL views of MRI volumes. · `utils/augmentations/ssl_transforms.py`

### Downstream analysis

Evaluation, features, and explainability outputs

- [Test](https://github.com/katherlab/braintrain/blob/main/test.py): evaluation · Computes metrics, curves, matrices, and confidence intervals. · `test.py`
- [Features](https://github.com/katherlab/braintrain/blob/main/features.py): representation analysis · Extracts embeddings and visualizes them with dimensionality reduction. · `features.py`
- [Heatmaps](https://github.com/katherlab/braintrain/blob/main/heatmap.py): explainability · Generates saliency, Grad-CAM, and atlas-based regional attributions. · `heatmap.py`
- [Atlas prep](https://github.com/katherlab/braintrain/blob/main/prepare_atlas.py): atlas utils · Prepares atlas and region masks for regional attribution analysis. · `prepare_atlas.py`

### Model internals

Backbone, heads, and adaptation layers

- [SFCN cls](https://github.com/katherlab/braintrain/blob/main/utils/architectures/sfcn_cls.py): backbone · Supervised SFCN variant for classification or regression. · `utils/architectures/sfcn_cls.py`
- [SFCN SSL](https://github.com/katherlab/braintrain/blob/main/utils/architectures/sfcn_ssl2.py): backbone · Self-supervised SFCN variant used for representation learning. · `utils/architectures/sfcn_ssl2.py`
- [Task head](https://github.com/katherlab/braintrain/blob/main/utils/architectures/head.py): output head · Task-specific prediction layers. · `utils/architectures/head.py`
- [LoRA layers](https://github.com/katherlab/braintrain/blob/main/utils/architectures/lora_layers.py): adaptation · Parameter-efficient fine-tuning layers. · `utils/architectures/lora_layers.py`

## Connections

- Preprocess → Train set: standardized volumes · Preprocessed NumPy volumes become supervised inputs.
- Preprocess → SSL set: standardized volumes · Preprocessed NumPy volumes also feed SSL datasets.
- Labels → Train set: targets · Label mapping supplies supervised training targets.
- Config → Preprocess: paths · Configuration defines preprocessing inputs and outputs.
- Config → Train: run config · Training uses configured devices, paths, and hyperparameters.
- Config → SSL train: run config · SSL pretraining uses the same central configuration.
- Config → Test: eval config · Evaluation reads checkpoint and output locations from config.
- Config → Heatmaps: viz config · Explainability outputs follow configured directories and cohorts.
- Train → Models: builds · Training requests model construction from the shared factory.
- SSL train → Models: builds · SSL training also constructs models through the shared factory.
- Models → SFCN cls: backbone · Factory can select the supervised SFCN backbone.
- Models → SFCN SSL: backbone · Factory can select the SSL backbone.
- Models → Task head: head · Factory attaches task-specific output heads.
- Models → LoRA layers: adaptation · Factory can enable LoRA-based adaptation.
- Train → Losses: objective · Supervised training uses shared task losses.
- SSL train → SSL augments: augments · SSL training uses specialized augmentation views.
- Train → Dist utils: optional · Training can use distributed execution utilities.
- SSL train → Dist utils: optional · SSL pretraining can use distributed execution utilities.
- Train → Test: checkpoints · Trained models are evaluated by the test pipeline.
- Train → Features: embeddings · Trained backbones can be queried for feature extraction.
- Train → Heatmaps: predictions · Trained models support attribution and heatmap generation.
- Heatmaps → Atlas prep: regions · Atlas preparation supplies masks for regional attribution.