V0-commit

src/BrainIAC/Brainage/README.md

@@ -0,0 +1,55 @@
+# Brain Age Prediction
+
+<p align="left">
+  <img src="brainage.jpeg" width="200" alt="Brain Age Prediction Example"/>
+</p>
+
+## Overview
+
+We present the brainage prediction training and inference code for BrainIAC as a downstream task. The pipeline is trained and infered on T1 scans, with MAE as evaluation metric.
+
+## Data Requirements
+
+- **Input**: T1-weighted MRI scans
+- **Format**: NIFTI (.nii.gz)
+- **Preprocessing**: Bias field corrected, registered to standard space, skull stripped
+- **CSV Structure**:
+  ```
+  pat_id,scandate,label
+  subject001,20240101,65    # brain age in years
+  ```
+refer to [ quickstart.ipynb](../quickstart.ipynb)  to find how to preprocess data and generate csv file.
+
+
+## Setup
+
+1. **Configuration**:
+change the [config.yml](../config.yml) file accordingly.
+   ```yaml
+   # config.yml
+   data:
+     train_csv: "path/to/train.csv"
+     val_csv: "path/to/val.csv"
+     test_csv: "path/to/test.csv"
+     root_dir: "../data/sample/processed"
+     collate: 1  # single scan framework
+    
+   checkpoints: "./checkpoints/brainage_model.00"     # for inference/testing 
+   
+   train:
+    finetune: 'yes'      # yes to finetune the entire model 
+    freeze: 'no'         # yes to freeze the resnet backbone 
+    weights: ./checkpoints/brainiac.ckpt  # path to brainiac weights
+
+   ```
+
+2. **Training**:
+   ```bash
+   python -m Brainage.train_brainage
+   ```
+
+3. **Inference**:
+   ```bash
+   python -m Brainage.infer_brainage
+   ```
+

src/BrainIAC/Brainage/infer_brainage.py

@@ -0,0 +1,85 @@
+import torch
+import pandas as pd
+import os
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from torch.cuda.amp import autocast
+from sklearn.metrics import mean_absolute_error
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from dataset2 import MedicalImageDatasetBalancedIntensity3D
+from model import Backbone, SingleScanModel, Classifier
+from utils import BaseConfig
+
+class BrainAgeInference(BaseConfig):
+    """
+    Inference class for brain age prediction model.
+    """
+
+    def __init__(self):
+        """Initialize the inference setup with model and data."""
+        super().__init__()
+        self.setup_model()
+        self.setup_data()
+        
+    def setup_model(self):
+        config = self.get_config()
+        self.backbone = Backbone()
+        self.classifier = Classifier(d_model=2048)
+        self.model = SingleScanModel(self.backbone, self.classifier)
+        
+        # Load weights
+        checkpoint = torch.load(config["infer"]["checkpoints"], map_location=self.device)
+        self.model.load_state_dict(checkpoint["model_state_dict"])
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        print("Model and checkpoint loaded!")
+        
+    ## spinup dataloaders    
+    def setup_data(self):
+        config = self.get_config()
+        self.test_dataset = MedicalImageDatasetBalancedIntensity3D(
+            csv_path=config["data"]["test_csv"],
+            root_dir=config["data"]["root_dir"]
+        )
+        self.test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self.custom_collate,
+            num_workers=1
+        )
+            
+    def infer(self):
+        """ Infer pass """
+        results_df = pd.DataFrame(columns=['PredictedAge', 'TrueAge'])
+        all_labels = []
+        all_predictions = []
+        
+        with torch.no_grad():
+            for sample in tqdm(self.test_loader, desc="Inference", unit="batch"):
+                inputs = sample['image'].to(self.device)
+                labels = sample['label'].float().to(self.device)
+                
+                with autocast():
+                    outputs = self.model(inputs)
+                
+                predictions = outputs.cpu().numpy().flatten()
+                all_labels.extend(labels.cpu().numpy().flatten())
+                all_predictions.extend(predictions)
+                
+                result = pd.DataFrame({
+                    'PredictedAge': predictions,
+                    'TrueAge': labels.cpu().numpy().flatten()
+                })
+                results_df = pd.concat([results_df, result], ignore_index=True)
+        
+        mae = mean_absolute_error(all_labels, all_predictions)
+        print(f"Mean Absolute Error (MAE): {mae:.4f} months")
+        results_df.to_csv('./data/output/brainage_output.csv', index=False)
+        
+        return mae
+
+if __name__ == "__main__":
+    inferencer = BrainAgeInference()
+    mae = inferencer.infer()

src/BrainIAC/Brainage/train_brainage.py

@@ -0,0 +1,230 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import wandb
+from tqdm import tqdm 
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
+from torch.cuda.amp import GradScaler, autocast
+from sklearn.metrics import mean_absolute_error
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from dataset2 import MedicalImageDatasetBalancedIntensity3D, TransformationMedicalImageDatasetBalancedIntensity3D
+from model import Backbone, SingleScanModel, Classifier
+from utils import BaseConfig
+
+
+class BrainAgeTrainer(BaseConfig):
+    """
+    A trainer class for brain age prediction models.
+    
+    This class handles the complete training pipeline including model setup,
+    data loading, training loop, and validation. 
+    Inherits from BaseConfig for configuration management.
+    """
+
+    def __init__(self):
+        """Initialize the trainer with model, data, and training setup."""
+        super().__init__()
+        self.setup_wandb()
+        self.setup_model()
+        self.setup_data()
+        self.setup_training()
+        
+    ## setup wandb logger    
+    def setup_wandb(self):
+        config = self.get_config()
+        wandb.init(
+            project=config['logger']['project_name'],
+            name=config['logger']['run_name'],
+            config=config
+        )
+        
+    def setup_model(self):
+        """
+        Set up the model architecture.
+        
+        Initializes the backbone and classifier blocks, and loads
+        checkpoints
+        """
+        self.backbone = Backbone()
+        self.classifier = Classifier(d_model=2048)
+        self.model = SingleScanModel(self.backbone, self.classifier)
+        
+        # Load BrainIACs weights 
+        config = self.get_config()
+        if config["train"]["finetune"] == "yes":
+            checkpoint = torch.load(config["train"]["weights"], map_location=self.device)
+            state_dict = checkpoint["state_dict"]
+            filtered_state_dict = {}
+            for key, value in state_dict.items():
+                new_key = key.replace("module.", "backbone.") if key.startswith("module.") else key
+                filtered_state_dict[new_key] = value
+            self.model.backbone.load_state_dict(filtered_state_dict, strict=False)
+            print("Pretrained weights loaded!")
+
+        # Freeze backbone if specified
+        if config["train"]["freeze"] == "yes":
+            for param in self.model.backbone.parameters():
+                param.requires_grad = False
+            print("Backbone weights frozen!")
+            
+        self.model = self.model.to(self.device)
+        
+    def setup_data(self):
+        """
+        Set up data loaders for training and validation.
+        Inherit configuration from the base config
+        """
+        config = self.get_config()
+        self.train_dataset = TransformationMedicalImageDatasetBalancedIntensity3D(
+            csv_path=config['data']['train_csv'],
+            root_dir=config["data"]["root_dir"]
+        )
+        self.val_dataset = MedicalImageDatasetBalancedIntensity3D(
+            csv_path=config['data']['val_csv'],
+            root_dir=config["data"]["root_dir"]
+        )
+        
+        self.train_loader = DataLoader(
+            self.train_dataset,
+            batch_size=config["data"]["batch_size"],
+            shuffle=True,
+            collate_fn=self.custom_collate,
+            num_workers=config["data"]["num_workers"]
+        )
+        self.val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self.custom_collate,
+            num_workers=1
+        )
+        
+    def setup_training(self):
+        """
+        Set up training config with loss, scheduler, optimizer.
+        """
+        config = self.get_config()
+        self.criterion = nn.MSELoss()
+        self.optimizer = optim.Adam(
+            self.model.parameters(),
+            lr=config['optim']['lr'],
+            weight_decay=config["optim"]["weight_decay"]
+        )
+        self.scheduler = CosineAnnealingWarmRestarts(self.optimizer, T_0=50, T_mult=2)
+        self.scaler = GradScaler()
+        
+    def train(self):
+        """
+        main training loop
+        """
+        config = self.get_config()
+        max_epochs = config['optim']['max_epochs']
+        best_val_loss = float('inf')
+        best_val_mae = float('inf')
+
+        for epoch in range(max_epochs):
+            train_loss = self.train_epoch(epoch, max_epochs)
+            val_loss, mae = self.validate_epoch(epoch, max_epochs)
+            
+            # Save best model
+            if (val_loss <= best_val_loss) and (mae <= best_val_mae):
+                print(f"Improved Val Loss from {best_val_loss:.4f} to {val_loss:.4f}")
+                print(f"Improved Val MAE from {best_val_mae:.4f} to {mae:.4f}")
+                best_val_loss = val_loss
+                best_val_mae = mae
+                self.save_checkpoint(epoch, val_loss, mae)
+                
+        wandb.finish()
+                
+    def train_epoch(self, epoch, max_epochs):
+        """
+        Train pass.
+        
+        Args:
+            epoch (int): Current epoch number
+            max_epochs (int): Total number of epochs
+            
+        Returns:
+            float: Average training loss for the epoch
+        """
+        self.model.train()
+        train_loss = 0.0
+        
+        for sample in tqdm(self.train_loader, desc=f"Training Epoch {epoch}/{max_epochs-1}"):
+            inputs = sample['image'].to(self.device)
+            labels = sample['label'].float().to(self.device)
+            
+            self.optimizer.zero_grad()
+            with autocast():
+                outputs = self.model(inputs)
+                loss = self.criterion(outputs, labels.unsqueeze(1))
+            
+            self.scaler.scale(loss).backward()
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            
+            train_loss += loss.item() * inputs.size(0)
+            
+        train_loss = train_loss / len(self.train_loader.dataset)
+        wandb.log({"Train Loss": train_loss})
+        return train_loss
+        
+    def validate_epoch(self, epoch, max_epochs):
+        """
+        Validation pass.
+        
+        Args:
+            epoch (int): Current epoch number
+            max_epochs (int): Total number of epochs
+            
+        Returns:
+            tuple: (validation_loss, mean_absolute_error)
+        """
+        self.model.eval()
+        val_loss = 0.0
+        all_labels = []
+        all_preds = []
+        
+        with torch.no_grad():
+            for sample in tqdm(self.val_loader, desc=f"Validation Epoch {epoch}/{max_epochs-1}"):
+                inputs = sample['image'].to(self.device)
+                labels = sample['label'].float().to(self.device)
+                
+                outputs = self.model(inputs)
+                loss = self.criterion(outputs, labels.unsqueeze(1))
+                
+                val_loss += loss.item() * inputs.size(0)
+                all_labels.extend(labels.cpu().numpy().flatten())
+                all_preds.extend(outputs.cpu().numpy().flatten())
+                
+        val_loss = val_loss / len(self.val_loader.dataset)
+        mae = mean_absolute_error(all_labels, all_preds)
+        
+        wandb.log({"Val Loss": val_loss, "MAE": mae})
+        self.scheduler.step(val_loss)
+        
+        print(f"Epoch {epoch}/{max_epochs-1} Val Loss: {val_loss:.4f} MAE: {mae:.4f}")
+        return val_loss, mae
+        
+    def save_checkpoint(self, epoch, loss, mae):
+        """
+        Save model checkpoint.
+        """
+        config = self.get_config()
+        checkpoint = {
+            'model_state_dict': self.model.state_dict(),
+            'loss': loss,
+            'epoch': epoch,
+        }
+        save_path = os.path.join(
+            config['logger']['save_dir'],
+            config['logger']['save_name'].format(epoch=epoch, loss=loss, metric=mae)
+        )
+        torch.save(checkpoint, save_path)
+
+if __name__ == "__main__":
+    trainer = BrainAgeTrainer()
+    trainer.train()

