Create skullstripping.py

captcha/preprocessing/skull_stripping.py

@@ -0,0 +1,69 @@
+"""
+File name: skull_stripping.py
+Author: ngocviendang
+Date created: July 13, 2020
+Updated: May 16, 2022
+
+This file removes the skull from the MRI images.
+
+Takes as input three .nii files : the original image, the mask, and the vessel_mask (which serves as the vessel labels).
+Outputs three .nii files : _img, _mask, _label with the skull removed.
+"""
+import argparse
+import sys
+import os
+import regex as re
+from captcha.utils import helper
+from captcha.utils.helper import getAllFiles
+
+def main(args):
+    original_data_dir = os.path.expanduser(args.original_data_dir)
+    target_dir = os.path.expanduser(args.target_dir)
+    if not os.path.exists(target_dir):
+        os.makedirs(target_dir)
+    # List filenames of data
+    unfiltered_filelist = getAllFiles(original_data_dir)
+    ## input_list is the full images in .nii format
+    input_list = [item for item in unfiltered_filelist if re.search('angio.nii', item)]
+    ## mask_list is the mask.nii files, the files with the brain masked out
+    mask_list = [item for item in unfiltered_filelist if re.search('\\\\mask.nii', item)]
+    # label_list is the vessel_mask files, the files that mask out the vessels
+    label_list = [item for item in unfiltered_filelist if re.search('vessel_mask.nii', item)]
+    input_list = sorted(input_list)
+    mask_list = sorted(mask_list)
+    label_list = sorted(label_list)
+    print(input_list)
+    print(mask_list)
+    print(label_list)
+    # load image, mask and label stacks as matrices	
+    for i,j in enumerate(input_list):
+        print('Loading image...')
+        img_mat = helper.load_nifti_mat_from_file(j)
+        print('Loading mask...')
+        mask_mat = helper.load_nifti_mat_from_file(mask_list[i])
+        print('Loading label...')
+        label_mat = helper.load_nifti_mat_from_file(label_list[i])      
+        # check the dimensions
+        assert img_mat.shape == mask_mat.shape == label_mat.shape, "The DIMENSIONS of image, mask and label are NOT " \
+                                                                   "SAME."
+
+        # mask images and labels (skull stripping)
+        img_mat = helper.aplly_mask(img_mat, mask_mat)
+        label_mat = helper.aplly_mask(label_mat, mask_mat)
+        # save to new file as masked version of original data 
+        helper.create_and_save_nifti(img_mat, target_dir + j.split(os.sep)[-1].split('_')[0] + '_img_NEW.nii')
+        helper.create_and_save_nifti(mask_mat, target_dir + j.split(os.sep)[-1].split('_')[0] + '_mask_NEW.nii')
+        helper.create_and_save_nifti(label_mat, target_dir + j.split(os.sep)[-1].split('_')[0] + '_label_NEW.nii')
+
+        print()
+    print('DONE')
+
+def parse_arguments(argv):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--original_data_dir", type=str,
+                        help='data dictionary.')
+    parser.add_argument("--target_dir", type=str, help='Directory for saving the images after the skull stripping process.')     
+    return parser.parse_args(argv)
+
+if __name__ == '__main__':
+    main(parse_arguments(sys.argv[1:]))