Merge pull request ellisdg#150 from ellisdg/update_to_brats18

Update modalities for brats18

Author: ellisdg

brats/preprocess.py

@@ -116,10 +116,17 @@ def normalize_image(in_file, out_file, bias_correction=True):
     return out_file
 
 
-def convert_brats_folder(in_folder, out_folder, truth_name="GlistrBoost_ManuallyCorrected",
-                         no_bias_correction_modalities=None):
+def convert_brats_folder(in_folder, out_folder, truth_name='seg', no_bias_correction_modalities=None):
     for name in config["all_modalities"]:
-        image_file = get_image(in_folder, name)
+        try:
+            image_file = get_image(in_folder, name)
+        except RuntimeError as error:
+            if name == 't1ce':
+                image_file = get_image(in_folder, 't1Gd')
+                truth_name = "GlistrBoost_ManuallyCorrected"
+            else:
+                raise error
+
         out_file = os.path.abspath(os.path.join(out_folder, name + ".nii.gz"))
         perform_bias_correction = no_bias_correction_modalities and name not in no_bias_correction_modalities
         normalize_image(image_file, out_file, bias_correction=perform_bias_correction)
@@ -128,6 +135,7 @@ def convert_brats_folder(in_folder, out_folder, truth_name="GlistrBoost_Manually
         truth_file = get_image(in_folder, truth_name)
     except RuntimeError:
         truth_file = get_image(in_folder, truth_name.split("_")[0])
+
     out_file = os.path.abspath(os.path.join(out_folder, "truth.nii.gz"))
     shutil.copy(truth_file, out_file)
     check_origin(out_file, get_image(in_folder, config["all_modalities"][0]))

brats/train.py

@@ -13,7 +13,7 @@
 config["patch_shape"] = (64, 64, 64)  # switch to None to train on the whole image
 config["labels"] = (1, 2, 4)  # the label numbers on the input image
 config["n_labels"] = len(config["labels"])
-config["all_modalities"] = ["t1", "t1Gd", "flair", "t2"]
+config["all_modalities"] = ["t1", "t1ce", "flair", "t2"]
 config["training_modalities"] = config["all_modalities"]  # change this if you want to only use some of the modalities
 config["nb_channels"] = len(config["training_modalities"])
 if "patch_shape" in config and config["patch_shape"] is not None:

brats/train_isensee2017.py

@@ -13,7 +13,7 @@
 config["labels"] = (1, 2, 4)  # the label numbers on the input image
 config["n_base_filters"] = 16
 config["n_labels"] = len(config["labels"])
-config["all_modalities"] = ["t1", "t1Gd", "flair", "t2"]
+config["all_modalities"] = ["t1", "t1ce", "flair", "t2"]
 config["training_modalities"] = config["all_modalities"]  # change this if you want to only use some of the modalities
 config["nb_channels"] = len(config["training_modalities"])
 if "patch_shape" in config and config["patch_shape"] is not None:

test/test_brats.py

@@ -0,0 +1,73 @@
+from unittest import TestCase
+import os
+import shutil
+import sys
+import nibabel as nib
+import numpy as np
+
+
+def create_blank_image(filename, image_shape=(140, 140, 140)):
+    data = np.zeros(image_shape, dtype=np.int16)
+    affine = np.diag(np.ones(4))
+    image = nib.Nifti1Image(dataobj=data, affine=affine)
+    image.to_filename(filename)
+
+
+class TestPreprocess(TestCase):
+    def setUp(self):
+        self.temp_brats_dir = os.path.abspath("temp_brats")
+        os.makedirs(self.temp_brats_dir)
+        sys.path.append('../brats')
+        self.preprocessed_dir = os.path.abspath("temp_preprocessed_brats")
+
+    def tearDown(self):
+        for directory in (self.temp_brats_dir, self.preprocessed_dir):
+            if os.path.exists(directory):
+                shutil.rmtree(directory)
+
+    def create_replica_dataset(self, subject_ids, scan_types, directory_name):
+        gbm_dir = os.path.join(self.temp_brats_dir, directory_name)
+        for subject_id in subject_ids:
+            subject_dir = os.path.join(gbm_dir, subject_id)
+            os.makedirs(subject_dir)
+            for scan_label in scan_types:
+                basename = '{}_{}.nii.gz'.format(subject_id, scan_label)
+                scan_filename = os.path.join(subject_dir, basename)
+                create_blank_image(scan_filename)
+
+    def create_replica_dataset_pre2018(self):
+        self.create_replica_dataset(subject_ids=('TCGA-00-000',),
+                                    scan_types=('flair',
+                                                'GlistrBoost',
+                                                'GlistrBoost_ManuallyCorrected',
+                                                't1',
+                                                't1Gd',
+                                                't2'),
+                                    directory_name='Pre-operative_TCGA_GBM_NIfTI_and_Segmentations')
+        self.create_replica_dataset(subject_ids=('TCGA-01-000',),
+                                    scan_types=('flair',
+                                                'GlistrBoost',
+                                                't1',
+                                                't1Gd',
+                                                't2'),
+                                    directory_name='Pre-operative_TCGA_GBM_NIfTI_and_Segmentations')
+
+    def create_replica_dataset_2018(self):
+        self.create_replica_dataset(subject_ids=('Brats18_1900_1_1',),
+                                    scan_types=('flair',
+                                                't1',
+                                                't1ce',
+                                                't2',
+                                                'seg'),
+                                    directory_name='HGG')
+
+    def test_preprocess_pre2018(self):
+        from preprocess import convert_brats_data
+        self.create_replica_dataset_pre2018()
+        convert_brats_data(self.temp_brats_dir, self.preprocessed_dir)
+
+    def test_preprocess_2018(self):
+        self.create_replica_dataset_2018()
+        from preprocess import convert_brats_data
+        convert_brats_data(self.temp_brats_dir, self.preprocessed_dir)
+

test/test_generator.py

@@ -188,8 +188,7 @@ def test_generator_with_permutations(self):
                                                             permute=True)
         training_generator, validation_generator, n_training_steps, n_validation_steps = generators
 
-        for x in training_generator:
-            break
+        _ = next(training_generator)
 
         self.rm_tmp_files()