HD-BET for windows

HD_BET/data_loading.py

@@ -84,7 +84,7 @@ def save_segmentation_nifti(segmentation, dct, out_fname, order=1, dtype=np.uint
                      bbox[2][0]:bbox[2][1]] = segmentation
     else:
         seg_old_size = segmentation
-    if np.any([i != j for i, j in zip(np.array(seg_old_size), np.array(dct['size'])[[2, 1, 0]])]):
+    if np.any(np.array(seg_old_size.shape) != 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

HD_BET/hd-bet.cmd

@@ -0,0 +1,2 @@
+@echo off
+python "%~dp0\hd-bet.py" %*

HD_BET/hd-bet.py

@@ -0,0 +1,119 @@
+import os
+from HD_BET.run import run_hd_bet
+from HD_BET.utils import maybe_mkdir_p, subfiles
+import HD_BET
+
+
+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()
+
+    input_file_or_dir = args.input
+    output_file_or_dir = args.output
+
+    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")
+
+    mode = args.mode
+    device = args.device
+    tta = args.tta
+    pp = args.pp
+    save_mask = args.save_mask
+    overwrite_existing = args.overwrite_existing
+
+    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='.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)

HD_BET/paths.py

@@ -1,4 +1,4 @@
 import os
 
 # please refer to the readme on where to get the parameters. Save them in this folder:
-folder_with_parameter_files = os.path.join(os.path.expanduser('~'), 'hd-bet_params')
+folder_with_parameter_files = os.path.join(os.path.expanduser('~'), '.local', 'share', 'hd-bet', 'hd-bet_params')

HD_BET/run.py

@@ -3,7 +3,7 @@
 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
+import importlib.util
 from HD_BET.utils import postprocess_prediction, SetNetworkToVal, get_params_fname, maybe_download_parameters
 import os
 import HD_BET
@@ -53,7 +53,9 @@ def run_hd_bet(mri_fnames, output_fnames, mode="accurate", config_file=os.path.j
 
     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)
+    spec = importlib.util.spec_from_file_location("cf", config_file)
+    cf = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(cf)
     cf = cf.config()
 
     net, _ = cf.get_network(cf.val_use_train_mode, None)

requirements.txt

@@ -1,4 +1,4 @@
-numpy>=1.14.5
+numpy>=1.14.5, <1.25.0
 torch>=0.4.0
 scikit-image>=0.14.0
 SimpleITK>=2.0.2

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 
 setup(name='HD_BET',
-      version='1.0',
+      version='1.01',
       description='Tool for brain extraction',
       url='https://github.com/MIC-DKFZ/hd-bet',
       python_requires='>=3.5',
@@ -15,13 +15,15 @@
       'scikit-image',
       'SimpleITK'
       ],
-      scripts=['HD_BET/hd-bet'],
+      scripts=['HD_BET/hd-bet',
+              'HD_BET/hd-bet.cmd',
+              'HD_BET/hd-bet.py'],
       packages=find_packages(include=['HD_BET']),
       classifiers=[
           'Intended Audience :: Science/Research',
           'Programming Language :: Python',
           'Topic :: Scientific/Engineering',
-          'Operating System :: Unix'
+          'Operating System :: Unix/Windows'
       ]
       )