Add BrainIAC Glioma Segmentation app with proper Docker setup

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,28 @@
 ---
 title: BrainIAC Glioma Segmentation
-emoji: 🏃
-colorFrom: red
-colorTo: yellow
+emoji: 🧠
+colorFrom: blue
+colorTo: red
 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: Glioma Segmentation
+
+A Vision Transformer UNETR model for glioma segmentation from FLAIR MRI scans.
+
+## Features
+- Upload FLAIR MRI NIfTI files
+- Optional preprocessing (debiasing + registration + skull stripping)
+- Interactive slice-by-slice visualization  
+- Download preprocessed images and segmentation masks
+- Real-time segmentation statistics
+
+## Usage
+1. Upload a FLAIR MRI scan (.nii or .nii.gz)
+2. Optionally enable preprocessing
+3. Adjust segmentation threshold
+4. View results and download files
+
+*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]))