Upload test_models.py with huggingface_hub

test_models.py

@@ -0,0 +1,354 @@
+#!/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()