src/BrainIAC/HD_BET/config.py

@@ -0,0 +1,121 @@
+import numpy as np
+import torch
+from HD_BET.utils import SetNetworkToVal, softmax_helper
+from abc import abstractmethod
+from HD_BET.network_architecture import Network
+
+
+class BaseConfig(object):
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def get_split(self, fold, random_state=12345):
+        pass
+
+    @abstractmethod
+    def get_network(self, mode="train"):
+        pass
+
+    @abstractmethod
+    def get_basic_generators(self, fold):
+        pass
+
+    @abstractmethod
+    def get_data_generators(self, fold):
+        pass
+
+    def preprocess(self, data):
+        return data
+
+    def __repr__(self):
+        res = ""
+        for v in vars(self):
+            if not v.startswith("__") and not v.startswith("_") and v != 'dataset':
+                res += (v + ": " + str(self.__getattribute__(v)) + "\n")
+        return res
+
+
+class HD_BET_Config(BaseConfig):
+    def __init__(self):
+        super(HD_BET_Config, self).__init__()
+
+        self.EXPERIMENT_NAME = self.__class__.__name__ # just a generic experiment name
+
+        # network parameters
+        self.net_base_num_layers = 21
+        self.BATCH_SIZE = 2
+        self.net_do_DS = True
+        self.net_dropout_p = 0.0
+        self.net_use_inst_norm = True
+        self.net_conv_use_bias = True
+        self.net_norm_use_affine = True
+        self.net_leaky_relu_slope = 1e-1
+
+        # hyperparameters
+        self.INPUT_PATCH_SIZE = (128, 128, 128)
+        self.num_classes = 2
+        self.selected_data_channels = range(1)
+
+        # data augmentation
+        self.da_mirror_axes = (2, 3, 4)
+
+        # validation
+        self.val_use_DO = False
+        self.val_use_train_mode = False # for dropout sampling
+        self.val_num_repeats = 1 # only useful if dropout sampling
+        self.val_batch_size = 1 # only useful if dropout sampling
+        self.val_save_npz = True
+        self.val_do_mirroring = True # test time data augmentation via mirroring
+        self.val_write_images = True
+        self.net_input_must_be_divisible_by = 16  # we could make a network class that has this as a property
+        self.val_min_size = self.INPUT_PATCH_SIZE
+        self.val_fn = None
+
+        # CAREFUL! THIS IS A HACK TO MAKE PYTORCH 0.3 STATE DICTS COMPATIBLE WITH PYTORCH 0.4 (setting keep_runnings_
+        # stats=True but not using them in validation. keep_runnings_stats was True before 0.3 but unused and defaults
+        # to false in 0.4)
+        self.val_use_moving_averages = False
+
+    def get_network(self, train=True, pretrained_weights=None):
+        net = Network(self.num_classes, len(self.selected_data_channels), self.net_base_num_layers,
+                               self.net_dropout_p, softmax_helper, self.net_leaky_relu_slope, self.net_conv_use_bias,
+                               self.net_norm_use_affine, True, self.net_do_DS)
+
+        if pretrained_weights is not None:
+            net.load_state_dict(
+                torch.load(pretrained_weights, map_location=lambda storage, loc: storage))
+
+        if train:
+            net.train(True)
+        else:
+            net.train(False)
+            net.apply(SetNetworkToVal(self.val_use_DO, self.val_use_moving_averages))
+            net.do_ds = False
+
+        optimizer = None
+        self.lr_scheduler = None
+        return net, optimizer
+
+    def get_data_generators(self, fold):
+        pass
+
+    def get_split(self, fold, random_state=12345):
+        pass
+
+    def get_basic_generators(self, fold):
+        pass
+
+    def on_epoch_end(self, epoch):
+        pass
+
+    def preprocess(self, data):
+        data = np.copy(data)
+        for c in range(data.shape[0]):
+            data[c] -= data[c].mean()
+            data[c] /= data[c].std()
+        return data
+
+
+config = HD_BET_Config
+

src/BrainIAC/HD_BET/data_loading.py

