finalising + testing claustrum part

README.md

@@ -8,7 +8,7 @@ WMH segmentation can be performed either using FLAIR only or both FLAIR and T1 s
 Fig.1: Segmentation result on Singapore 34 [1]. From left to right: FLAIR MR image, the associated ground truth, segmentation result using FLAIR modality only and segmentation result using FLAIR and T1 modalities. In column SegF LAIR and SegF LAIR+T1, the green area is the overlap between the segmentation maps and the ground-truth, the red pixels are the false negatives and the black ones are the false positives.
 
 ## Claustrum Segmentation
-For Claustrum Segmentation the T1 scan with .nii extension must be provided. 
+For Claustrum Segmentation the T1 scan must be provided. 
 
 # How to:
 Run deepNeuroSeg either in command line or Python.
@@ -19,7 +19,7 @@ deepNeuroSeg --type wmh --flair YOUR_PATH.nii.gz --t1 YOUR_PATH.nii.gz --o YOUR_
 ```
 Or: 
 ```ruby
-deepNeuroSeg --type c --t1 YOUR_PATH.nii --o YOUR_PATH
+deepNeuroSeg --type c --t1 YOUR_PATH.nii.gz --o YOUR_PATH
 ```
 
 For more details see:
@@ -28,13 +28,13 @@ deepNeuroSeg --help
 Options:
   --type [wmh|c]  Either 'wmh' (White Matter Hyperintensities) or 'c'
                   (Claustrum)
-  --flair PATH    Path to file of a FLAIR scan.
-  --t1 PATH       Path to file of a T1 scan.
-  --o TEXT        Directory path where to save the resulting segmentation.
+  --flair PATH    Path to nii.gz file with a FLAIR scan.
+  --t1 PATH       Path to nii.gz file with a T1 scan.
+  --o TEXT        Path where to save the resulting segmentation. Directory path or specific nii.gz file path.
                   [required]
   --help          Show this message and exit.
 ```
-The resulting mask will be saved in the user-specified directory under the name <code>out_mask.nii.gz</code>.
+The resulting mask will be saved with user-specified .nii.gz file name or in the user-specified directory under the name <code>out_mask.nii.gz</code>.
 
 ## Python
 In Python user will have to follow the next steps:
