katherlab/braintrain

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.