@@ -0,0 +1,121 @@
+import SimpleITK as sitk
+import numpy as np
+from skimage.transform import resize
+
+
+def resize_image(image, old_spacing, new_spacing, order=3):
+    new_shape = (int(np.round(old_spacing[0]/new_spacing[0]*float(image.shape[0]))),
+                 int(np.round(old_spacing[1]/new_spacing[1]*float(image.shape[1]))),
+                 int(np.round(old_spacing[2]/new_spacing[2]*float(image.shape[2]))))
+    return resize(image, new_shape, order=order, mode='edge', cval=0, anti_aliasing=False)
+
+
+def preprocess_image(itk_image, is_seg=False, spacing_target=(1, 0.5, 0.5)):
+    spacing = np.array(itk_image.GetSpacing())[[2, 1, 0]]
+    image = sitk.GetArrayFromImage(itk_image).astype(float)
+
+    assert len(image.shape) == 3, "The image has unsupported number of dimensions. Only 3D images are allowed"
+
+    if not is_seg:
+        if np.any([[i != j] for i, j in zip(spacing, spacing_target)]):
+            image = resize_image(image, spacing, spacing_target).astype(np.float32)
+
+        image -= image.mean()
+        image /= image.std()
+    else:
+        new_shape = (int(np.round(spacing[0] / spacing_target[0] * float(image.shape[0]))),
+                     int(np.round(spacing[1] / spacing_target[1] * float(image.shape[1]))),
+                     int(np.round(spacing[2] / spacing_target[2] * float(image.shape[2]))))
+        image = resize_segmentation(image, new_shape, 1)
+    return image
+
+
+def load_and_preprocess(mri_file):
+    images = {}
+    # t1
+    images["T1"] = sitk.ReadImage(mri_file)
+
+    properties_dict = {
+        "spacing": images["T1"].GetSpacing(),
+        "direction": images["T1"].GetDirection(),
+        "size": images["T1"].GetSize(),
+        "origin": images["T1"].GetOrigin()
+    }
+
+    for k in images.keys():
+        images[k] = preprocess_image(images[k], is_seg=False, spacing_target=(1.5, 1.5, 1.5))
+
+    properties_dict['size_before_cropping'] = images["T1"].shape
+
+    imgs = []
+    for seq in ['T1']:
+        imgs.append(images[seq][None])
+    all_data = np.vstack(imgs)
+    print("image shape after preprocessing: ", str(all_data[0].shape))
+    return all_data, properties_dict
+
+
+def save_segmentation_nifti(segmentation, dct, out_fname, order=1):
+    '''
+    segmentation must have the same spacing as the original nifti (for now). segmentation may have been cropped out
+    of the original image
+
+    dct:
+    size_before_cropping
+    brain_bbox
+    size -> this is the original size of the dataset, if the image was not resampled, this is the same as size_before_cropping
+    spacing
+    origin
+    direction
+
+    :param segmentation:
+    :param dct:
+    :param out_fname:
+    :return:
+    '''
+    old_size = dct.get('size_before_cropping')
+    bbox = dct.get('brain_bbox')
+    if bbox is not None:
+        seg_old_size = np.zeros(old_size)
+        for c in range(3):
+            bbox[c][1] = np.min((bbox[c][0] + segmentation.shape[c], old_size[c]))
+        seg_old_size[bbox[0][0]:bbox[0][1],
+                     bbox[1][0]:bbox[1][1],
+                     bbox[2][0]:bbox[2][1]] = segmentation
+    else:
+        seg_old_size = segmentation
+    if np.any(np.array(seg_old_size) != np.array(dct['size'])[[2, 1, 0]]):
+        seg_old_spacing = resize_segmentation(seg_old_size, np.array(dct['size'])[[2, 1, 0]], order=order)
+    else:
+        seg_old_spacing = seg_old_size
+    seg_resized_itk = sitk.GetImageFromArray(seg_old_spacing.astype(np.int32))
+    seg_resized_itk.SetSpacing(np.array(dct['spacing'])[[0, 1, 2]])
+    seg_resized_itk.SetOrigin(dct['origin'])
+    seg_resized_itk.SetDirection(dct['direction'])
+    sitk.WriteImage(seg_resized_itk, out_fname)
+
+
+def resize_segmentation(segmentation, new_shape, order=3, cval=0):
+    '''
+    Taken from batchgenerators (https://github.com/MIC-DKFZ/batchgenerators) to prevent dependency
+
+    Resizes a segmentation map. Supports all orders (see skimage documentation). Will transform segmentation map to one
+    hot encoding which is resized and transformed back to a segmentation map.
+    This prevents interpolation artifacts ([0, 0, 2] -> [0, 1, 2])
+    :param segmentation:
+    :param new_shape:
+    :param order:
+    :return:
+    '''
+    tpe = segmentation.dtype
+    unique_labels = np.unique(segmentation)
+    assert len(segmentation.shape) == len(new_shape), "new shape must have same dimensionality as segmentation"
+    if order == 0:
+        return resize(segmentation, new_shape, order, mode="constant", cval=cval, clip=True, anti_aliasing=False).astype(tpe)
+    else:
+        reshaped = np.zeros(new_shape, dtype=segmentation.dtype)
+
+        for i, c in enumerate(unique_labels):
+            reshaped_multihot = resize((segmentation == c).astype(float), new_shape, order, mode="edge", clip=True, anti_aliasing=False)
+            reshaped[reshaped_multihot >= 0.5] = c
+        return reshaped

src/BrainIAC/HD_BET/hd_bet.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+import os
+import sys
+sys.path.append("/mnt/93E8-0534/AIDAN/HDBET/")
+from HD_BET.run import run_hd_bet
+from HD_BET.utils import maybe_mkdir_p, subfiles
+import HD_BET
+
+def hd_bet(input_file_or_dir,output_file_or_dir,mode,device,tta,pp=1,save_mask=0,overwrite_existing=1):
+
+    if output_file_or_dir is None:
+        output_file_or_dir = os.path.join(os.path.dirname(input_file_or_dir),
+                                          os.path.basename(input_file_or_dir).split(".")[0] + "_bet")
+
+    
+    params_file = os.path.join(HD_BET.__path__[0], "model_final.py")
+    config_file = os.path.join(HD_BET.__path__[0], "config.py")
+
+    assert os.path.abspath(input_file_or_dir) != os.path.abspath(output_file_or_dir), "output must be different from input"
+
+    if device == 'cpu':
+        pass
+    else:
+        device = int(device)
+
+    if os.path.isdir(input_file_or_dir):
+        maybe_mkdir_p(output_file_or_dir)
+        input_files = subfiles(input_file_or_dir, suffix='_0000.nii.gz', join=False)
+
+        if len(input_files) == 0:
+            raise RuntimeError("input is a folder but no nifti files (.nii.gz) were found in here")
+
+        output_files = [os.path.join(output_file_or_dir, i) for i in input_files]
+        input_files = [os.path.join(input_file_or_dir, i) for i in input_files]
+    else:
+        if not output_file_or_dir.endswith('.nii.gz'):
+            output_file_or_dir += '.nii.gz'
+            assert os.path.abspath(input_file_or_dir) != os.path.abspath(output_file_or_dir), "output must be different from input"
+        
+        output_files = [output_file_or_dir]
+        input_files = [input_file_or_dir]
+
+    if tta == 0:
+        tta = False
+    elif tta == 1:
+        tta = True
+    else:
+        raise ValueError("Unknown value for tta: %s. Expected: 0 or 1" % str(tta))
+
+    if overwrite_existing == 0:
+        overwrite_existing = False
+    elif overwrite_existing == 1:
+        overwrite_existing = True
+    else:
+        raise ValueError("Unknown value for overwrite_existing: %s. Expected: 0 or 1" % str(overwrite_existing))
+
+    if pp == 0:
+        pp = False
+    elif pp == 1:
+        pp = True
+    else:
+        raise ValueError("Unknown value for pp: %s. Expected: 0 or 1" % str(pp))
+
+    if save_mask == 0:
+        save_mask = False
+    elif save_mask == 1:
+        save_mask = True
+    else:
+        raise ValueError("Unknown value for pp: %s. Expected: 0 or 1" % str(pp))
+
+    run_hd_bet(input_files, output_files, mode, config_file, device, pp, tta, save_mask, overwrite_existing)
+
+
+if __name__ == "__main__":
+    print("\n########################")
+    print("If you are using hd-bet, please cite the following paper:")
+    print("Isensee F, Schell M, Tursunova I, Brugnara G, Bonekamp D, Neuberger U, Wick A, Schlemmer HP, Heiland S, Wick W,"
+           "Bendszus M, Maier-Hein KH, Kickingereder P. Automated brain extraction of multi-sequence MRI using artificial"
+           "neural networks. arXiv preprint arXiv:1901.11341, 2019.")
+    print("########################\n")
+
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--input', help='input. Can be either a single file name or an input folder. If file: must be '
+                                       'nifti (.nii.gz) and can only be 3D. No support for 4d images, use fslsplit to '
+                                       'split 4d sequences into 3d images. If folder: all files ending with .nii.gz '
+                                       'within that folder will be brain extracted.', required=True, type=str)
+    parser.add_argument('-o', '--output', help='output. Can be either a filename or a folder. If it does not exist, the folder'
+                                     ' will be created', required=False, type=str)
+    parser.add_argument('-mode', type=str, default='accurate', help='can be either \'fast\' or \'accurate\'. Fast will '
+                                                                'use only one set of parameters whereas accurate will '
+                                                                'use the five sets of parameters that resulted from '
+                                                                'our cross-validation as an ensemble. Default: '
+                                                                    'accurate',
+                        required=False)
+    parser.add_argument('-device', default='0', type=str, help='used to set on which device the prediction will run. '
+                                                               'Must be either int or str. Use int for GPU id or '
+                                                               '\'cpu\' to run on CPU. When using CPU you should '
+                                                               'consider disabling tta. Default for -device is: 0',
+                        required=False)
+    parser.add_argument('-tta', default=1, required=False, type=int, help='whether to use test time data augmentation '
+                                                                          '(mirroring). 1= True, 0=False. Disable this '
+                                                                          'if you are using CPU to speed things up! '
+                                                                          'Default: 1')
+    parser.add_argument('-pp', default=1, type=int, required=False, help='set to 0 to disabe postprocessing (remove all'
+                                                                         ' but the largest connected component in '
+                                                                         'the prediction. Default: 1')
+    parser.add_argument('-s', '--save_mask', default=1, type=int, required=False, help='if set to 0 the segmentation '
+                                                                                       'mask will not be '
+                                                                                       'saved')
+    parser.add_argument('--overwrite_existing', default=1, type=int, required=False, help="set this to 0 if you don't "
+                                                                                          "want to overwrite existing "
+                                                                                          "predictions")
+
+    args = parser.parse_args()
+
+    hd_bet(args.input,args.output,args.mode,args.device,args.tta,args.pp,args.save_mask,args.overwrite_existing)
+    

