Diff of /src/registration.py [000000] .. [e918fa]

Switch to side-by-side view

--- a
+++ b/src/registration.py
@@ -0,0 +1,114 @@
+#!/usr/bin/env python
+import os
+from multiprocessing import Process, Manager, Value, Pool
+import multiprocessing as mp
+import subprocess
+import argparse
+import pathlib
+import sys
+
+
+def ANTS_rigid_body_trans(b0_nii, result, mask_file, reference):
+
+    print("Performing ants rigid body transformation...")
+    input_file = b0_nii
+    case_name = os.path.basename(input_file)
+    output_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-'
+    output_file = os.path.join(os.path.dirname(input_file), output_name)
+
+    trans_matrix = "antsRegistrationSyNQuick.sh -d 3 -f " + reference + " -m " + input_file + " -t r -o " + output_file
+    output1 = subprocess.check_output(trans_matrix, shell=True)
+
+    omat_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-0GenericAffine.mat'
+    omat_file = os.path.join(os.path.dirname(input_file), omat_name)
+
+    output_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-Warped.nii.gz'
+    transformed_file = os.path.join(os.path.dirname(input_file), output_name)
+
+    result.append((transformed_file, omat_file, mask_file))
+
+SUFFIX_TXT = "txt"
+SUFFIX_NIFTI_GZ = "nii.gz"
+parser = argparse.ArgumentParser()
+parser.add_argument('-b0', action='store', dest='b0', type=str,
+                        help="txt file containing list of /path/to/b0, one path in each line")
+parser.add_argument('-mask', action='store', dest='mask', type=str,
+                        help="txt file containing list of /path/to/mask, one path in each line")
+parser.add_argument('-ref', action='store', dest='ref', type=str,
+                        help="reference b0 file for registration")
+
+args = parser.parse_args()
+reference = str(args.ref)
+
+try:
+    args = parser.parse_args()
+    if len(sys.argv) == 1:
+        parser.print_help()
+        parser.error('too few arguments')
+        sys.exit(0)
+
+except SystemExit:
+    sys.exit(0)
+
+if args.b0:
+    f = pathlib.Path(args.b0)
+    if f.exists():
+        print ("File exist")
+        filename = args.b0
+    else:
+        print ("File not found")
+        sys.exit(1)
+
+    # Input caselist.txt
+    if filename.endswith(SUFFIX_TXT):
+        with open(filename) as f:
+            target_list = f.read().splitlines()
+
+if args.mask:
+    f = pathlib.Path(args.mask)
+    if f.exists():
+        print ("File exist")
+        filename = args.mask
+    else:
+        print ("File not found")
+        sys.exit(1)
+
+    # Input caselist.txt
+    if filename.endswith(SUFFIX_TXT):
+        with open(filename) as f:
+            mask_list = f.read().splitlines()
+
+with Manager() as manager:
+    result = manager.list()              
+    ants_jobs = []
+    for i in range(0, len(target_list)):
+        p_ants = mp.Process(target=ANTS_rigid_body_trans, args=(target_list[i],
+                                                             result, mask_list[i], reference))
+        ants_jobs.append(p_ants)
+        p_ants.start()
+        
+    for process in ants_jobs:
+        process.join()
+
+    result = list(result)
+
+transformed_cases = []
+omat_list = []
+masks_new_list = []
+
+for subject_ANTS in result:
+    transformed_cases.append(subject_ANTS[0])
+    omat_list.append(subject_ANTS[1])
+    masks_new_list.append(subject_ANTS[2])
+
+# Apply the same tranformation to the mask file
+for i in range(0, len(transformed_cases)):
+
+    input_file = transformed_cases[i]
+    case_name = os.path.basename(input_file)
+    output_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-mask.nii.gz'
+    output_file = os.path.join(os.path.dirname(masks_new_list[i]), output_name)
+    apply_mask_trans = "antsApplyTransforms -d 3 -i " + masks_new_list[i] + " -r " + input_file + " -o " \
+                            + output_file + " --transform [" + omat_list[i] + "]"
+
+    output2 = subprocess.check_output(apply_mask_trans, shell=True)