Add BrainIAC IDH Classification app with Vision Transformer model

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+.DS_Store

Dockerfile

@@ -0,0 +1,37 @@
+# Use an official Python runtime as a parent image
+FROM python:3.10-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Install necessary system dependencies (kept minimal)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    git \
+    libgl1 \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy the requirements file first to leverage Docker cache
+COPY requirements.txt ./
+
+# Install Python packages
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the entire project
+COPY . /app/
+
+# Create a non-root user (HF Spaces requirement)
+RUN useradd -m -u 1000 user
+USER user
+
+# Make sure the user owns the app directory
+COPY --chown=user:user . /app/
+
+# Expose Gradio default port
+EXPOSE 7860
+
+ENV PYTHONUNBUFFERED=1
+
+# Run the app from the src/IDH directory
+CMD ["python", "src/IDH/app_gradio.py"]

README.md

@@ -1,11 +1,29 @@
 ---
-title: IDH Classification BrainIAC
-emoji: 🐠
-colorFrom: purple
-colorTo: red
+title: BrainIAC IDH Classification
+emoji: 🧠
+colorFrom: blue
+colorTo: purple
 sdk: docker
 pinned: false
-license: cc-by-4.0
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# BrainIAC: IDH Classification
+
+A Vision Transformer model for predicting IDH mutation status from dual MRI sequences (FLAIR + T1c).
+
+## Features
+- Upload FLAIR and T1c MRI NIfTI files
+- Optional preprocessing (registration + enhancement + skull stripping)
+- Vision Transformer attention map visualization
+- Interactive slice-by-slice attention viewing
+- Real-time IDH mutation prediction with confidence scores
+
+## Usage
+1. Upload FLAIR and T1c MRI scans (.nii or .nii.gz)
+2. Optionally enable preprocessing
+3. Enable saliency maps for attention visualization
+4. Adjust prediction threshold
+5. View results and attention maps
+
+*Research use only*

requirements.txt

@@ -0,0 +1,18 @@
+monai==1.3.2
+nibabel==5.2.1
+numpy==1.23.5
+pydicom
+PyYAML
+pytorch-lightning==2.3.3
+scipy==1.10.1
+SimpleITK==2.4.0
+torch==2.6.0
+tqdm
+gradio
+pandas
+scikit-image==0.21.0
+opencv-python
+itk-elastix
+dicom2nifti
+einops
+matplotlib

src/IDH/HD_BET/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/IDH/HD_BET/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/IDH/HD_BET/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/IDH/HD_BET/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/IDH/HD_BET/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/IDH/hdbet_model" 

src/IDH/HD_BET/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/IDH/HD_BET/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/IDH/HD_BET/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/IDH/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/IDH/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/IDH/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/IDH/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/IDH/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/IDH/hdbet_model" 

src/IDH/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/IDH/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/IDH/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/IDH/app_gradio.py