src/BrainIAC/HD_BET/network_architecture.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from HD_BET.utils import softmax_helper
+
+
+class EncodingModule(nn.Module):
+    def __init__(self, in_channels, out_channels, filter_size=3, dropout_p=0.3, leakiness=1e-2, conv_bias=True,
+                 inst_norm_affine=True, lrelu_inplace=True):
+        nn.Module.__init__(self)
+        self.dropout_p = dropout_p
+        self.lrelu_inplace = lrelu_inplace
+        self.inst_norm_affine = inst_norm_affine
+        self.conv_bias = conv_bias
+        self.leakiness = leakiness
+        self.bn_1 = nn.InstanceNorm3d(in_channels, affine=self.inst_norm_affine, track_running_stats=True)
+        self.conv1 = nn.Conv3d(in_channels, out_channels, filter_size, 1, (filter_size - 1) // 2, bias=self.conv_bias)
+        self.dropout = nn.Dropout3d(dropout_p)
+        self.bn_2 = nn.InstanceNorm3d(in_channels, affine=self.inst_norm_affine, track_running_stats=True)
+        self.conv2 = nn.Conv3d(out_channels, out_channels, filter_size, 1, (filter_size - 1) // 2, bias=self.conv_bias)
+
+    def forward(self, x):
+        skip = x
+        x = F.leaky_relu(self.bn_1(x), negative_slope=self.leakiness, inplace=self.lrelu_inplace)
+        x = self.conv1(x)
+        if self.dropout_p is not None and self.dropout_p > 0:
+            x = self.dropout(x)
+        x = F.leaky_relu(self.bn_2(x), negative_slope=self.leakiness, inplace=self.lrelu_inplace)
+        x = self.conv2(x)
+        x = x + skip
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, size=None, scale_factor=None, mode='nearest', align_corners=True):
+        super(Upsample, self).__init__()
+        self.align_corners = align_corners
+        self.mode = mode
+        self.scale_factor = scale_factor
+        self.size = size
+
+    def forward(self, x):
+        return nn.functional.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode,
+                                         align_corners=self.align_corners)
+
+
+class LocalizationModule(nn.Module):
+    def __init__(self, in_channels, out_channels, leakiness=1e-2, conv_bias=True, inst_norm_affine=True,
+                 lrelu_inplace=True):
+        nn.Module.__init__(self)
+        self.lrelu_inplace = lrelu_inplace
+        self.inst_norm_affine = inst_norm_affine
+        self.conv_bias = conv_bias
+        self.leakiness = leakiness
+        self.conv1 = nn.Conv3d(in_channels, in_channels, 3, 1, 1, bias=self.conv_bias)
+        self.bn_1 = nn.InstanceNorm3d(in_channels, affine=self.inst_norm_affine, track_running_stats=True)
+        self.conv2 = nn.Conv3d(in_channels, out_channels, 1, 1, 0, bias=self.conv_bias)
+        self.bn_2 = nn.InstanceNorm3d(out_channels, affine=self.inst_norm_affine, track_running_stats=True)
+
+    def forward(self, x):
+        x = F.leaky_relu(self.bn_1(self.conv1(x)), negative_slope=self.leakiness, inplace=self.lrelu_inplace)
+        x = F.leaky_relu(self.bn_2(self.conv2(x)), negative_slope=self.leakiness, inplace=self.lrelu_inplace)
+        return x
+
+
+class UpsamplingModule(nn.Module):
+    def __init__(self, in_channels, out_channels, leakiness=1e-2, conv_bias=True, inst_norm_affine=True,
+                 lrelu_inplace=True):
+        nn.Module.__init__(self)
+        self.lrelu_inplace = lrelu_inplace
+        self.inst_norm_affine = inst_norm_affine
+        self.conv_bias = conv_bias
+        self.leakiness = leakiness
+        self.upsample = Upsample(scale_factor=2, mode="trilinear", align_corners=True)
+        self.upsample_conv = nn.Conv3d(in_channels, out_channels, 3, 1, 1, bias=self.conv_bias)
+        self.bn = nn.InstanceNorm3d(out_channels, affine=self.inst_norm_affine, track_running_stats=True)
+
+    def forward(self, x):
+        x = F.leaky_relu(self.bn(self.upsample_conv(self.upsample(x))), negative_slope=self.leakiness,
+                         inplace=self.lrelu_inplace)
+        return x
+
+
+class DownsamplingModule(nn.Module):
+    def __init__(self, in_channels, out_channels, leakiness=1e-2, conv_bias=True, inst_norm_affine=True,
+                 lrelu_inplace=True):
+        nn.Module.__init__(self)
+        self.lrelu_inplace = lrelu_inplace
+        self.inst_norm_affine = inst_norm_affine
+        self.conv_bias = conv_bias
+        self.leakiness = leakiness
+        self.bn = nn.InstanceNorm3d(in_channels, affine=self.inst_norm_affine, track_running_stats=True)
+        self.downsample = nn.Conv3d(in_channels, out_channels, 3, 2, 1, bias=self.conv_bias)
+
+    def forward(self, x):
+        x = F.leaky_relu(self.bn(x), negative_slope=self.leakiness, inplace=self.lrelu_inplace)
+        b = self.downsample(x)
+        return x, b
+
+
+class Network(nn.Module):
+    def __init__(self, num_classes=4, num_input_channels=4, base_filters=16, dropout_p=0.3,
+                 final_nonlin=softmax_helper, leakiness=1e-2, conv_bias=True, inst_norm_affine=True,
+                 lrelu_inplace=True, do_ds=True):
+        super(Network, self).__init__()
+
+        self.do_ds = do_ds
+        self.lrelu_inplace = lrelu_inplace
+        self.inst_norm_affine = inst_norm_affine
+        self.conv_bias = conv_bias
+        self.leakiness = leakiness
+        self.final_nonlin = final_nonlin
+        self.init_conv = nn.Conv3d(num_input_channels, base_filters, 3, 1, 1, bias=self.conv_bias)
+
+        self.context1 = EncodingModule(base_filters, base_filters, 3, dropout_p, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.down1 = DownsamplingModule(base_filters, base_filters * 2, leakiness=1e-2, conv_bias=True,
+                                        inst_norm_affine=True, lrelu_inplace=True)
+
+        self.context2 = EncodingModule(2 * base_filters, 2 * base_filters, 3, dropout_p, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.down2 = DownsamplingModule(2 * base_filters, base_filters * 4, leakiness=1e-2, conv_bias=True,
+                                        inst_norm_affine=True, lrelu_inplace=True)
+
+        self.context3 = EncodingModule(4 * base_filters, 4 * base_filters, 3, dropout_p, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.down3 = DownsamplingModule(4 * base_filters, base_filters * 8, leakiness=1e-2, conv_bias=True,
+                                        inst_norm_affine=True, lrelu_inplace=True)
+
+        self.context4 = EncodingModule(8 * base_filters, 8 * base_filters, 3, dropout_p, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.down4 = DownsamplingModule(8 * base_filters, base_filters * 16, leakiness=1e-2, conv_bias=True,
+                                        inst_norm_affine=True, lrelu_inplace=True)
+
+        self.context5 = EncodingModule(16 * base_filters, 16 * base_filters, 3, dropout_p, leakiness=1e-2,
+                                       conv_bias=True, inst_norm_affine=True, lrelu_inplace=True)
+
+        self.bn_after_context5 = nn.InstanceNorm3d(16 * base_filters, affine=self.inst_norm_affine, track_running_stats=True)
+        self.up1 = UpsamplingModule(16 * base_filters, 8 * base_filters, leakiness=1e-2, conv_bias=True,
+                                    inst_norm_affine=True, lrelu_inplace=True)
+
+        self.loc1 = LocalizationModule(16 * base_filters, 8 * base_filters, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.up2 = UpsamplingModule(8 * base_filters, 4 * base_filters, leakiness=1e-2, conv_bias=True,
+                                    inst_norm_affine=True, lrelu_inplace=True)
+
+        self.loc2 = LocalizationModule(8 * base_filters, 4 * base_filters, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.loc2_seg = nn.Conv3d(4 * base_filters, num_classes, 1, 1, 0, bias=False)
+        self.up3 = UpsamplingModule(4 * base_filters, 2 * base_filters, leakiness=1e-2, conv_bias=True,
+                                    inst_norm_affine=True, lrelu_inplace=True)
+
+        self.loc3 = LocalizationModule(4 * base_filters, 2 * base_filters, leakiness=1e-2, conv_bias=True,
+                                       inst_norm_affine=True, lrelu_inplace=True)
+        self.loc3_seg = nn.Conv3d(2 * base_filters, num_classes, 1, 1, 0, bias=False)
+        self.up4 = UpsamplingModule(2 * base_filters, 1 * base_filters, leakiness=1e-2, conv_bias=True,
+                                    inst_norm_affine=True, lrelu_inplace=True)
+
+        self.end_conv_1 = nn.Conv3d(2 * base_filters, 2 * base_filters, 3, 1, 1, bias=self.conv_bias)
+        self.end_conv_1_bn = nn.InstanceNorm3d(2 * base_filters, affine=self.inst_norm_affine, track_running_stats=True)
+        self.end_conv_2 = nn.Conv3d(2 * base_filters, 2 * base_filters, 3, 1, 1, bias=self.conv_bias)
+        self.end_conv_2_bn = nn.InstanceNorm3d(2 * base_filters, affine=self.inst_norm_affine, track_running_stats=True)
+        self.seg_layer = nn.Conv3d(2 * base_filters, num_classes, 1, 1, 0, bias=False)
+
+    def forward(self, x):
+        seg_outputs = []
+
+        x = self.init_conv(x)
+        x = self.context1(x)
+
+        skip1, x = self.down1(x)
+        x = self.context2(x)
+
+        skip2, x = self.down2(x)
+        x = self.context3(x)
+
+        skip3, x = self.down3(x)
+        x = self.context4(x)
+
+        skip4, x = self.down4(x)
+        x = self.context5(x)
+
+        x = F.leaky_relu(self.bn_after_context5(x), negative_slope=self.leakiness, inplace=self.lrelu_inplace)
+        x = self.up1(x)
+
+        x = torch.cat((skip4, x), dim=1)
+        x = self.loc1(x)
+        x = self.up2(x)
+
+        x = torch.cat((skip3, x), dim=1)
+        x = self.loc2(x)
+        loc2_seg = self.final_nonlin(self.loc2_seg(x))
+        seg_outputs.append(loc2_seg)
+        x = self.up3(x)
+
+        x = torch.cat((skip2, x), dim=1)
+        x = self.loc3(x)
+        loc3_seg = self.final_nonlin(self.loc3_seg(x))
+        seg_outputs.append(loc3_seg)
+        x = self.up4(x)
+
+        x = torch.cat((skip1, x), dim=1)
+        x = F.leaky_relu(self.end_conv_1_bn(self.end_conv_1(x)), negative_slope=self.leakiness,
+                         inplace=self.lrelu_inplace)
+        x = F.leaky_relu(self.end_conv_2_bn(self.end_conv_2(x)), negative_slope=self.leakiness,
+                         inplace=self.lrelu_inplace)
+        x = self.final_nonlin(self.seg_layer(x))
+        seg_outputs.append(x)
+
+        if self.do_ds:
+            return seg_outputs[::-1]
+        else:
+            return seg_outputs[-1]

src/BrainIAC/HD_BET/paths.py

@@ -0,0 +1,6 @@
+import os
+
+# please refer to the readme on where to get the parameters. Save them in this folder:
+# Original Path: "/media/sdb/divyanshu/divyanshu/aidan_segmentation/nnUNet_pLGG/home/divyanshu/hd-bet_params"
+# Updated path for Docker container:
+folder_with_parameter_files = "/app/BrainIAC/hdbet_model" 

src/BrainIAC/HD_BET/predict_case.py

@@ -0,0 +1,126 @@
+import torch
+import numpy as np
+
+
+def pad_patient_3D(patient, shape_must_be_divisible_by=16, min_size=None):
+    if not (isinstance(shape_must_be_divisible_by, list) or isinstance(shape_must_be_divisible_by, tuple)):
+        shape_must_be_divisible_by = [shape_must_be_divisible_by] * 3
+    shp = patient.shape
+    new_shp = [shp[0] + shape_must_be_divisible_by[0] - shp[0] % shape_must_be_divisible_by[0],
+               shp[1] + shape_must_be_divisible_by[1] - shp[1] % shape_must_be_divisible_by[1],
+               shp[2] + shape_must_be_divisible_by[2] - shp[2] % shape_must_be_divisible_by[2]]
+    for i in range(len(shp)):
+        if shp[i] % shape_must_be_divisible_by[i] == 0:
+            new_shp[i] -= shape_must_be_divisible_by[i]
+    if min_size is not None:
+        new_shp = np.max(np.vstack((np.array(new_shp), np.array(min_size))), 0)
+    return reshape_by_padding_upper_coords(patient, new_shp, 0), shp
+
+
+def reshape_by_padding_upper_coords(image, new_shape, pad_value=None):
+    shape = tuple(list(image.shape))
+    new_shape = tuple(np.max(np.concatenate((shape, new_shape)).reshape((2,len(shape))), axis=0))
+    if pad_value is None:
+        if len(shape) == 2:
+            pad_value = image[0,0]
+        elif len(shape) == 3:
+            pad_value = image[0, 0, 0]
+        else:
+            raise ValueError("Image must be either 2 or 3 dimensional")
+    res = np.ones(list(new_shape), dtype=image.dtype) * pad_value
+    if len(shape) == 2:
+        res[0:0+int(shape[0]), 0:0+int(shape[1])] = image
+    elif len(shape) == 3:
+        res[0:0+int(shape[0]), 0:0+int(shape[1]), 0:0+int(shape[2])] = image
+    return res
+
+
+def predict_case_3D_net(net, patient_data, do_mirroring, num_repeats, BATCH_SIZE=None,
+                           new_shape_must_be_divisible_by=16, min_size=None, main_device=0, mirror_axes=(2, 3, 4)):
+    with torch.no_grad():
+        pad_res = []
+        for i in range(patient_data.shape[0]):
+            t, old_shape = pad_patient_3D(patient_data[i], new_shape_must_be_divisible_by, min_size)
+            pad_res.append(t[None])
+
+        patient_data = np.vstack(pad_res)
+
+        new_shp = patient_data.shape
+
+        data = np.zeros(tuple([1] + list(new_shp)), dtype=np.float32)
+
+        data[0] = patient_data
+
+        if BATCH_SIZE is not None:
+            data = np.vstack([data] * BATCH_SIZE)
+
+        a = torch.rand(data.shape).float()
+
+        if main_device == 'cpu':
+            pass
+        else:
+            a = a.cuda(main_device)
+
+        if do_mirroring:
+            x = 8
+        else:
+            x = 1
+        all_preds = []
+        for i in range(num_repeats):
+            for m in range(x):
+                data_for_net = np.array(data)
+                do_stuff = False
+                if m == 0:
+                    do_stuff = True
+                    pass
+                if m == 1 and (4 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, :, :, ::-1]
+                if m == 2 and (3 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, :, ::-1, :]
+                if m == 3 and (4 in mirror_axes) and (3 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, :, ::-1, ::-1]
+                if m == 4 and (2 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, ::-1, :, :]
+                if m == 5 and (2 in mirror_axes) and (4 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, ::-1, :, ::-1]
+                if m == 6 and (2 in mirror_axes) and (3 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, ::-1, ::-1, :]
+                if m == 7 and (2 in mirror_axes) and (3 in mirror_axes) and (4 in mirror_axes):
+                    do_stuff = True
+                    data_for_net = data_for_net[:, :, ::-1, ::-1, ::-1]
+
+                if do_stuff:
+                    _ = a.data.copy_(torch.from_numpy(np.copy(data_for_net)))
+                    p = net(a)  # np.copy is necessary because ::-1 creates just a view i think
+                    p = p.data.cpu().numpy()
+
+                    if m == 0:
+                        pass
+                    if m == 1 and (4 in mirror_axes):
+                        p = p[:, :, :, :, ::-1]
+                    if m == 2 and (3 in mirror_axes):
+                        p = p[:, :, :, ::-1, :]
+                    if m == 3 and (4 in mirror_axes) and (3 in mirror_axes):
+                        p = p[:, :, :, ::-1, ::-1]
+                    if m == 4 and (2 in mirror_axes):
+                        p = p[:, :, ::-1, :, :]
+                    if m == 5 and (2 in mirror_axes) and (4 in mirror_axes):
+                        p = p[:, :, ::-1, :, ::-1]
+                    if m == 6 and (2 in mirror_axes) and (3 in mirror_axes):
+                        p = p[:, :, ::-1, ::-1, :]
+                    if m == 7 and (2 in mirror_axes) and (3 in mirror_axes) and (4 in mirror_axes):
+                        p = p[:, :, ::-1, ::-1, ::-1]
+                    all_preds.append(p)
+
+        stacked = np.vstack(all_preds)[:, :, :old_shape[0], :old_shape[1], :old_shape[2]]
+        predicted_segmentation = stacked.mean(0).argmax(0)
+        uncertainty = stacked.var(0)
+        bayesian_predictions = stacked
+        softmax_pred = stacked.mean(0)
+    return predicted_segmentation, bayesian_predictions, softmax_pred, uncertainty

src/BrainIAC/HD_BET/run.py

@@ -0,0 +1,117 @@
+import torch
+import numpy as np
+import SimpleITK as sitk
+from HD_BET.data_loading import load_and_preprocess, save_segmentation_nifti
+from HD_BET.predict_case import predict_case_3D_net
+import imp
+from HD_BET.utils import postprocess_prediction, SetNetworkToVal, get_params_fname, maybe_download_parameters
+import os
+import HD_BET
+
+
+def apply_bet(img, bet, out_fname):
+    img_itk = sitk.ReadImage(img)
+    img_npy = sitk.GetArrayFromImage(img_itk)
+    img_bet = sitk.GetArrayFromImage(sitk.ReadImage(bet))
+    img_npy[img_bet == 0] = 0
+    out = sitk.GetImageFromArray(img_npy)
+    out.CopyInformation(img_itk)
+    sitk.WriteImage(out, out_fname)
+
+
+def run_hd_bet(mri_fnames, output_fnames, mode="accurate", config_file=os.path.join(HD_BET.__path__[0], "config.py"), device=0,
+               postprocess=False, do_tta=True, keep_mask=True, overwrite=True):
+    """
+
+    :param mri_fnames: str or list/tuple of str
+    :param output_fnames: str or list/tuple of str. If list: must have the same length as output_fnames
+    :param mode: fast or accurate
+    :param config_file: config.py
+    :param device: either int (for device id) or 'cpu'
+    :param postprocess: whether to do postprocessing or not. Postprocessing here consists of simply discarding all
+    but the largest predicted connected component. Default False
+    :param do_tta: whether to do test time data augmentation by mirroring along all axes. Default: True. If you use
+    CPU you may want to turn that off to speed things up
+    :return:
+    """
+
+    list_of_param_files = []
+
+    if mode == 'fast':
+        params_file = get_params_fname(0)
+        maybe_download_parameters(0)
+
+        list_of_param_files.append(params_file)
+    elif mode == 'accurate':
+        for i in range(5):
+            params_file = get_params_fname(i)
+            maybe_download_parameters(i)
+
+            list_of_param_files.append(params_file)
+    else:
+        raise ValueError("Unknown value for mode: %s. Expected: fast or accurate" % mode)
+
+    assert all([os.path.isfile(i) for i in list_of_param_files]), "Could not find parameter files"
+
+    cf = imp.load_source('cf', config_file)
+    cf = cf.config()
+
+    net, _ = cf.get_network(cf.val_use_train_mode, None)
+    if device == "cpu":
+        net = net.cpu()
+    else:
+        net.cuda(device)
+
+    if not isinstance(mri_fnames, (list, tuple)):
+        mri_fnames = [mri_fnames]
+
+    if not isinstance(output_fnames, (list, tuple)):
+        output_fnames = [output_fnames]
+
+    assert len(mri_fnames) == len(output_fnames), "mri_fnames and output_fnames must have the same length"
+
+    params = []
+    for p in list_of_param_files:
+        params.append(torch.load(p, map_location=lambda storage, loc: storage))
+
+    for in_fname, out_fname in zip(mri_fnames, output_fnames):
+        mask_fname = out_fname[:-7] + "_mask.nii.gz"
+        if overwrite or (not (os.path.isfile(mask_fname) and keep_mask) or not os.path.isfile(out_fname)):
+            print("File:", in_fname)
+            print("preprocessing...")
+            try:
+                data, data_dict = load_and_preprocess(in_fname)
+            except RuntimeError:
+                print("\nERROR\nCould not read file", in_fname, "\n")
+                continue
+            except AssertionError as e:
+                print(e)
+                continue
+
+            softmax_preds = []
+
+            print("prediction (CNN id)...")
+            for i, p in enumerate(params):
+                print(i)
+                net.load_state_dict(p)
+                net.eval()
+                net.apply(SetNetworkToVal(False, False))
+                _, _, softmax_pred, _ = predict_case_3D_net(net, data, do_tta, cf.val_num_repeats,
+                                                            cf.val_batch_size, cf.net_input_must_be_divisible_by,
+                                                            cf.val_min_size, device, cf.da_mirror_axes)
+                softmax_preds.append(softmax_pred[None])
+
+            seg = np.argmax(np.vstack(softmax_preds).mean(0), 0)
+
+            if postprocess:
+                seg = postprocess_prediction(seg)
+
+            print("exporting segmentation...")
+            save_segmentation_nifti(seg, data_dict, mask_fname)
+
+            apply_bet(in_fname, mask_fname, out_fname)
+
+            if not keep_mask:
+                os.remove(mask_fname)
+
+

src/BrainIAC/HD_BET/utils.py

@@ -0,0 +1,115 @@
+from urllib.request import urlopen
+import torch
+from torch import nn
+import numpy as np
+from skimage.morphology import label
+import os
+from HD_BET.paths import folder_with_parameter_files
+
+
+def get_params_fname(fold):
+    return os.path.join(folder_with_parameter_files, "%d.model" % fold)
+
+
+def maybe_download_parameters(fold=0, force_overwrite=False):
+    """
+    Downloads the parameters for some fold if it is not present yet.
+    :param fold:
+    :param force_overwrite: if True the old parameter file will be deleted (if present) prior to download
+    :return:
+    """
+
+    assert 0 <= fold <= 4, "fold must be between 0 and 4"
+
+    if not os.path.isdir(folder_with_parameter_files):
+        maybe_mkdir_p(folder_with_parameter_files)
+
+    out_filename = get_params_fname(fold)
+
+    if force_overwrite and os.path.isfile(out_filename):
+        os.remove(out_filename)
+
+    if not os.path.isfile(out_filename):
+        url = "https://zenodo.org/record/2540695/files/%d.model?download=1" % fold
+        print("Downloading", url, "...")
+        data = urlopen(url).read()
+        #out_filename = "/media/sdb/divyanshu/divyanshu/aidan_segmentation/nnUNet_pLGG/home/divyanshu/hd-bet_params/0.model"
+        with open(out_filename, 'wb') as f:
+            f.write(data)
+
+
+def init_weights(module):
+    if isinstance(module, nn.Conv3d):
+        module.weight = nn.init.kaiming_normal(module.weight, a=1e-2)
+        if module.bias is not None:
+            module.bias = nn.init.constant(module.bias, 0)
+
+
+def softmax_helper(x):
+    rpt = [1 for _ in range(len(x.size()))]
+    rpt[1] = x.size(1)
+    x_max = x.max(1, keepdim=True)[0].repeat(*rpt)
+    e_x = torch.exp(x - x_max)
+    return e_x / e_x.sum(1, keepdim=True).repeat(*rpt)
+
+
+class SetNetworkToVal(object):
+    def __init__(self, use_dropout_sampling=False, norm_use_average=True):
+        self.norm_use_average = norm_use_average
+        self.use_dropout_sampling = use_dropout_sampling
+
+    def __call__(self, module):
+        if isinstance(module, nn.Dropout3d) or isinstance(module, nn.Dropout2d) or isinstance(module, nn.Dropout):
+            module.train(self.use_dropout_sampling)
+        elif isinstance(module, nn.InstanceNorm3d) or isinstance(module, nn.InstanceNorm2d) or \
+                isinstance(module, nn.InstanceNorm1d) \
+                or isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.BatchNorm3d) or \
+                isinstance(module, nn.BatchNorm1d):
+            module.train(not self.norm_use_average)
+
+
+def postprocess_prediction(seg):
+    # basically look for connected components and choose the largest one, delete everything else
+    print("running postprocessing... ")
+    mask = seg != 0
+    lbls = label(mask, connectivity=mask.ndim)
+    lbls_sizes = [np.sum(lbls == i) for i in np.unique(lbls)]
+    largest_region = np.argmax(lbls_sizes[1:]) + 1
+    seg[lbls != largest_region] = 0
+    return seg
+
+
+def subdirs(folder, join=True, prefix=None, suffix=None, sort=True):
+    if join:
+        l = os.path.join
+    else:
+        l = lambda x, y: y
+    res = [l(folder, i) for i in os.listdir(folder) if os.path.isdir(os.path.join(folder, i))
+           and (prefix is None or i.startswith(prefix))
+           and (suffix is None or i.endswith(suffix))]
+    if sort:
+        res.sort()
+    return res
+
+
+def subfiles(folder, join=True, prefix=None, suffix=None, sort=True):
+    if join:
+        l = os.path.join
+    else:
+        l = lambda x, y: y
+    res = [l(folder, i) for i in os.listdir(folder) if os.path.isfile(os.path.join(folder, i))
+           and (prefix is None or i.startswith(prefix))
+           and (suffix is None or i.endswith(suffix))]
+    if sort:
+        res.sort()
+    return res
+
+
+subfolders = subdirs  # I am tired of confusing those
+
+
+def maybe_mkdir_p(directory):
+    splits = directory.split("/")[1:]
+    for i in range(0, len(splits)):
+        if not os.path.isdir(os.path.join("", *splits[:i+1])):
+            os.mkdir(os.path.join("", *splits[:i+1]))

src/BrainIAC/MCIclassification/README.md

@@ -0,0 +1,52 @@
+# MCI Classification
+
+<p align="left">
+  <img src="mci.jpeg" width="200" alt="MCI Classification Example"/>
+</p>
+
+## Overview
+
+We present the MCI classification training and inference code for BrainIAC as a downstream task. The pipeline is trained and infered on T1 scans, with AUC and F1 as evaluation metric.
+
+## Data Requirements
+
+- **Input**: T1-weighted MR scans
+- **Format**: NIFTI (.nii.gz)
+- **Preprocessing**: Bias field corrected, registered to standard space, skull stripped, histogram normalized (optional)
+- **CSV Structure**:
+  ```
+  pat_id,scandate,label
+  subject001,20240101,1    # 1 for MCI, 0 for healthy control
+  ```
+refer to [ quickstart.ipynb](../quickstart.ipynb) to find how to preprocess data and generate csv file.
+
+## Setup
+
+1. **Configuration**:
+change the [config.yml](../config.yml) file accordingly.
+   ```yaml
+   # config.yml
+   data:
+     train_csv: "path/to/train.csv"
+     val_csv: "path/to/val.csv"
+     test_csv: "path/to/test.csv"
+     root_dir: "../data/sample/processed"
+     collate: 1  # single scan framework
+    
+   checkpoints: "./checkpoints/mci_model.00"     # for inference/testing 
+   
+   train:
+    finetune: 'yes'      # yes to finetune the entire model 
+    freeze: 'no'         # yes to freeze the resnet backbone 
+    weights: ./checkpoints/brainiac.ckpt  # path to brainiac weights
+   ```
+
+2. **Training**:
+   ```bash
+   python -m MCIclassification.train_mci
+   ```
+
+3. **Inference**:
+   ```bash
+   python -m MCIclassification.infer_mci
+   ```

src/BrainIAC/MCIclassification/infer_mci.py

@@ -0,0 +1,142 @@
+import torch
+import pandas as pd
+import os
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from torch.cuda.amp import autocast
+from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score, accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+import numpy as np
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from dataset2 import MedicalImageDatasetBalancedIntensity3D
+from model import Backbone, SingleScanModelBP, Classifier
+from utils import BaseConfig
+
+
+
+def calculate_metrics(pred_probs, pred_labels, true_labels):
+    """
+    classification metrics.
+    Args:
+        pred_probs (numpy.ndarray): Predicted probabilities
+        pred_labels (numpy.ndarray): Predicted labels
+        true_labels (numpy.ndarray): Ground truth labels
+        
+    Returns:
+        dict: Dictionary containing accuracy, precision, recall, F1, and AUC metrics
+    """
+    accuracy = accuracy_score(true_labels, pred_labels)
+    precision = precision_score(true_labels, pred_labels)
+    recall = recall_score(true_labels, pred_labels)
+    f1 = f1_score(true_labels, pred_labels)
+    auc = roc_auc_score(true_labels, pred_probs)
+    
+    return {
+        'accuracy': accuracy,
+        'precision': precision,
+        'recall': recall,
+        'f1': f1,
+        'auc': auc
+    }
+
+
+#============================
+#  INFERENCE CLASS
+#============================
+class MCIInference(BaseConfig):
+    """
+    Inference class for MCI classification model.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.setup_model()
+        self.setup_data()
+        
+    def setup_model(self):
+        config = self.get_config()
+        self.backbone = Backbone()
+        self.classifier = Classifier(d_model=2048, num_classes=1)  # Binary classification
+        self.model = SingleScanModelBP(self.backbone, self.classifier)
+        
+        # Load weights
+        checkpoint = torch.load(config["infer"]["checkpoints"], map_location=self.device, weights_only=False)
+        self.model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        print("Model and checkpoint loaded!")
+        
+    ## spin up data loaders
+    def setup_data(self):
+        config = self.get_config()
+        self.test_dataset = MedicalImageDatasetBalancedIntensity3D(
+            csv_path=config["data"]["test_csv"],
+            root_dir=config["data"]["root_dir"]
+        )
+        self.test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self.custom_collate,
+            num_workers=1
+        )
+            
+    def infer(self):
+        """
+        Run inference pass
+        
+        Returns:
+            dict: Dictionary with evaluation metrics
+        """
+        results_df = pd.DataFrame(columns=['PredictedProb', 'PredictedLabel', 'TrueLabel'])
+        all_labels = []
+        all_predictions = []
+        all_probs = []
+        
+        with torch.no_grad():
+            for sample in tqdm(self.test_loader, desc="Inference", unit="batch"):
+                inputs = sample['image'].to(self.device)
+                labels = sample['label'].float().to(self.device)
+                
+                with autocast():
+                    outputs = self.model(inputs)
+                
+                probs = torch.sigmoid(outputs).cpu().numpy().flatten()
+                preds = (probs > 0.5).astype(int)
+                
+                all_labels.extend(labels.cpu().numpy().flatten())
+                all_predictions.extend(preds)
+                all_probs.extend(probs)
+                
+                result = pd.DataFrame({
+                    'PredictedProb': probs,
+                    'PredictedLabel': preds,
+                    'TrueLabel': labels.cpu().numpy().flatten()
+                })
+                
+                results_df = pd.concat([results_df, result], ignore_index=True)
+        
+        # log metrics 
+        """metrics = calculate_metrics(
+            np.array(all_probs),
+            np.array(all_predictions),
+            np.array(all_labels)
+        )
+        
+    
+        print("\nTest Set Metrics:")
+        print(f"Accuracy: {metrics['accuracy']:.4f}")
+        print(f"Precision: {metrics['precision']:.4f}")
+        print(f"Recall: {metrics['recall']:.4f}")
+        print(f"F1 Score: {metrics['f1']:.4f}")
+        print(f"AUC: {metrics['auc']:.4f}")"""
+        
+        # Save results
+        print("PredictedLabel", results_df["PredictedLabel"][0])
+        results_df.to_csv('./data/output/mci_classification_predictions.csv', index=False)
+        
+        return None
+
+if __name__ == "__main__":
+    inferencer = MCIInference()
+    metrics = inferencer.infer()

src/BrainIAC/MCIclassification/train_mci.py

@@ -0,0 +1,265 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import wandb
+from tqdm import tqdm 
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.cuda.amp import GradScaler, autocast
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from dataset2 import MedicalImageDatasetBalancedIntensity3D, TransformationMedicalImageDatasetBalancedIntensity3D
+from model import Backbone, SingleScanModel, Classifier
+from utils import BaseConfig
+import numpy as np
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+
+
+def calculate_metrics(pred_probs, pred_labels, true_labels):
+    """
+    classification metrics.
+    
+    Args:
+        pred_probs (numpy.ndarray): Predicted probabilities
+        pred_labels (numpy.ndarray): Predicted labels
+        true_labels (numpy.ndarray): Ground truth labels
+        
+    Returns:
+        dict: Dictionary containing accuracy, precision, recall, F1, and AUC
+    """
+    accuracy = accuracy_score(true_labels, pred_labels)
+    precision = precision_score(true_labels, pred_labels)
+    recall = recall_score(true_labels, pred_labels)
+    f1 = f1_score(true_labels, pred_labels)
+    auc = roc_auc_score(true_labels, pred_probs)
+    
+    return {
+        'accuracy': accuracy,
+        'precision': precision,
+        'recall': recall,
+        'f1': f1,
+        'auc': auc
+    }
+
+#============================
+#  TRAINER CLASS
+#============================
+
+class MCITrainer(BaseConfig):
+    """
+    trainer class for MCI classification 
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.setup_wandb()
+        self.setup_model()
+        self.setup_data()
+        self.setup_training()
+        
+    def setup_wandb(self):
+        config = self.get_config()
+        wandb.init(
+            project=config['logger']['project_name'],
+            name=config['logger']['run_name'],
+            config=config
+        )
+        
+    def setup_model(self):
+        self.backbone = Backbone()
+        # Change classifier to output 1 value for binary classification
+        self.classifier = Classifier(d_model=2048, num_classes=1)
+        self.model = SingleScanModel(self.backbone, self.classifier)
+        
+        # Load weights from brainiac
+        config = self.get_config()
+        if config["train"]["finetune"] == "yes":
+            checkpoint = torch.load(config["train"]["weights"], map_location=self.device)
+            state_dict = checkpoint["state_dict"]
+            filtered_state_dict = {}
+            for key, value in state_dict.items():
+                new_key = key.replace("module.", "backbone.") if key.startswith("module.") else key
+                filtered_state_dict[new_key] = value
+            self.model.backbone.load_state_dict(filtered_state_dict, strict=False)
+            print("Pretrained weights loaded!")
+
+        if config["train"]["freeze"] == "yes":
+            for param in self.model.backbone.parameters():
+                param.requires_grad = False
+            print("Backbone weights frozen!")
+            
+        self.model = self.model.to(self.device)
+        
+    ## spinup dataloaders
+    def setup_data(self):
+        config = self.get_config()
+        self.train_dataset = TransformationMedicalImageDatasetBalancedIntensity3D(
+            csv_path=config['data']['train_csv'],
+            root_dir=config["data"]["root_dir"]
+        )
+        self.val_dataset = MedicalImageDatasetBalancedIntensity3D(
+            csv_path=config['data']['val_csv'],
+            root_dir=config["data"]["root_dir"]
+        )
+        
+        self.train_loader = DataLoader(
+            self.train_dataset,
+            batch_size=config["data"]["batch_size"],
+            shuffle=True,
+            collate_fn=self.custom_collate,
+            num_workers=config["data"]["num_workers"]
+        )
+        self.val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self.custom_collate,
+            num_workers=1
+        )
+        
+    def setup_training(self):
+        """
+        training setup
+        """
+        config = self.get_config()
+        # BCE loss  
+        self.criterion = nn.BCEWithLogitsLoss().to(self.device)
+        self.optimizer = optim.AdamW(
+            self.model.parameters(),
+            lr=config['optim']['lr'],
+            weight_decay=config["optim"]["weight_decay"]
+        )
+        self.scheduler = OneCycleLR(
+            self.optimizer,
+            max_lr=config['optim']['lr'],
+            epochs=config['optim']['max_epochs'],
+            steps_per_epoch=len(self.train_loader)
+        )
+        self.scaler = GradScaler()
+    
+    ## main training loop
+    def train(self):
+        config = self.get_config()
+        max_epochs = config['optim']['max_epochs']
+        best_metrics = {
+            'val_loss': float('inf'),
+            'accuracy': 0,
+            'precision': 0,
+            'recall': 0,
+            'f1': 0,
+            'auc': 0
+        }
+
+        for epoch in range(max_epochs):
+            train_loss = self.train_epoch(epoch, max_epochs)
+            val_loss, metrics = self.validate_epoch(epoch, max_epochs)
+            
+            # Save best model based on validation loss and F1 score
+            if  metrics['auc'] > best_metrics['auc']:
+                print(f"New best model found!")
+                print(f"Improved Val Loss from {best_metrics['val_loss']:.4f} to {val_loss:.4f}")
+                print(f"Improved F1 from {best_metrics['f1']:.4f} to {metrics['f1']:.4f}")
+                best_metrics.update(metrics)
+                best_metrics['val_loss'] = val_loss
+                self.save_checkpoint(epoch, val_loss, metrics)
+                
+        wandb.finish()
+
+    ## training pass
+    def train_epoch(self, epoch, max_epochs):
+        self.model.train()
+        train_loss = 0.0
+        
+        for sample in tqdm(self.train_loader, desc=f"Training Epoch {epoch}/{max_epochs-1}"):
+            inputs = sample['image'].to(self.device)
+            labels = sample['label'].float().to(self.device)
+            
+            self.optimizer.zero_grad(set_to_none=True)
+            with autocast():
+                outputs = self.model(inputs)
+                loss = self.criterion(outputs, labels.unsqueeze(1))
+            
+            self.scaler.scale(loss).backward()
+            
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.scheduler.step()
+            
+            train_loss += loss.item() * inputs.size(0)
+            
+        train_loss = train_loss / len(self.train_loader.dataset)
+        wandb.log({"Train Loss": train_loss})
+        return train_loss
+    
+    ## validation pass
+    def validate_epoch(self, epoch, max_epochs):
+        self.model.eval()
+        val_loss = 0.0
+        all_labels = []
+        all_preds = []
+        all_probs = []
+        
+        with torch.no_grad():
+            for sample in tqdm(self.val_loader, desc=f"Validation Epoch {epoch}/{max_epochs-1}"):
+                inputs = sample['image'].to(self.device)
+                labels = sample['label'].float().to(self.device)
+                
+                outputs = self.model(inputs)
+                loss = self.criterion(outputs, labels.unsqueeze(1))
+                
+                # Get probabilities and predictions
+                probs = torch.sigmoid(outputs).cpu().numpy()
+                preds = (probs > 0.5).astype(int)
+                
+                val_loss += loss.item() * inputs.size(0)
+                all_labels.extend(labels.cpu().numpy().flatten())
+                all_preds.extend(preds.flatten())
+                all_probs.extend(probs.flatten())
+                
+        val_loss = val_loss / len(self.val_loader.dataset)
+        metrics = calculate_metrics(
+            np.array(all_probs),
+            np.array(all_preds),
+            np.array(all_labels)
+        )
+        
+        wandb.log({
+            "Val Loss": val_loss,
+            "Accuracy": metrics['accuracy'],
+            "Precision": metrics['precision'],
+            "Recall": metrics['recall'],
+            "F1 Score": metrics['f1'],
+            "AUC": metrics['auc']
+        })
+        
+        print(f"Epoch {epoch}/{max_epochs-1}")
+        print(f"Val Loss: {val_loss:.4f}")
+        print(f"Accuracy: {metrics['accuracy']:.4f}")
+        print(f"Precision: {metrics['precision']:.4f}")
+        print(f"Recall: {metrics['recall']:.4f}")
+        print(f"F1 Score: {metrics['f1']:.4f}")
+        print(f"AUC: {metrics['auc']:.4f}")
+        
+        return val_loss, metrics
+    
+    ## save best model
+    def save_checkpoint(self, epoch, loss, metrics):
+        config = self.get_config()
+        checkpoint = {
+            'epoch': epoch,
+            'model_state_dict': self.model.state_dict(),
+            'metrics': metrics
+        }
+        save_path = os.path.join(
+            config['logger']['save_dir'],
+            config['logger']['save_name'].format(epoch=epoch, loss=loss, metric=metrics['f1'])
+        )
+        torch.save(checkpoint, save_path)
+
+if __name__ == "__main__":
+    trainer = MCITrainer()
+    trainer.train()