@@ -56,7 +56,7 @@ segmenter = SegmentationFactory.create_segmenter(SegmentationType.Claustrum,
 
 3. Next the segmentation can be performed.
 
-Option 1: The user can specify the output directory directly in <code>perform_segmentation</code> method.
+Option 1: The user can specify the output path directly in <code>perform_segmentation</code> method.
 ```ruby
 prediction = segmenter.perform_segmentation(outputDir='YOUR_PATH')
 ```
@@ -66,7 +66,16 @@ Option 2: The output numpy array can be inspected first, and then saved with <co
 prediction = segmenter.perform_segmentation()
 segmenter.save_segmentation(mask=prediction, outputDir='YOUR_PATH')
 ```
-In both cases, the output mask will be saved to a specified directory under the name <code>out_mask.nii.gz</code>.
+In both cases, the output mask will be saved with user-specified .nii.gz file name or in user-specified directory under the name <code>out_mask.nii.gz</code>.
+
+**Special feature of Claustrum Segmentation:**
+
+The user can check the orientation of the coronal and axial images by selecting the special feature in <code>perform_segmentation</code> method:
+```ruby
+prediction = segmenter.perform_segmentation(check_orientation=True)
+```
+<code>check_orientation=True</code> will save the coronal and axial images under ~/.deepNeuroSeg/images/.
+
 
 # References:
 

deepNeuroSeg/__main__.py

@@ -3,9 +3,9 @@
 
 @click.command()
 @click.option('--type', type=click.Choice(['wmh', 'c'], case_sensitive=True),  help="Either 'wmh' (White Matter Hyperintensities) or 'c' (Claustrum)")
-@click.option('--flair', help="Path to file of a FLAIR scan.", required=False, default=None, type=click.Path(exists=True))
-@click.option('--t1', help="Path to file of a T1 scan.", required=False, default=None, type=click.Path(exists=True))
-@click.option('--o', help="Directory path where to save the resulting segmentation.", required=True)
+@click.option('--flair', help="Path to nii.gz file with a FLAIR scan.", required=False, default=None, type=click.Path(exists=True))
+@click.option('--t1', help="Path to nii.gz file with a T1 scan.", required=False, default=None, type=click.Path(exists=True))
+@click.option('--o', help="Path where to save the resulting segmentation. Directory path or specific nii.gz file path.", required=True)
 def run(type, flair, t1, o):
     if type=='wmh':
         if flair is None:
@@ -15,11 +15,11 @@ def run(type, flair, t1, o):
         if t1 is not None and not t1.endswith('.nii.gz'):
             raise NameError('Invalide T1 file expension. Must end with .nii.gz')
         segmenter = SegmentationFactory.create_segmenter(SegmentationType.WMH, FLAIR_path=flair, T1_path=t1)
-        _ = segmenter.perform_segmentation(outputDir=o)
+        _ = segmenter.perform_segmentation(outputPath=o)
     elif type=='c':
         if t1 is None:
             raise TypeError('T1 scan is needed for \'c\' (Claustrum) Segmentation.')
-        if not t1.endswith('.nii'):
-            raise NameError('Invalide T1 file expension. Must end with .nii')
+        if not t1.endswith('.nii.gz'):
+            raise NameError('Invalide T1 file expension. Must end with .nii.gz')
         segmenter = SegmentationFactory.create_segmenter(SegmentationType.Claustrum, T1_path=t1)
-        _ = segmenter.perform_segmentation(outputDir=o)
+        _ = segmenter.perform_segmentation(outputPath=o)

deepNeuroSeg/claustrum_segmenter.py

@@ -1,6 +1,7 @@
 import os
 import scipy
 import numpy as np
+from PIL import Image
 from tensorflow import keras
 from tensorflow.keras.models import Model
 from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, Cropping2D, ZeroPadding2D, Activation
@@ -23,14 +24,17 @@ class ClaustrumSegmentation(AbstractSegmenter):
     def __init__(self, T1_path):
         self.T1_path = T1_path
 
+    def get_T1_path(self):
+        return self.T1_path
+
     def _get_links(self):
         return 'pretrained_T1_claustrum', ClaustrumSegmentation.c_dict['pretrained_T1_claustrum']
 
-    def perform_segmentation(self, outputDir=None):
+    def perform_segmentation(self, outputPath=None, check_orientation=False):
         """Performs claustrum segmentation by loading required models from ./~deepNeuroSeg/pretrained_T1_claustrum cache directory.
 
         Args:
-            outputDir (str, optional): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz. Defaults to None meaning not saved.
+            outputPath (str, optional): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz. Defaults to None meaning not saved.
 
         Returns:
             numpy.ndarray: the predicted mask.
@@ -51,6 +55,30 @@ def perform_segmentation(self, outputDir=None):
         corona_array =  self.pre_processing(corona_array, img_shape)
         axial_array =  self.pre_processing(axial_array, img_shape)
 
+        if check_orientation:
+            #this is to check the orientation of your images is right or not. Please check /images/coronal and /images/axial
+            cache_dir = os.path.realpath(os.path.expanduser('~/.deepNeuroSeg'))
+            image_path = os.path.join(cache_dir,'images')
+            direction_1 = 'coronal'
+            direction_2 = 'axial'
+            if not os.path.exists(image_path):
+                os.makedirs(image_path)
+            if not os.path.exists(os.path.join(image_path, direction_1)):
+                os.makedirs(os.path.join(image_path, direction_1))
+            if not os.path.exists(os.path.join(image_path, direction_2)):
+                os.makedirs(os.path.join(image_path, direction_2))
+            for ss in range(np.shape(corona_array)[0]):
+                slice_ = 255*(corona_array[ss] - np.min(corona_array[ss]))/(np.max(corona_array[ss]) - np.min(corona_array[ss]))
+                np_slice_ = np.squeeze(slice_, axis=2)
+                im = Image.fromarray(np.uint8(np_slice_))
+                im.save(os.path.join(image_path, direction_1, str(ss)+'.png'))
+
+            for ss in range(np.shape(axial_array)[0]):
+                slice_ = 255*(axial_array[ss] - np.min(axial_array[ss]))/(np.max(axial_array[ss]) - np.min(axial_array[ss]))
+                np_slice_ = np.squeeze(slice_, axis=2)
+                im = Image.fromarray(np.uint8(np_slice_))
+                im.save(os.path.join(image_path, direction_2, str(ss)+'.png')) 
+
         pred_a, pred_c = self.predict(axial_array, corona_array, self.get_unet(img_shape))
 
         # transform them to their original size and orientations
@@ -62,22 +90,29 @@ def perform_segmentation(self, outputDir=None):
         pred[pred > 0.40] = 1.
         pred[pred <= 0.40] = 0.
 
-        if outputDir:
+        if outputPath:
             #save the masks
-            self.save_segmentation(pred, outputDir)
+            self.save_segmentation(pred, outputPath)
 
         return pred
 
-    def save_segmentation(self, mask, outputDir):
+    def save_segmentation(self, mask, outputPath):
         """Saves provided mask as out_mask.nii.gz in the given directory.
 
         Args:
             mask (numpy.ndarray): the mask.
-            outputDir ([type]): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz
+            outputPath ([type]): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz
         """
-        if not os.path.exists(outputDir):
-            os.mkdir(outputDir)
-        filename_resultImage = os.path.join(outputDir,'out_mask.nii.gz')
+        if os.path.isdir(outputPath):
+            if not os.path.exists(outputPath):
+                os.mkdir(outputPath)
+            filename_resultImage = os.path.join(outputPath,'out_mask.nii.gz')
+        else:
+            if outputPath.endswith('nii.gz'):
+                filename_resultImage = outputPath
+            else:
+                raise NameError('Invalide file expension. Must end with .nii.gz')
+
         img_out = sitk.GetImageFromArray(mask)
         sitk.WriteImage(img_out, filename_resultImage)
 

deepNeuroSeg/wmh_segmenter.py

@@ -27,33 +27,45 @@ def __init__(self, FLAIR_path, T1_path=None):
         self.FLAIR_path = FLAIR_path
         self.T1_path = T1_path
 
-    def perform_segmentation(self, outputDir=None):
+    def get_FLAIR_path(self):
+        return self.FLAIR_path
+
+    def get_T1_path(self):
+        return self.T1_path
+
+    def perform_segmentation(self, outputPath=None):
         """Performs segmentation by loading three required models from ./~deepNeuroSeg cache directory.
 
         Args:
-            outputDir (str, optional): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz. Defaults to None meaning not saved.
+            outputPath (str, optional): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz. Defaults to None meaning not saved.
 
         Returns:
             numpy.ndarray: the predicted mask.
         """
         img_shape, imgs_test, model_dir, FLAIR_array = read_data(self.FLAIR_path, self.T1_path)
         original_pred = load_model(img_shape, imgs_test, model_dir, FLAIR_array)
 
-        if outputDir:
-            self.save_segmentation(original_pred, outputDir)
+        if outputPath:
+            self.save_segmentation(original_pred, outputPath)
 
         return original_pred
 
-    def save_segmentation(self, mask, outputDir):
+    def save_segmentation(self, mask, outputPath):
         """Saves provided mask as out_mask.nii.gz in the given directory.
 
         Args:
             mask (numpy.ndarray): the mask.
-            outputDir ([type]): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz
+            outputPath ([type]): the desired directory path where the resulting mask will be saved under the name out_mask.nii.gz
         """
-        if not os.path.exists(outputDir):
-            os.mkdir(outputDir)
-        filename_resultImage = os.path.join(outputDir,'out_mask.nii.gz')
+        if os.path.isdir(outputPath):
+            if not os.path.exists(outputPath):
+                os.mkdir(outputPath)
+            filename_resultImage = os.path.join(outputPath,'out_mask.nii.gz')
+        else:
+            if outputPath.endswith('nii.gz'):
+                filename_resultImage = outputPath
+            else:
+                raise NameError('Invalide file expension. Must end with .nii.gz')
         img_out = sitk.GetImageFromArray(mask)
         FLAIR_image = sitk.ReadImage(self.FLAIR_path)
         img_out.CopyInformation(FLAIR_image) #copy the meta information (voxel size, etc.) from the input raw image

setup.py

@@ -27,6 +27,7 @@
         'scipy==1.6.3',
         'pytest==6.2.4',
         'click==8.0.0',
+        'Pillow==8.2.0',
     ],
     entry_points={
         'console_scripts': [

tests/test_segmenters.py

@@ -0,0 +1,35 @@
+import pytest
+import os
+import deepNeuroSeg
+from deepNeuroSeg import SegmentationFactory, SegmentationType
+
+def test_SegmentationFactory_SegmentationTypeWMH():
+    test_path = '../test_your_data_WMH/input_dir/FLAIR.nii.gz'
+    concrete_strategy_a = SegmentationFactory.create_segmenter(SegmentationType.WMH,  FLAIR_path=test_path)
+    assert isinstance(concrete_strategy_a, deepNeuroSeg.wmh_segmenter.WMHSegmentation)
+
+def test_SegmentationFactory_SegmentationTypeClaustrum():
+    test_path = '../Claustrum_Demo_GitHub/data/sub-0051456/sub-0051456_T1w_denoised.nii'
+    concrete_strategy_b = SegmentationFactory.create_segmenter(SegmentationType.Claustrum,  T1_path=test_path)
+    assert isinstance(concrete_strategy_b, deepNeuroSeg.claustrum_segmenter.ClaustrumSegmentation)
+
+def test_SegmentationFactory_WNH_initialization():
+    test_path = '../test_your_data_WMH/input_dir/FLAIR.nii.gz'
+    concrete_strategy_a = SegmentationFactory.create_segmenter(SegmentationType.WMH, FLAIR_path=test_path)
+    assert concrete_strategy_a.get_FLAIR_path() == test_path and concrete_strategy_a.get_T1_path() == None
+
+def test_wmh_segmenter():
+    concrete_strategy_a = SegmentationFactory.create_segmenter(SegmentationType.WMH,
+                                                               FLAIR_path='../test_your_data_WMH/input_dir/FLAIR.nii.gz',
+                                                               T1_path='../test_your_data_WMH/input_dir/T1.nii.gz')
+    pred_image = concrete_strategy_a.perform_segmentation()
+    test_path = '../test_your_data_WMH/result/test.nii.gz'
+    concrete_strategy_a.save_segmentation(mask=pred_image, outputPath=test_path)
+    assert os.path.exists(test_path)
+
+def test_claustrum_segmenter():
+    concrete_strategy_b = SegmentationFactory.create_segmenter(SegmentationType.Claustrum,
+                                                               T1_path='/Users/rita/Uni/pmsd/Claustrum_Demo_GitHub/data/sub-0051456/sub-0051456_T1w_denoised.nii')
+    _ = concrete_strategy_b.perform_segmentation(check_orientation=True)
+    test_path = os.path.join( os.path.realpath(os.path.expanduser('~/.deepNeuroSeg')), 'images')
+    assert os.path.exists(test_path) and not len(test_path) == 0