@@ -0,0 +1,867 @@
+import os
+import yaml
+import torch
+import nibabel as nib
+import numpy as np
+import gradio as gr
+from typing import Tuple
+import tempfile
+import shutil
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use('Agg') # Use non-interactive backend
+import cv2 # For Gaussian Blur
+import io # For saving plots to memory
+import base64 # For encoding plots
+import uuid # For unique IDs
+import traceback # For detailed error printing
+
+import SimpleITK as sitk
+import itk
+from scipy.signal import medfilt
+import skimage.filters
+
+from monai.transforms import Compose, LoadImaged, EnsureChannelFirstd, Resized, NormalizeIntensityd, ToTensord
+
+from model import ViTBackboneNet, Classifier, SingleScanModelBP
+
+# Optional HD-BET import (packaged locally like in MCI app)
+try:
+    from HD_BET.run import run_hd_bet
+    from HD_BET.hd_bet import hd_bet
+except Exception as e:
+    print(f"Warning: HD_BET not available: {e}")
+    run_hd_bet = None
+    hd_bet = None
+
+
+APP_DIR = os.path.dirname(__file__)
+TEMPLATE_DIR = os.path.join(APP_DIR, "golden_image", "mni_templates")
+PARAMS_RIGID_PATH = os.path.join(TEMPLATE_DIR, "Parameters_Rigid.txt")
+DEFAULT_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "temp_head.nii.gz")
+FLAIR_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "nihpd_asym_04.5-18.5_t2w.nii")
+T1C_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "nihpd_asym_13.0-18.5_t1w.nii")
+HD_BET_CONFIG_PATH = os.path.join(APP_DIR, "HD_BET", "config.py")
+HD_BET_MODEL_DIR = os.path.join(APP_DIR, "hdbet_model")
+
+
+def load_config() -> dict:
+    cfg_path = os.path.join(APP_DIR, "config.yml")
+    if os.path.exists(cfg_path):
+        with open(cfg_path, "r") as f:
+            return yaml.safe_load(f)
+    # Defaults
+    return {
+        "gpu": {"device": "cpu"},
+        "infer": {
+            "checkpoints": "./checkpoints/idh_model.ckpt",
+            "simclr_checkpoint": "./checkpoints/simclr_vitb.ckpt",
+            "threshold": 0.5,
+            "image_size": [96, 96, 96],
+        },
+    }
+
+
+def build_model(cfg: dict):
+    device = torch.device(cfg.get("gpu", {}).get("device", "cpu"))
+    infer_cfg = cfg.get("infer", {})
+    simclr_path = os.path.join(APP_DIR, infer_cfg.get("simclr_checkpoint", ""))
+    ckpt_path = os.path.join(APP_DIR, infer_cfg.get("checkpoints", ""))
+
+    backbone = ViTBackboneNet(simclr_ckpt_path=simclr_path)
+    classifier = Classifier(d_model=768, num_classes=1)
+    model = SingleScanModelBP(backbone, classifier)
+
+    # Load finetuned checkpoint (Lightning or plain state_dict)
+    if os.path.exists(ckpt_path):
+        checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        if "state_dict" in checkpoint:
+            state_dict = checkpoint["state_dict"]
+            new_state_dict = {}
+            for key, value in state_dict.items():
+                if key.startswith("model."):
+                    new_state_dict[key[len("model."):]] = value
+                else:
+                    new_state_dict[key] = value
+        else:
+            new_state_dict = checkpoint
+        model.load_state_dict(new_state_dict, strict=False)
+    else:
+        print(f"Warning: Finetuned checkpoint not found at {ckpt_path}. Model will use backbone-only weights.")
+
+    model.to(device)
+    model.eval()
+    return model, device
+
+
+# ---------------- Preprocessing (Registration + Enhancement + Skull Stripping) ----------------
+
+def bias_field_correction(img_array: np.ndarray) -> np.ndarray:
+    image = sitk.GetImageFromArray(img_array.astype(np.float32))
+    if image.GetPixelID() != sitk.sitkFloat32:
+        image = sitk.Cast(image, sitk.sitkFloat32)
+    maskImage = sitk.OtsuThreshold(image, 0, 1, 200)
+    corrector = sitk.N4BiasFieldCorrectionImageFilter()
+    numberFittingLevels = 4
+    max_iters = [min(50 * (2 ** i), 200) for i in range(numberFittingLevels)]
+    corrector.SetMaximumNumberOfIterations(max_iters)
+    corrected_image = corrector.Execute(image, maskImage)
+    return sitk.GetArrayFromImage(corrected_image)
+
+
+def denoise(volume: np.ndarray, kernel_size: int = 3) -> np.ndarray:
+    return medfilt(volume, kernel_size)
+
+
+def rescale_intensity(volume: np.ndarray, percentils=[0.5, 99.5], bins_num=256) -> np.ndarray:
+    volume_float = volume.astype(np.float32)
+    try:
+        t = skimage.filters.threshold_otsu(volume_float, nbins=256)
+        volume_masked = np.copy(volume_float)
+        volume_masked[volume_masked < t] = 0
+        obj_volume = volume_masked[np.where(volume_masked > 0)]
+    except ValueError:
+        obj_volume = volume_float.flatten()
+    if obj_volume.size == 0:
+        obj_volume = volume_float.flatten()
+        min_value = np.min(obj_volume)
+        max_value = np.max(obj_volume)
+    else:
+        min_value = np.percentile(obj_volume, percentils[0])
+        max_value = np.percentile(obj_volume, percentils[1])
+    denom = max_value - min_value
+    if denom < 1e-6:
+        denom = 1e-6
+    if bins_num == 0:
+        output_volume = (volume_float - min_value) / denom
+        output_volume = np.clip(output_volume, 0.0, 1.0)
+    else:
+        output_volume = np.round((volume_float - min_value) / denom * (bins_num - 1))
+        output_volume = np.clip(output_volume, 0, bins_num - 1)
+    return output_volume.astype(np.float32)
+
+
+def equalize_hist(volume: np.ndarray, bins_num=256) -> np.ndarray:
+    mask = volume > 1e-6
+    obj_volume = volume[mask]
+    if obj_volume.size == 0:
+        return volume
+    hist, bins = np.histogram(obj_volume, bins_num, range=(obj_volume.min(), obj_volume.max()))
+    cdf = hist.cumsum()
+    cdf_normalized = (bins_num - 1) * cdf / float(cdf[-1])
+    equalized_obj_volume = np.interp(obj_volume, bins[:-1], cdf_normalized)
+    equalized_volume = np.copy(volume)
+    equalized_volume[mask] = equalized_obj_volume
+    return equalized_volume.astype(np.float32)
+
+
+def run_enhance_on_file(input_nifti_path: str, output_nifti_path: str):
+    """
+    Simplified enhancement - just copy the file since N4 is now done in registration.
+    This maintains compatibility with the existing preprocessing pipeline.
+    """
+    print(f"Enhancement step (N4 already applied during registration): {input_nifti_path}")
+    # Since N4 bias correction is now handled in registration, just copy the file
+    import shutil
+    shutil.copy2(input_nifti_path, output_nifti_path)
+    print(f"Enhancement complete (passthrough): {output_nifti_path}")
+
+
+def register_image_sitk(input_nifti_path: str, output_nifti_path: str, template_path: str, interp_type='linear'):
+    """
+    MRI registration with SimpleITK matching the provided script approach.
+    
+    Args:
+        input_nifti_path: Path to input NIfTI file
+        output_nifti_path: Path to save registered output
+        template_path: Path to template image
+        interp_type: Interpolation type ('linear', 'bspline', 'nearest_neighbor')
+    """
+    print(f"Registering {input_nifti_path} to template {template_path}")
+    
+    # Read template and moving images
+    fixed_img = sitk.ReadImage(template_path, sitk.sitkFloat32)
+    moving_img = sitk.ReadImage(input_nifti_path, sitk.sitkFloat32)
+    
+    # Apply N4 bias correction to moving image
+    moving_img = sitk.N4BiasFieldCorrection(moving_img)
+    
+    # Resample fixed image to 1mm isotropic
+    old_size = fixed_img.GetSize()
+    old_spacing = fixed_img.GetSpacing()
+    new_spacing = (1, 1, 1)
+    new_size = [
+        int(round((old_size[0] * old_spacing[0]) / float(new_spacing[0]))),
+        int(round((old_size[1] * old_spacing[1]) / float(new_spacing[1]))),
+        int(round((old_size[2] * old_spacing[2]) / float(new_spacing[2])))
+    ]
+    
+    # Set interpolation type
+    if interp_type == 'linear':
+        interp_type = sitk.sitkLinear
+    elif interp_type == 'bspline':
+        interp_type = sitk.sitkBSpline
+    elif interp_type == 'nearest_neighbor':
+        interp_type = sitk.sitkNearestNeighbor
+    else:
+        interp_type = sitk.sitkLinear
+    
+    # Resample fixed image
+    resample = sitk.ResampleImageFilter()
+    resample.SetOutputSpacing(new_spacing)
+    resample.SetSize(new_size)
+    resample.SetOutputOrigin(fixed_img.GetOrigin())
+    resample.SetOutputDirection(fixed_img.GetDirection())
+    resample.SetInterpolator(interp_type)
+    resample.SetDefaultPixelValue(fixed_img.GetPixelIDValue())
+    resample.SetOutputPixelType(sitk.sitkFloat32)
+    fixed_img = resample.Execute(fixed_img)
+    
+    # Initialize transform
+    transform = sitk.CenteredTransformInitializer(
+        fixed_img, 
+        moving_img, 
+        sitk.Euler3DTransform(), 
+        sitk.CenteredTransformInitializerFilter.GEOMETRY)
+    
+    # Set up registration method
+    registration_method = sitk.ImageRegistrationMethod()
+    registration_method.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)
+    registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)
+    registration_method.SetMetricSamplingPercentage(0.01)
+    registration_method.SetInterpolator(sitk.sitkLinear)
+    registration_method.SetOptimizerAsGradientDescent(
+        learningRate=1.0, 
+        numberOfIterations=100, 
+        convergenceMinimumValue=1e-6, 
+        convergenceWindowSize=10)
+    registration_method.SetOptimizerScalesFromPhysicalShift()
+    registration_method.SetShrinkFactorsPerLevel(shrinkFactors=[4, 2, 1])
+    registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas=[2, 1, 0])
+    registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()
+    registration_method.SetInitialTransform(transform)
+    
+    # Execute registration
+    final_transform = registration_method.Execute(fixed_img, moving_img)
+    
+    # Apply transform and save registered image
+    moving_img_resampled = sitk.Resample(
+        moving_img, 
+        fixed_img, 
+        final_transform, 
+        sitk.sitkLinear, 
+        0.0, 
+        moving_img.GetPixelID())
+    
+    sitk.WriteImage(moving_img_resampled, output_nifti_path)
+    print(f"Registration complete. Saved to: {output_nifti_path}")
+
+
+def register_image(input_nifti_path: str, output_nifti_path: str):
+    """Wrapper to maintain compatibility - now uses SimpleITK registration."""
+    if not os.path.exists(DEFAULT_TEMPLATE_PATH):
+        raise FileNotFoundError(f"Template file missing: {DEFAULT_TEMPLATE_PATH}")
+    register_image_sitk(input_nifti_path, output_nifti_path, DEFAULT_TEMPLATE_PATH)
+
+
+def run_skull_stripping(input_nifti_path: str, output_dir: str):
+    """
+    Brain extraction using HD-BET direct integration matching the script approach.
+    
+    Args:
+        input_nifti_path: Path to input NIfTI file
+        output_dir: Directory to save skull-stripped output
+        
+    Returns:
+        tuple: (output_file_path, output_mask_path)
+    """
+    print(f"Running HD-BET skull stripping on {input_nifti_path}")
+    
+    if hd_bet is None:
+        raise RuntimeError("HD-BET not available. Please include HD_BET and hdbet_model in src/IDH.")
+    
+    if not os.path.exists(HD_BET_MODEL_DIR):
+        raise FileNotFoundError(f"HD-BET models not found at {HD_BET_MODEL_DIR}")
+
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Get base filename and prepare HD-BET compatible naming
+    base_name = os.path.basename(input_nifti_path).replace('.nii.gz', '').replace('.nii', '')
+    
+    # HD-BET expects files with _0000 suffix - create temporary file if needed
+    temp_input_dir = os.path.join(output_dir, "temp_input")
+    os.makedirs(temp_input_dir, exist_ok=True)
+    
+    # Copy input file with _0000 suffix for HD-BET
+    temp_input_path = os.path.join(temp_input_dir, f"{base_name}_0000.nii.gz")
+    shutil.copy2(input_nifti_path, temp_input_path)
+    
+    # Set device
+    device = "0" if torch.cuda.is_available() else "cpu"
+    
+    try:
+        # Also try setting the specific model file path
+        model_file = os.path.join(HD_BET_MODEL_DIR, '0.model')
+        
+        if os.path.exists(model_file):
+            print(f"Local model file exists at: {model_file}")
+        else:
+            print(f"Warning: Model file not found at: {model_file}")
+            # List directory contents for debugging
+            if os.path.exists(HD_BET_MODEL_DIR):
+                print(f"Contents of {HD_BET_MODEL_DIR}: {os.listdir(HD_BET_MODEL_DIR)}")
+            else:
+                print(f"Directory {HD_BET_MODEL_DIR} does not exist")
+        
+        # Run HD-BET directly on the temporary directory
+        print(f"Running hd_bet with input_dir: {temp_input_dir}, output_dir: {output_dir}")
+        hd_bet(temp_input_dir, output_dir, device=device, mode='fast', tta=0)
+        
+        # HD-BET outputs files with original naming convention
+        output_file_path = os.path.join(output_dir, f"{base_name}_0000.nii.gz")
+        output_mask_path = os.path.join(output_dir, f"{base_name}_0000_mask.nii.gz")
+        
+        # Rename to expected format for compatibility
+        final_output_path = os.path.join(output_dir, f"{base_name}_bet.nii.gz")
+        final_mask_path = os.path.join(output_dir, f"{base_name}_bet_mask.nii.gz")
+        
+        if os.path.exists(output_file_path):
+            shutil.move(output_file_path, final_output_path)
+        if os.path.exists(output_mask_path):
+            shutil.move(output_mask_path, final_mask_path)
+            
+        # Clean up temporary directory
+        shutil.rmtree(temp_input_dir, ignore_errors=True)
+        
+        if not os.path.exists(final_output_path):
+            raise RuntimeError(f"HD-BET did not produce output file: {final_output_path}")
+            
+        print(f"Skull stripping complete. Output saved to: {final_output_path}")
+        return final_output_path, final_mask_path
+        
+    except Exception as e:
+        # Clean up on error
+        shutil.rmtree(temp_input_dir, ignore_errors=True)
+        raise RuntimeError(f"HD-BET skull stripping failed: {str(e)}")
+
+
+# ---------------- Saliency Generation ----------------
+
+def extract_attention_map(vit_model, image, layer_idx=-1, img_size=(96, 96, 96), patch_size=16):
+    """
+    Extracts the attention map from a Vision Transformer (ViT) model.
+
+    This function wraps the attention blocks of the ViT to capture the attention
+    weights during a forward pass. It then processes these weights to generate
+    a 3D saliency map corresponding to the model's focus on the input image.
+    """
+    attention_maps = {}
+    original_attns = {}
+
+    # A wrapper class to intercept and store attention weights from a ViT block.
+    class AttentionWithWeights(torch.nn.Module):
+        def __init__(self, original_attn_module):
+            super().__init__()
+            self.original_attn_module = original_attn_module
+            self.attn_weights = None
+
+        def forward(self, x):
+            # The original implementation of the attention module may not return
+            # the attention weights. This wrapper recalculates them to ensure they
+            # are captured. This is based on the standard ViT attention mechanism.
+            output = self.original_attn_module(x)
+            if hasattr(self.original_attn_module, 'qkv'):
+                qkv = self.original_attn_module.qkv(x)
+                batch_size, seq_len, _ = x.shape
+                # Assuming qkv has been fused and has shape (batch_size, seq_len, 3 * num_heads * head_dim)
+                qkv = qkv.reshape(batch_size, seq_len, 3, self.original_attn_module.num_heads, -1)
+                qkv = qkv.permute(2, 0, 3, 1, 4)
+                q, k, v = qkv[0], qkv[1], qkv[2]
+                attn = (q @ k.transpose(-2, -1)) * self.original_attn_module.scale
+                self.attn_weights = attn.softmax(dim=-1)
+            return output
+
+    # Store original attention modules and replace with wrappers
+    for i, block in enumerate(vit_model.blocks):
+        if hasattr(block, 'attn'):
+            original_attns[i] = block.attn
+            block.attn = AttentionWithWeights(block.attn)
+
+    try:
+        # Perform a forward pass to execute the wrapped modules and capture weights
+        with torch.no_grad():
+            _ = vit_model(image)
+
+        # Collect the captured attention weights from each block
+        for i, block in enumerate(vit_model.blocks):
+            if hasattr(block.attn, 'attn_weights') and block.attn.attn_weights is not None:
+                attention_maps[f"layer_{i}"] = block.attn.attn_weights.detach()
+
+    finally:
+        # Restore original attention modules
+        for i, original_attn in original_attns.items():
+            vit_model.blocks[i].attn = original_attn
+
+    if not attention_maps:
+        raise RuntimeError("Could not extract any attention maps. Please check the ViT model structure.")
+
+    # Select the attention map from the specified layer
+    if layer_idx < 0:
+        layer_idx = len(attention_maps) + layer_idx
+    layer_name = f"layer_{layer_idx}"
+    if layer_name not in attention_maps:
+        raise ValueError(f"Layer {layer_idx} not found. Available layers: {list(attention_maps.keys())}")
+
+    layer_attn = attention_maps[layer_name]
+    # Average attention across all heads
+    head_attn = layer_attn[0].mean(dim=0)
+    # Get attention from the [CLS] token to all other image patches
+    cls_attn = head_attn[0, 1:]
+
+    # Reshape the 1D attention vector into a 3D volume
+    patches_per_dim = img_size[0] // patch_size
+    total_patches = patches_per_dim ** 3
+    
+    # Pad or truncate if the number of patches doesn't align
+    if cls_attn.shape[0] != total_patches:
+        if cls_attn.shape[0] > total_patches:
+            cls_attn = cls_attn[:total_patches]
+        else:
+            padded = torch.zeros(total_patches, device=cls_attn.device)
+            padded[:cls_attn.shape[0]] = cls_attn
+            cls_attn = padded
+
+    cls_attn_3d = cls_attn.reshape(patches_per_dim, patches_per_dim, patches_per_dim)
+    cls_attn_3d = cls_attn_3d.unsqueeze(0).unsqueeze(0)  # Add batch and channel dims
+
+    # Upsample the attention map to the full image resolution
+    upsampled_attn = torch.nn.functional.interpolate(
+        cls_attn_3d,
+        size=img_size,
+        mode='trilinear',
+        align_corners=False
+    ).squeeze()
+
+    # Normalize the map to [0, 1] for visualization
+    upsampled_attn = upsampled_attn.cpu().numpy()
+    upsampled_attn = (upsampled_attn - upsampled_attn.min()) / (upsampled_attn.max() - upsampled_attn.min())
+    return upsampled_attn
+
+
+def generate_saliency_dual(model, input_tensor, layer_idx=-1):
+    """
+    Generate saliency maps for dual-input IDH model.
+    
+    Args:
+        model: The complete IDH model
+        input_tensor: Dual input tensor (batch_size, 2, C, D, H, W)
+        layer_idx: ViT layer to visualize
+    
+    Returns:
+        tuple: (flair_input_3d, t1c_input_3d, flair_saliency_3d)
+    """
+    print("Generating saliency maps for dual input...")
+    
+    try:
+        # Extract individual images from dual input
+        # input_tensor shape: [batch_size, 2, C, D, H, W]
+        flair_tensor = input_tensor[:, 0]  # [batch, C, D, H, W]
+        t1c_tensor = input_tensor[:, 1]    # [batch, C, D, H, W]
+        
+        # Get the ViT backbone
+        vit_model = model.backbone.backbone
+        
+        # Generate attention map only for FLAIR
+        flair_attn = extract_attention_map(vit_model, flair_tensor, layer_idx)
+        
+        # Convert input tensors to numpy for visualization
+        flair_input_3d = flair_tensor.squeeze().cpu().detach().numpy()
+        t1c_input_3d = t1c_tensor.squeeze().cpu().detach().numpy()
+        
+        print("Saliency maps generated successfully.")
+        return flair_input_3d, t1c_input_3d, flair_attn
+        
+    except Exception as e:
+        print(f"Error during saliency generation: {e}")
+        traceback.print_exc()
+        return None, None, None
+
+
+# ---------------- Visualization Functions ----------------
+
+def create_slice_plots_dual(flair_data_3d, t1c_data_3d, flair_saliency_3d, slice_index):
+    """Create slice plots for simplified dual input visualization: T1c, FLAIR, FLAIR attention."""
+    print(f"Generating plots for slice index: {slice_index}")
+    
+    if any(data is None for data in [flair_data_3d, t1c_data_3d, flair_saliency_3d]):
+        return None, None, None
+
+    # Check bounds - using axis 2 for axial slices
+    if not (0 <= slice_index < flair_data_3d.shape[2]):
+        print(f"Error: Slice index {slice_index} out of bounds (0-{flair_data_3d.shape[2]-1}).")
+        return None, None, None
+
+    def save_plot_to_numpy(fig):
+        with io.BytesIO() as buf:
+            fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75)
+            plt.close(fig)
+            buf.seek(0)
+            img_arr = plt.imread(buf, format='png')
+        return (img_arr * 255).astype(np.uint8)
+
+    try:
+        # Extract axial slices - using axis 2 (last dimension)
+        flair_slice = flair_data_3d[:, :, slice_index]
+        t1c_slice = t1c_data_3d[:, :, slice_index]
+        flair_saliency_slice = flair_saliency_3d[:, :, slice_index]
+
+        # Normalize input slices
+        def normalize_slice(slice_data, volume_data):
+            p1, p99 = np.percentile(volume_data, (1, 99))
+            denom = max(p99 - p1, 1e-6)
+            return np.clip((slice_data - p1) / denom, 0, 1)
+
+        flair_slice_norm = normalize_slice(flair_slice, flair_data_3d)
+        t1c_slice_norm = normalize_slice(t1c_slice, t1c_data_3d)
+
+        # Process saliency slice
+        def process_saliency_slice(saliency_slice, saliency_volume):
+            saliency_slice = np.copy(saliency_slice)
+            saliency_slice[saliency_slice < 0] = 0
+            saliency_slice_blurred = cv2.GaussianBlur(saliency_slice, (15, 15), 0)
+            s_max = max(np.max(saliency_volume[saliency_volume >= 0]), 1e-6)
+            saliency_slice_norm = saliency_slice_blurred / s_max
+            return np.where(saliency_slice_norm > 0.0, saliency_slice_norm, 0)
+
+        flair_sal_processed = process_saliency_slice(flair_saliency_slice, flair_saliency_3d)
+
+        # Create plots
+        plots = []
+        
+        # T1c Input
+        fig1, ax1 = plt.subplots(figsize=(6, 6))
+        ax1.imshow(t1c_slice_norm, cmap='gray', interpolation='none', origin='lower')
+        ax1.axis('off')
+        ax1.set_title('T1c Input', fontsize=14, color='white', pad=10)
+        plots.append(save_plot_to_numpy(fig1))
+
+        # FLAIR Input  
+        fig2, ax2 = plt.subplots(figsize=(6, 6))
+        ax2.imshow(flair_slice_norm, cmap='gray', interpolation='none', origin='lower')
+        ax2.axis('off')
+        ax2.set_title('FLAIR Input', fontsize=14, color='white', pad=10)
+        plots.append(save_plot_to_numpy(fig2))
+
+        # FLAIR Attention
+        fig3, ax3 = plt.subplots(figsize=(6, 6))
+        ax3.imshow(flair_sal_processed, cmap='magma', interpolation='none', origin='lower', vmin=0)
+        ax3.axis('off')
+        ax3.set_title('FLAIR Attention', fontsize=14, color='white', pad=10)
+        plots.append(save_plot_to_numpy(fig3))
+
+        print(f"Generated 3 plots successfully for axial slice {slice_index}.")
+        return tuple(plots)
+
+    except Exception as e:
+        print(f"Error generating plots for slice {slice_index}: {e}")
+        traceback.print_exc()
+        return tuple([None] * 3)
+
+
+# ---------------- Inference ----------------
+
+def get_dual_validation_transform(image_size: Tuple[int, int, int]):
+    return Compose([
+        LoadImaged(keys=["image1", "image2"]),
+        EnsureChannelFirstd(keys=["image1", "image2"]),
+        Resized(keys=["image1", "image2"], spatial_size=tuple(image_size), mode="trilinear"),
+        NormalizeIntensityd(keys=["image1", "image2"], nonzero=True, channel_wise=True),
+        ToTensord(keys=["image1", "image2"]),
+    ])
+
+
+def preprocess_dual_nifti(flair_path: str, t1c_path: str, image_size: Tuple[int, int, int], device: torch.device) -> torch.Tensor:
+    transform = get_dual_validation_transform(image_size)
+    sample = {"image1": flair_path, "image2": t1c_path}
+    sample = transform(sample)
+    img1 = sample["image1"]  # (C, D, H, W)
+    img2 = sample["image2"]  # (C, D, H, W)
+    images = torch.stack([img1, img2], dim=0).unsqueeze(0).to(device)
+    return images
+
+
+def predict_idh(flair_file, t1c_file, threshold: float, do_preprocess: bool, generate_saliency: bool, cfg: dict, model, device):
+    try:
+        if flair_file is None or t1c_file is None:
+            return {"error": "Please upload both FLAIR and T1c NIfTI files (.nii.gz)."}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "saliency_paths": None, "num_slices": 0}
+
+        flair_path = flair_file.name if hasattr(flair_file, 'name') else flair_file
+        t1c_path = t1c_file.name if hasattr(t1c_file, 'name') else t1c_file
+        
+        if not (flair_path.endswith(".nii") or flair_path.endswith(".nii.gz")):
+            return {"error": "FLAIR must be a NIfTI file (.nii or .nii.gz)."}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "saliency_paths": None, "num_slices": 0}
+        if not (t1c_path.endswith(".nii") or t1c_path.endswith(".nii.gz")):
+            return {"error": "T1c must be a NIfTI file (.nii or .nii.gz)."}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "saliency_paths": None, "num_slices": 0}
+
+        work_dir = tempfile.mkdtemp()
+        flair_final_path, t1c_final_path = flair_path, t1c_path
+        
+        try:
+            # Optional preprocessing pipeline
+            if do_preprocess:
+                # Registration (use modality-specific templates)
+                flair_reg = os.path.join(work_dir, "flair_registered.nii.gz")
+                t1c_reg = os.path.join(work_dir, "t1c_registered.nii.gz")
+                register_image_sitk(flair_path, flair_reg, FLAIR_TEMPLATE_PATH)
+                register_image_sitk(t1c_path, t1c_reg, T1C_TEMPLATE_PATH)
+                # Enhancement
+                flair_enh = os.path.join(work_dir, "flair_enhanced.nii.gz")
+                t1c_enh = os.path.join(work_dir, "t1c_enhanced.nii.gz")
+                run_enhance_on_file(flair_reg, flair_enh)
+                run_enhance_on_file(t1c_reg, t1c_enh)
+                # Skull stripping
+                skullstrip_dir = os.path.join(work_dir, "skullstripped")
+                flair_bet, _ = run_skull_stripping(flair_enh, skullstrip_dir)
+                t1c_bet, _ = run_skull_stripping(t1c_enh, skullstrip_dir)
+                flair_final_path, t1c_final_path = flair_bet, t1c_bet
+
+            # Prediction
+            image_size = cfg.get("infer", {}).get("image_size", [96, 96, 96])
+            input_tensor = preprocess_dual_nifti(flair_final_path, t1c_final_path, image_size, device)
+            
+            with torch.no_grad():
+                logits = model(input_tensor)
+                prob = torch.sigmoid(logits).cpu().numpy().flatten()[0].item()
+                predicted_class = int(prob >= threshold)
+
+            prediction_result = {
+                "IDH_mutant_probability": float(prob),
+                "threshold": float(threshold),
+                "predicted_class": int(predicted_class),
+                "preprocessing": bool(do_preprocess),
+                "class_label": "IDH-mutant" if predicted_class == 1 else "IDH-wildtype"
+            }
+
+            # Initialize saliency outputs
+            t1c_input_img = flair_input_img = flair_attn_img = None
+            slider_update = gr.Slider(visible=False)
+            saliency_state = {"input_paths": None, "saliency_paths": None, "num_slices": 0}
+
+            # Generate saliency maps if requested
+            if generate_saliency:
+                print("--- Generating Saliency Maps ---")
+                try:
+                    flair_input_3d, t1c_input_3d, flair_saliency_3d = generate_saliency_dual(model, input_tensor, layer_idx=-1)
+                    
+                    if all(data is not None for data in [flair_input_3d, t1c_input_3d, flair_saliency_3d]):
+                        num_slices = flair_input_3d.shape[2]  # Use axis 2 for axial slices
+                        center_slice_index = num_slices // 2
+
+                        # Save numpy arrays for slider callback
+                        unique_id = str(uuid.uuid4())
+                        temp_paths = []
+                        for name, data in [("flair_input", flair_input_3d), ("t1c_input", t1c_input_3d), 
+                                         ("flair_saliency", flair_saliency_3d)]:
+                            path = os.path.join(work_dir, f"{unique_id}_{name}.npy")
+                            np.save(path, data)
+                            temp_paths.append(path)
+
+                        # Generate initial plots for center slice
+                        plots = create_slice_plots_dual(flair_input_3d, t1c_input_3d, flair_saliency_3d, center_slice_index)
+                        if plots and all(p is not None for p in plots):
+                            t1c_input_img, flair_input_img, flair_attn_img = plots
+
+                        # Update state and slider
+                        saliency_state = {
+                            "input_paths": temp_paths[:2],  # [flair_input, t1c_input]
+                            "saliency_paths": temp_paths[2:],  # [flair_saliency]
+                            "num_slices": num_slices
+                        }
+                        slider_update = gr.Slider(value=center_slice_index, minimum=0, maximum=num_slices-1, step=1, label="Select Slice", visible=True)
+                        print("--- Saliency Generation Complete ---")
+                    else:
+                        print("Warning: Saliency generation failed - some outputs were None")
+                        
+                except Exception as e:
+                    print(f"Error during saliency generation: {e}")
+                    traceback.print_exc()
+
+            return (prediction_result, t1c_input_img, flair_input_img, flair_attn_img, slider_update, saliency_state)
+
+        except Exception as e:
+            shutil.rmtree(work_dir, ignore_errors=True)
+            return {"error": f"Processing failed: {str(e)}"}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "saliency_paths": None, "num_slices": 0}
+
+    except Exception as e:
+        return {"error": str(e)}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "saliency_paths": None, "num_slices": 0}
+
+
+def update_slice_viewer_dual(slice_index, current_state):
+    """Update slice viewer for dual input saliency visualization."""
+    input_paths = current_state.get("input_paths", [])
+    saliency_paths = current_state.get("saliency_paths", [])
+
+    if not input_paths or not saliency_paths or len(input_paths) != 2 or len(saliency_paths) != 1:
+        print(f"Warning: Invalid state for slice viewer update: {current_state}")
+        return None, None, None
+
+    try:
+        # Load numpy arrays
+        flair_input_3d = np.load(input_paths[0])
+        t1c_input_3d = np.load(input_paths[1])
+        flair_saliency_3d = np.load(saliency_paths[0])
+
+        # Validate slice index
+        slice_index = int(slice_index)
+        if not (0 <= slice_index < flair_input_3d.shape[2]):  # Use axis 2 for axial slices
+            print(f"Warning: Invalid slice index {slice_index}")
+            return None, None, None
+
+        # Generate new plots
+        plots = create_slice_plots_dual(flair_input_3d, t1c_input_3d, flair_saliency_3d, slice_index)
+        return plots if plots else tuple([None] * 3)
+
+    except Exception as e:
+        print(f"Error updating slice viewer for index {slice_index}: {e}")
+        traceback.print_exc()
+        return tuple([None] * 3)
+
+
+def build_interface():
+    cfg = load_config()
+    model, device = build_model(cfg)
+    default_threshold = float(cfg.get("infer", {}).get("threshold", 0.5))
+
+    with gr.Blocks(title="BrainIAC: IDH Classification", css="""
+#header-row {
+    min-height: 150px;
+    align-items: center;
+}
+.logo-img img {
+    height: 150px;
+    object-fit: contain;
+}
+""") as demo:
+        # --- Header with Logos ---
+        with gr.Row(elem_id="header-row"):
+            with gr.Column(scale=1):
+                gr.Image(os.path.join(APP_DIR, "static/images/kannlab.png"),
+                         show_label=False, interactive=False,
+                         show_download_button=False,
+                         container=False,
+                         elem_classes=["logo-img"])
+            with gr.Column(scale=3):
+                gr.Markdown(
+                    "<h1 style='text-align: center; margin-bottom: 2.5rem'>"
+                    "BrainIAC: IDH Classification"
+                    "</h1>"
+                )
+            with gr.Column(scale=1):
+                gr.Image(os.path.join(APP_DIR, "static/images/brainiac.jpeg"),
+                         show_label=False, interactive=False,
+                         show_download_button=False,
+                         container=False,
+                         elem_classes=["logo-img"])
+
+        # --- Add model description section ---
+        with gr.Accordion("ℹ️ Model Details and Usage Guide", open=False):
+            gr.Markdown("""
+### 🧠 BrainIAC: IDH Classification
+
+**Model Description**  
+A Vision Transformer (ViT) model with BrainIAC as pre-trained backbone designed to predict IDH mutation status from dual MRI sequences (FLAIR + T1c).
+
+**Training Dataset**  
+- **Subjects**: Trained on FLAIR and T1c MRI scans from glioma patients from UCSF-PDGM dataset
+- **Imaging Modalities**: FLAIR and T1c (contrast-enhanced T1-weighted)
+- **Preprocessing**: N4 bias correction, MNI registration, and skull stripping (HD-BET)
+
+**Input**  
+- Format: NIfTI (.nii or .nii.gz)  
+- Required sequences: FLAIR and T1c (both required)
+- Image size: Automatically resized to 96×96×96 voxels
+
+**Output**  
+- Binary classification: IDH-mutant or IDH-wildtype
+- Probability score for IDH mutation
+- Attention map visualization
+
+**Intended Use**  
+- Research use only!
+
+**NOTE**  
+- Requires both FLAIR and T1c sequences 
+- Not validated on other MRI sequences
+- Not validated for other brain pathologies beyond gliomas
+- Upload PHI data at own risk!
+- The model is hosted on a cloud-based CPU instance
+- The data is not stored, shared or collected for any purpose!
+
+**Preprocessing Pipeline**  
+When enabled, the preprocessing performs:
+1. **Registration**: SimpleITK-based registration to template space with mutual information metric and 1mm isotropic resampling
+2. **N4 Bias Correction**: Applied during registration step
+3. **Skull Stripping**: Remove non-brain tissue using HD-BET direct integration
+
+**Attention Maps**  
+When enabled, generates ViT attention maps showing which brain regions the model focuses on for prediction.
+
+""")
+
+        # Use gr.State to store paths to numpy arrays for the slider callback
+        saliency_state = gr.State({"input_paths": None, "saliency_paths": None, "num_slices": 0})
+
+        # Main Content
+        gr.Markdown("**Upload FLAIR and T1c NIfTI volumes** — Optional preprocessing performs registration to MNI, enhancement, and skull stripping.")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                with gr.Group():
+                    gr.Markdown("### Controls")
+                    flair_input = gr.File(label="FLAIR (.nii or .nii.gz)")
+                    t1c_input = gr.File(label="T1c (.nii or .nii.gz)")
+                    preprocess_checkbox = gr.Checkbox(value=False, label="Preprocess NIfTI (registration + enhancement + skull stripping)")
+                    generate_saliency_checkbox = gr.Checkbox(value=True, label="Generate Attention Maps")
+                    threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value=default_threshold, step=0.01, label="Decision Threshold")
+                    predict_btn = gr.Button("Predict IDH Status", variant="primary")
+            
+            with gr.Column(scale=2):
+                with gr.Group():
+                    gr.Markdown("### Classification Result")
+                    output_json = gr.JSON(label="Prediction")
+
+        # Saliency visualization section
+        with gr.Group():
+            gr.Markdown("### Attention Map Viewer (Axial Slice)")
+            slice_slider = gr.Slider(label="Select Slice", minimum=0, maximum=0, step=1, value=0, visible=False)
+            
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("<p style='text-align: center;'>T1c Input</p>")
+                    t1c_input_img = gr.Image(label="T1c Input", type="numpy", show_label=False)
+                with gr.Column():
+                    gr.Markdown("<p style='text-align: center;'>FLAIR Input</p>")
+                    flair_input_img = gr.Image(label="FLAIR Input", type="numpy", show_label=False)
+                with gr.Column():
+                    gr.Markdown("<p style='text-align: center;'>FLAIR Attention</p>")
+                    flair_attn_img = gr.Image(label="Attention Mask", type="numpy", show_label=False)
+
+        # Wire components
+        predict_btn.click(
+            fn=lambda f, t, prep, gen_sal, thr: predict_idh(f, t, thr, prep, gen_sal, cfg, model, device),
+            inputs=[flair_input, t1c_input, preprocess_checkbox, generate_saliency_checkbox, threshold_input],
+            outputs=[output_json, t1c_input_img, flair_input_img, flair_attn_img, slice_slider, saliency_state],
+        )
+
+        slice_slider.change(
+            fn=update_slice_viewer_dual,
+            inputs=[slice_slider, saliency_state],
+            outputs=[t1c_input_img, flair_input_img, flair_attn_img]
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    iface = build_interface()
+    iface.launch(server_name="0.0.0.0", server_port=7860) 

src/IDH/checkpoints/idh_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4fbbddda1f0a1c38dd3c6241725b3ed4806705085e1b812b19464867518b5bf
+size 353461323

src/IDH/config.yml

@@ -0,0 +1,7 @@
+gpu:
+  device: cpu
+infer:
+  checkpoints: ./checkpoints/idh_model.ckpt
+  simclr_checkpoint: ./checkpoints/simclr_vitb.ckpt
+  threshold: 0.5
+  image_size: [96, 96, 96] 

src/IDH/golden_image/mni_templates/Parameters_Rigid.txt

@@ -0,0 +1,141 @@
+// Example parameter file for rotation registration
+// C-style comments: //
+
+// The internal pixel type, used for internal computations
+// Leave to float in general. 
+// NB: this is not the type of the input images! The pixel 
+// type of the input images is automatically read from the 
+// images themselves.
+// This setting can be changed to "short" to save some memory
+// in case of very large 3D images.
+(FixedInternalImagePixelType "float")
+(MovingInternalImagePixelType "float")
+
+// **************** Main Components **************************
+
+// The following components should usually be left as they are:
+(Registration "MultiResolutionRegistration")
+(Interpolator "BSplineInterpolator")
+(ResampleInterpolator "FinalBSplineInterpolator")
+(Resampler "DefaultResampler")
+
+// These may be changed to Fixed/MovingSmoothingImagePyramid.
+// See the manual.
+(FixedImagePyramid "FixedRecursiveImagePyramid")
+(MovingImagePyramid "MovingRecursiveImagePyramid")
+
+// The following components are most important:
+// The optimizer AdaptiveStochasticGradientDescent (ASGD) works
+// quite ok in general. The Transform and Metric are important
+// and need to be chosen careful for each application. See manual.
+(Optimizer "AdaptiveStochasticGradientDescent")
+(Transform "EulerTransform")
+(Metric "AdvancedMattesMutualInformation")
+
+// ***************** Transformation **************************
+
+// Scales the rotations compared to the translations, to make
+// sure they are in the same range. In general, it's best to  
+// use automatic scales estimation:
+(AutomaticScalesEstimation "true")
+
+// Automatically guess an initial translation by aligning the
+// geometric centers of the fixed and moving.
+(AutomaticTransformInitialization "true")
+
+// Whether transforms are combined by composition or by addition.
+// In generally, Compose is the best option in most cases.
+// It does not influence the results very much.
+(HowToCombineTransforms "Compose")
+
+// ******************* Similarity measure *********************
+
+// Number of grey level bins in each resolution level,
+// for the mutual information. 16 or 32 usually works fine.
+// You could also employ a hierarchical strategy:
+//(NumberOfHistogramBins 16 32 64)
+(NumberOfHistogramBins 32)
+
+// If you use a mask, this option is important. 
+// If the mask serves as region of interest, set it to false.
+// If the mask indicates which pixels are valid, then set it to true.
+// If you do not use a mask, the option doesn't matter.
+(ErodeMask "false")
+
+// ******************** Multiresolution **********************
+
+// The number of resolutions. 1 Is only enough if the expected
+// deformations are small. 3 or 4 mostly works fine. For large
+// images and large deformations, 5 or 6 may even be useful.
+(NumberOfResolutions 4)
+
+// The downsampling/blurring factors for the image pyramids.
+// By default, the images are downsampled by a factor of 2
+// compared to the next resolution.
+// So, in 2D, with 4 resolutions, the following schedule is used:
+//(ImagePyramidSchedule 8 8  4 4  2 2  1 1 )
+// And in 3D:
+//(ImagePyramidSchedule 8 8 8  4 4 4  2 2 2  1 1 1 )
+// You can specify any schedule, for example:
+//(ImagePyramidSchedule 4 4  4 3  2 1  1 1 )
+// Make sure that the number of elements equals the number
+// of resolutions times the image dimension.
+
+// ******************* Optimizer ****************************
+
+// Maximum number of iterations in each resolution level:
+// 200-500 works usually fine for rigid registration.
+// For more robustness, you may increase this to 1000-2000.
+(MaximumNumberOfIterations 250)
+
+// The step size of the optimizer, in mm. By default the voxel size is used.
+// which usually works well. In case of unusual high-resolution images
+// (eg histology) it is necessary to increase this value a bit, to the size
+// of the "smallest visible structure" in the image:
+//(MaximumStepLength 1.0)
+
+// **************** Image sampling **********************
+
+// Number of spatial samples used to compute the mutual
+// information (and its derivative) in each iteration.
+// With an AdaptiveStochasticGradientDescent optimizer,
+// in combination with the two options below, around 2000
+// samples may already suffice.
+(NumberOfSpatialSamples 2048)
+
+// Refresh these spatial samples in every iteration, and select
+// them randomly. See the manual for information on other sampling
+// strategies.
+(NewSamplesEveryIteration "true")
+(ImageSampler "Random")
+
+// ************* Interpolation and Resampling ****************
+
+// Order of B-Spline interpolation used during registration/optimisation.
+// It may improve accuracy if you set this to 3. Never use 0.
+// An order of 1 gives linear interpolation. This is in most 
+// applications a good choice.
+(BSplineInterpolationOrder 1)
+
+// Order of B-Spline interpolation used for applying the final
+// deformation.
+// 3 gives good accuracy; recommended in most cases.
+// 1 gives worse accuracy (linear interpolation)
+// 0 gives worst accuracy, but is appropriate for binary images
+// (masks, segmentations); equivalent to nearest neighbor interpolation.
+(FinalBSplineInterpolationOrder 3)
+
+//Default pixel value for pixels that come from outside the picture:
+(DefaultPixelValue 0)
+
+// Choose whether to generate the deformed moving image.
+// You can save some time by setting this to false, if you are
+// only interested in the final (nonrigidly) deformed moving image
+// for example.
+(WriteResultImage "true")
+
+// The pixel type and format of the resulting deformed moving image
+(ResultImagePixelType "short")
+(ResultImageFormat "mhd")
+
+

src/IDH/golden_image/mni_templates/nihpd_asym_04.5-18.5_t2w.nii

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0bd55fefd8a74deadca777c362bbdb7db9ae15ee135dd9fe045c8838af58d3ee
+size 17350930

src/IDH/golden_image/mni_templates/nihpd_asym_13.0-18.5_t1w.nii

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f10804664000688f0ddc124b39ce3ae27f2c339a40583bd6ff916727e97b77d0
+size 17350930

src/IDH/hdbet_model/0.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f75233753c4750672815e2b7a86db754995ae44b8f1cd77bccfc37becd2d83c
+size 65443735

src/IDH/hdbet_model/hdbet_model/0.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f75233753c4750672815e2b7a86db754995ae44b8f1cd77bccfc37becd2d83c
+size 65443735

src/IDH/model.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+from monai.networks.nets import ViT
+import os
+
+
+class ViTBackboneNet(nn.Module):
+    def __init__(self, simclr_ckpt_path: str):
+        super().__init__()
+        self.backbone = ViT(
+            in_channels=1,
+            img_size=(96, 96, 96),
+            patch_size=(16, 16, 16),
+            hidden_size=768,
+            mlp_dim=3072,
+            num_layers=12,
+            num_heads=12,
+            save_attn=True,
+        )
+        # Load pretrained weights from SimCLR checkpoint if provided
+        if simclr_ckpt_path and os.path.exists(simclr_ckpt_path):
+            ckpt = torch.load(simclr_ckpt_path, map_location="cpu", weights_only=False)
+            state_dict = ckpt.get("state_dict", ckpt)
+            backbone_state_dict = {}
+            for key, value in state_dict.items():
+                if key.startswith("backbone."):
+                    new_key = key[len("backbone."):]
+                    backbone_state_dict[new_key] = value
+            missing, unexpected = self.backbone.load_state_dict(backbone_state_dict, strict=False)
+            print(f"Loaded SimCLR backbone weights. Missing: {len(missing)}, Unexpected: {len(unexpected)}")
+        else:
+            print("Warning: SimCLR checkpoint not found or not provided. Using randomly initialized backbone.")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        features = self.backbone(x)
+        cls_token = features[0][:, 0]
+        return cls_token
+
+
+class Classifier(nn.Module):
+    def __init__(self, d_model: int = 768, num_classes: int = 1):
+        super().__init__()
+        self.fc = nn.Linear(d_model, num_classes)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc(x)
+
+
+class SingleScanModelBP(nn.Module):
+    def __init__(self, backbone: nn.Module, classifier: nn.Module):
+        super().__init__()
+        self.backbone = backbone
+        self.classifier = classifier
+        self.dropout = nn.Dropout(p=0.2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x shape: (batch_size, 2, C, D, H, W)
+        scan_features_list = []
+        for scan_tensor_with_extra_dim in x.split(1, dim=1):
+            squeezed_scan_tensor = scan_tensor_with_extra_dim.squeeze(1)
+            feature = self.backbone(squeezed_scan_tensor)
+            scan_features_list.append(feature)
+        stacked_features = torch.stack(scan_features_list, dim=1)
+        merged_features = torch.mean(stacked_features, dim=1)
+        merged_features = self.dropout(merged_features)
+        output = self.classifier(merged_features)
+        return output