Diff of /n4correction.py [000000] .. [e44b03]

Switch to side-by-side view

--- a
+++ b/n4correction.py
@@ -0,0 +1,36 @@
+from __future__ import division
+import os, sys, glob
+import numpy as np
+import nibabel as nib
+from skimage.transform import resize
+from multiprocessing import Pool, cpu_count
+
+def n4_correction(im_input):
+    command = 'N4BiasFieldCorrection -d 3 -i ' + im_input + ' ' + ' -s 3 -c [50x50x30x20] -b [300] -o ' +  im_input.replace('.nii.gz', '_corrected.nii.gz')
+    os.system(command)
+
+def batch_works(k):
+    if k == n_processes - 1:
+        paths = all_paths[k * int(len(all_paths) / n_processes) : ]
+    else:
+        paths = all_paths[k * int(len(all_paths) / n_processes) : (k + 1) * int(len(all_paths) / n_processes)]
+        
+    for path in paths:
+        n4_correction(glob.glob(os.path.join(path, '*_t1.nii.gz'))[0])
+        n4_correction(glob.glob(os.path.join(path, '*_t1ce.nii.gz'))[0])
+        n4_correction(glob.glob(os.path.join(path, '*_t2.nii.gz'))[0])
+        n4_correction(glob.glob(os.path.join(path, '*_flair.nii.gz'))[0])
+    
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        raise Exception("Need at least the input data directory")
+    input_path = sys.argv[1]
+        
+    all_paths = []
+    for dirpath, dirnames, files in os.walk(input_path):
+        if os.path.basename(dirpath)[0:7] == 'Brats17':
+            all_paths.append(dirpath)
+            
+    n_processes = cpu_count()
+    pool = Pool(processes=n_processes)
+    pool.map(batch_works, range(n_processes))
\ No newline at end of file