smilelab
/

BrainSegFounder

@@ -1,354 +1,279 @@
-#!/usr/bin/env python3
-"""
-Test script to validate all BrainSegFounder model weights.
-This script checks if models can be loaded and perform inference.
-"""
-import os
-import sys
 import torch
 from argparse import Namespace
-from pathlib import Path
-try:
-    from huggingface_hub import hf_hub_download
-    HF_AVAILABLE = True
-except ImportError:
-    HF_AVAILABLE = False
-    print("Warning: huggingface_hub not available. Will only test local files.")
-try:
-    from monai.networks.nets import SwinUNETR
-    MONAI_AVAILABLE = True
-except ImportError:
-    MONAI_AVAILABLE = False
-    print("Warning: MONAI not available. Install with: pip install git+https://github.com/Project-MONAI/MONAI.git@a23c7f54")
-# Try to import SSL_Head
-try:
-    from SSL_Head import SSLHead
-    SSL_HEAD_AVAILABLE = True
-except ImportError:
-    SSL_HEAD_AVAILABLE = False
-    print("Warning: SSL_Head.py not found in current directory or Python path.")
-class ModelTester:
-    """Test suite for BrainSegFounder models."""
-    def __init__(self, use_local=True, use_hf=False):
-        """
-        Initialize the model tester.
-        Args:
-            use_local: Test local model files
-            use_hf: Download and test from Hugging Face
-        """
-        self.use_local = use_local
-        self.use_hf = use_hf
-        self.repo_id = "smilelab/BrainSegFounder"
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Using device: {self.device}")
-        print("=" * 70)
-    def test_ssl_pretrain_model(self, model_path, model_name):
-        """Test SSL pretraining models (SSLHead)."""
-        if not SSL_HEAD_AVAILABLE:
-            print(f"⚠ Skipping {model_name}: SSL_Head not available")
-            return False
-        print(f"\n[Testing] {model_name}")
-        print("-" * 70)
-        try:
-            # Configure model
-            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
-            )
-            # Load model
-            print(f"  Loading model from: {model_path}")
-            model = SSLHead(args)
-            checkpoint = torch.load(model_path, map_location="cpu")
-            # Handle checkpoint format (may be nested under 'state_dict' key)
-            if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-                print(f"  Detected checkpoint format with nested state_dict")
-                state_dict = checkpoint['state_dict']
-            else:
-                state_dict = checkpoint
-            model.load_state_dict(state_dict)
-            model.to(self.device)
-            model.eval()
-            # Test forward pass with dummy input
-            print(f"  Creating dummy input (2 channels, 96x96x96)...")
-            dummy_input = torch.randn(1, 2, 96, 96, 96).to(self.device)
-            print(f"  Running forward pass...")
-            with torch.no_grad():
-                x_rot, x_contrastive, x_rec = model(dummy_input)
-            # Validate outputs
-            assert x_rot.shape == (1, 4), f"Expected rotation output shape (1, 4), got {x_rot.shape}"
-            assert x_contrastive.shape == (1, 512), f"Expected contrastive output shape (1, 512), got {x_contrastive.shape}"
-            assert x_rec.shape == (1, 2, 96, 96, 96), f"Expected reconstruction output shape (1, 2, 96, 96, 96), got {x_rec.shape}"
-            print(f"  ✓ Rotation output shape: {x_rot.shape}")
-            print(f"  ✓ Contrastive output shape: {x_contrastive.shape}")
-            print(f"  ✓ Reconstruction output shape: {x_rec.shape}")
-            print(f"  ✓ Model parameters: {sum(p.numel() for p in model.parameters()):,}")
-            print(f"✓ {model_name} passed all tests!")
-            # Clean up
-            del model, state_dict, dummy_input, x_rot, x_contrastive, x_rec
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            return True
-        except Exception as e:
-            print(f"✗ {model_name} failed: {str(e)}")
-            return False
-    def test_swinunetr_model(self, model_path, model_name):
-        """Test finetuned segmentation models (SwinUNETR)."""
-        if not MONAI_AVAILABLE:
-            print(f"⚠ Skipping {model_name}: MONAI not available")
-            return False
-        print(f"\n[Testing] {model_name}")
-        print("-" * 70)
-        try:
-            # Configure model
-            depths = [2, 2, 2, 2]
-            num_heads = [3, 6, 12, 24]
-            print(f"  Loading model from: {model_path}")
-            # Try with img_size first (older MONAI versions)
-            try:
-                model = SwinUNETR(
-                    img_size=(96, 96, 96),
-                    in_channels=4,
-                    out_channels=3,
-                    feature_size=48,
-                    use_checkpoint=False,
-                    depths=depths,
-                    num_heads=num_heads
-                )
-            except TypeError:
-                # Newer MONAI versions use spatial_size instead of img_size
-                print(f"  Using spatial_size parameter (newer MONAI)")
-                model = SwinUNETR(
-                    spatial_size=(96, 96, 96),
-                    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")
-            # Handle checkpoint format (may be nested under 'state_dict' key)
-            if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-                print(f"  Detected checkpoint format with nested state_dict")
-                state_dict = checkpoint['state_dict']
-            else:
-                state_dict = checkpoint
-            model.load_state_dict(state_dict)
-            model.to(self.device)
-            model.eval()
-            # Test forward pass with dummy input
-            print(f"  Creating dummy input (4 channels, 96x96x96)...")
-            dummy_input = torch.randn(1, 4, 96, 96, 96).to(self.device)
-            print(f"  Running forward pass...")
-            with torch.no_grad():
-                output = model(dummy_input)
-            # Validate output
-            assert output.shape == (1, 3, 96, 96, 96), f"Expected output shape (1, 3, 96, 96, 96), got {output.shape}"
-            print(f"  ✓ Output shape: {output.shape}")
-            print(f"  ✓ Model parameters: {sum(p.numel() for p in model.parameters()):,}")
-            print(f"✓ {model_name} passed all tests!")
-            # Clean up
-            del model, state_dict, dummy_input, output
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            return True
-        except Exception as e:
-            print(f"✗ {model_name} failed: {str(e)}")
-            return False
-    def test_local_models(self):
-        """Test all local model files."""
-        print("\n" + "=" * 70)
-        print("TESTING LOCAL MODEL FILES")
-        print("=" * 70)
-        models = {
-            "model_weights_UKB-pretrain.pt": ("ssl", "UKB Pretrain"),
-            "model_weights_BRATS-pretrain.pt": ("ssl", "BRATS Pretrain"),
-            "model_weights_ATLAS-pretrain.pt": ("ssl", "ATLAS Pretrain"),
-            "model_weights_BRATS-finetune.pt": ("swinunetr", "BRATS Finetune"),
-            "model_weights_ATLAS-finetune.pt": ("swinunetr", "ATLAS Finetune"),
-        }
-        results = {}
-        for filename, (model_type, display_name) in models.items():
-            if not os.path.exists(filename):
-                print(f"\n⚠ Skipping {display_name}: {filename} not found locally")
-                results[display_name] = "not_found"
-                continue
-            if model_type == "ssl":
-                success = self.test_ssl_pretrain_model(filename, display_name)
-            else:
-                success = self.test_swinunetr_model(filename, display_name)
-            results[display_name] = "passed" if success else "failed"
-        return results
-    def test_huggingface_models(self):
-        """Test models downloaded from Hugging Face."""
-        if not HF_AVAILABLE:
-            print("\n⚠ Skipping Hugging Face tests: huggingface_hub not installed")
-            return {}
-        print("\n" + "=" * 70)
-        print("TESTING HUGGING FACE MODELS")
-        print("=" * 70)
-        models = {
-            "model_weights_UKB-pretrain.pt": ("ssl", "UKB Pretrain (HF)"),
-            "model_weights_BRATS-pretrain.pt": ("ssl", "BRATS Pretrain (HF)"),
-            "model_weights_ATLAS-pretrain.pt": ("ssl", "ATLAS Pretrain (HF)"),
-            "model_weights_BRATS-finetune.pt": ("swinunetr", "BRATS Finetune (HF)"),
-            "model_weights_ATLAS-finetune.pt": ("swinunetr", "ATLAS Finetune (HF)"),
-        }
-        results = {}
-        for filename, (model_type, display_name) in models.items():
-            try:
-                print(f"\n[Downloading] {display_name} from Hugging Face...")
-                model_path = hf_hub_download(
-                    repo_id=self.repo_id,
-                    filename=filename,
-                    cache_dir=".hf_cache"
-                )
-                if model_type == "ssl":
-                    success = self.test_ssl_pretrain_model(model_path, display_name)
-                else:
-                    success = self.test_swinunetr_model(model_path, display_name)
-                results[display_name] = "passed" if success else "failed"
-            except Exception as e:
-                print(f"✗ Failed to download/test {display_name}: {str(e)}")
-                results[display_name] = "failed"
-        return results
-    def print_summary(self, local_results, hf_results):
-        """Print test summary."""
-        print("\n" + "=" * 70)
-        print("TEST SUMMARY")
-        print("=" * 70)
-        all_results = {}
-        if local_results:
-            print("\nLocal Models:")
-            for name, status in local_results.items():
-                symbol = "✓" if status == "passed" else "✗" if status == "failed" else "⚠"
-                print(f"  {symbol} {name}: {status}")
-                all_results[name] = status
-        if hf_results:
-            print("\nHugging Face Models:")
-            for name, status in hf_results.items():
-                symbol = "✓" if status == "passed" else "✗"
-                print(f"  {symbol} {name}: {status}")
-                all_results[name] = status
-        # Overall statistics
-        passed = sum(1 for s in all_results.values() if s == "passed")
-        failed = sum(1 for s in all_results.values() if s == "failed")
-        skipped = sum(1 for s in all_results.values() if s == "not_found")
-        total = len(all_results)
-        print(f"\nOverall: {passed}/{total} passed, {failed} failed, {skipped} skipped")
-        return failed == 0 and passed > 0
-def main():
-    """Main test function."""
-    import argparse
-    parser = argparse.ArgumentParser(
-        description="Test BrainSegFounder model weights"
-    )
-    parser.add_argument(
-        "--local",
-        action="store_true",
-        default=True,
-        help="Test local model files (default: True)"
-    )
-    parser.add_argument(
-        "--hf",
-        action="store_true",
-        help="Download and test models from Hugging Face"
-    )
-    parser.add_argument(
-        "--no-local",
-        action="store_true",
-        help="Skip testing local files"
-    )
-    args = parser.parse_args()
-    use_local = not args.no_local
-    use_hf = args.hf
-    if not use_local and not use_hf:
-        print("Error: Must test either local or Hugging Face models (or both)")
-        sys.exit(1)
-    tester = ModelTester(use_local=use_local, use_hf=use_hf)
-    local_results = {}
-    hf_results = {}
-    if use_local:
-        local_results = tester.test_local_models()
-    if use_hf:
-        hf_results = tester.test_huggingface_models()
-    success = tester.print_summary(local_results, hf_results)
-    sys.exit(0 if success else 1)
-if __name__ == "__main__":
-    main()

+---
+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.