smilelab
/

BrainSegFounder

@@ -1,279 +0,0 @@
----
-license: other
-license_name: uk-biobank-mta
-license_link: https://www.ukbiobank.ac.uk/media/p3el3p0f/mta-v20-final-29mar2021-clean.pdf
-tags:
-- medical
-- mri
-- brain-segmentation
-- self-supervised-learning
-- swin-transformer
-- monai
-- uk-biobank
-- brats
-- atlas
-pipeline_tag: image-segmentation
----
-# BrainSegFounder: UK Biobank Brain MRI Foundation Models
-## Overview
-This repository contains pretrained model weights derived from UK Biobank MRI data and downstream medical imaging datasets. These weights correspond to the self-supervised pretraining and supervised finetuning stages described in the **BrainSegFounder** framework.
-**Paper:** [BrainSegFounder: Self-supervised Learning for Brain MRI Segmentation](https://doi.org/10.1016/j.media.2024.103301)
-**Please cite this paper if you use these model weights.**
-### Available Model Weights
-This repository includes:
-| Model File | Description | Training Data | Subjects |
-|------------|-------------|---------------|----------|
-| `model_weights_UKB-pretrain.pt` | SSL pretraining weights | UK Biobank MRI (fields 20252-20253) | ~41,000 |
-| `model_weights_BRATS-pretrain.pt` | SSL pretraining weights | UKB + BraTS (multimodal MRI) | 42,470 |
-| `model_weights_BRATS-finetune.pt` | SwinUNETR finetuning weights | BraTS tumor segmentation | 1,470 |
-| `model_weights_ATLAS-pretrain.pt` | SSL pretraining weights | UKB + ATLAS v2.0 (multimodal MRI) | 42,271 |
-| `model_weights_ATLAS-finetune.pt` | SwinUNETR finetuning weights | ATLAS stroke lesion segmentation | 1,271 |
-| `SSL_Head.py` | Source code | SSL head implementation for pretraining | - |
-This README provides instructions for loading the weights, architecture descriptions, dataset information, and required UK Biobank privacy/policy statements.
-## Quick Start
-### Installation
-All models use MONAI components and require:
-```bash
-pip install git+https://github.com/Project-MONAI/MONAI.git@a23c7f54
-pip install torch
-```
-### Usage Examples
-#### Loading UKB-only Pretraining Weights
-```python
-import torch
-from huggingface_hub import hf_hub_download
-from SSL_Head import SSLHead  # Download SSL_Head.py from this repo
-from argparse import Namespace
-# Download model weights
-model_path = hf_hub_download(
-    repo_id="smilelab/BrainSegFounder",
-    filename="model_weights_UKB-pretrain.pt"
-)
-# Configure model
-args = Namespace(
-    in_channels=2,  # T1 + T2
-    spatial_dims=3,
-    bottleneck_depth=768,
-    feature_size=48,
-    num_swin_blocks_per_stage=[2,2,2,2],
-    num_heads_per_stage=[3,6,12,24],
-    dropout_path_rate=0.0,
-    use_checkpoint=False
-)
-# Load model
-model = SSLHead(args)
-checkpoint = torch.load(model_path, map_location="cpu")
-# Extract state_dict (checkpoints have nested structure)
-if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-    state_dict = checkpoint['state_dict']
-else:
-    state_dict = checkpoint
-model.load_state_dict(state_dict)
-model.eval()
-```
-#### Loading Multimodal Pretraining Weights (UKB + BraTS/ATLAS)
-```python
-import torch
-from huggingface_hub import hf_hub_download
-from SSL_Head import SSLHead
-from argparse import Namespace
-# Download BRATS or ATLAS pretrain weights
-model_path = hf_hub_download(
-    repo_id="smilelab/BrainSegFounder",
-    filename="model_weights_BRATS-pretrain.pt"  # or model_weights_ATLAS-pretrain.pt
-)
-args = Namespace(
-    in_channels=2,
-    spatial_dims=3,
-    bottleneck_depth=768,
-    feature_size=48,
-    num_swin_blocks_per_stage=[2,2,2,2],
-    num_heads_per_stage=[3,6,12,24],
-    dropout_path_rate=0.0,
-    use_checkpoint=False
-)
-model = SSLHead(args)
-checkpoint = torch.load(model_path, map_location="cpu")
-# Extract state_dict (checkpoints have nested structure)
-if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-    state_dict = checkpoint['state_dict']
-else:
-    state_dict = checkpoint
-model.load_state_dict(state_dict)
-model.eval()
-```
-#### Loading Finetuned Segmentation Weights (SwinUNETR)
-```python
-import torch
-from huggingface_hub import hf_hub_download
-from monai.networks.nets import SwinUNETR
-# Download BRATS or ATLAS finetune weights
-model_path = hf_hub_download(
-    repo_id="smilelab/BrainSegFounder",
-    filename="model_weights_BRATS-finetune.pt"  # or model_weights_ATLAS-finetune.pt
-)
-# Configure SwinUNETR
-depths = [2, 2, 2, 2]
-num_heads = [3, 6, 12, 24]
-# MONAI 1.x+ uses spatial_dims=3, older versions use img_size=(96,96,96)
-model = SwinUNETR(
-    spatial_dims=3,  # Use img_size=(96, 96, 96) for MONAI 0.x
-    in_channels=4,
-    out_channels=3,
-    feature_size=48,
-    use_checkpoint=False,
-    depths=depths,
-    num_heads=num_heads
-)
-checkpoint = torch.load(model_path, map_location="cpu")
-# Extract state_dict (checkpoints have nested structure)
-if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-    state_dict = checkpoint['state_dict']
-else:
-    state_dict = checkpoint
-model.load_state_dict(state_dict)
-model.eval()
-```
-## Architecture Description
-This project uses two architectures:
-1. **SSLHead** - Self-supervised pretraining model built on a 3D Swin Transformer encoder
-2. **SwinUNETR** - Supervised segmentation model for downstream tasks
-### SSLHead (Pretraining Model)
-The SSL pretraining model is a 3D SwinViT encoder equipped with self-supervised learning heads for **rotation prediction** and **contrastive representation learning**, as well as a VAE-style trilinear decoder for volume reconstruction.
-**Key components:**
-- **Backbone**: 3D Swin Transformer (`SwinViT` from MONAI)
-  - Patch size: `[2,2,2]`
-  - Window size: `[7,7,7]`
-  - Embedding dimension: 48
-  - Depths: `[2,2,2,2]`
-  - Number of heads: `[3,6,12,24]`
-- **Bottleneck dimension**: 768
-- **Self-supervised heads**:
-  - Rotation prediction: `nn.Linear(768, 4)`
-  - Contrastive learning: `nn.Linear(768, 512)`
-- **Reconstruction decoder**: 5-stage upsampling with 3D convolutions + InstanceNorm + LeakyReLU (trilinear interpolation)
-The SSLHead combines three self-supervised learning objectives in equal proportion:
-- **Reconstruction loss**: VAE-style image reconstruction
-- **Rotation prediction**: 4-way rotation classification
-- **Contrastive loss**: Feature representation learning
-See [SSL_Head.py](SSL_Head.py) for the complete implementation.
-### SwinUNETR (Finetuning Model)
-The finetuning model uses MONAI's 3D **SwinUNETR**, which integrates a Swin Transformer encoder with a U-Net-style hierarchical decoder for dense segmentation tasks.
-**Key components:**
-- Patch-based 3D Swin Transformer encoder
-- U-Net-style symmetric decoder with skip connections
-- Depths: `[2,2,2,2]`
-- Number of heads: `[3,6,12,24]`
-- Input image size: `(96, 96, 96)`
-- Output channels: 3 (task-dependent segmentation classes)
-Both architectures are implemented using **MONAI** (commit `a23c7f54`).
-## Training Data
-### UK Biobank (Pretraining)
-- **Subjects**: ~41,000
-- **Data fields**: 20252 (T1-weighted MRI), 20253 (T2-weighted MRI)
-- **Preprocessing**: Standard MONAI transforms (resizing and normalization)
-- **Used in**: `model_weights_UKB-pretrain.pt`
-### BraTS (Pretraining and Finetuning)
-- **Subjects**: 1,470
-- **Task**: Brain tumor segmentation
-- **Modalities**: Multi-modal 3D MRI
-- **Preprocessing**: Standard normalization and cropping
-- **Used in**: `model_weights_BRATS-pretrain.pt`, `model_weights_BRATS-finetune.pt`
-### ATLAS v2.0 (Pretraining and Finetuning)
-- **Subjects**: 1,271
-- **Task**: Stroke lesion segmentation
-- **Preprocessing**: Standard MONAI transforms
-- **Used in**: `model_weights_ATLAS-pretrain.pt`, `model_weights_ATLAS-finetune.pt`
-## Citation
-If you use these model weights, please cite:
-```bibtex
-@article{brainsegfounder2024,
-  title={BrainSegFounder: Self-supervised Learning for Brain MRI Segmentation},
-  journal={Medical Image Analysis},
-  year={2024},
-  doi={10.1016/j.media.2024.103301},
-  url={https://doi.org/10.1016/j.media.2024.103301}
-}
-```
-## Privacy and Data Protection Statement
-The returned model parameters do **not** contain any UK Biobank participant-level data, do not embed identifiable features, and cannot be used to reconstruct or infer individual-level MRI volumes.
-The models store only aggregated statistical representations learned across tens of thousands of participants. They do not contain per-participant embeddings, IDs, or reconstructed images.
-This submission complies with the UK Biobank requirement that *any parameters derived from UKB data be returned as derived variables*, while ensuring that no participant-level data or re-identifiable elements are included.
-## License and Allowed Use
-These derived variables (model weights) may be accessed by future approved UK Biobank researchers under UK Biobank's standard access procedures.
-The model architectures and research software used to train these models remain the intellectual property of the submitting researcher. The learned parameters, however, are derived from UK Biobank data and are therefore returned in accordance with:
-- The UK Biobank Material Transfer Agreement (MTA)
-- The UK Biobank "Use of Artificial Intelligence Applications and Models" policy
-**License**: The model weights are subject to the [UK Biobank Material Transfer Agreement](https://www.ukbiobank.ac.uk/media/p3el3p0f/mta-v20-final-29mar2021-clean.pdf).
-## Contact
-For questions about these models or collaborations, please open an issue in this repository or contact the authors through the paper.