--- 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)