--- a +++ b/code/motionEstimation.py @@ -0,0 +1,175 @@ +#!/usr/bin/python +import os, glob +from multiprocessing import Pool +from functools import partial + + +# Perform motion tracking for cine MR +def track_cine(data_dir, par_dir, tamplate_dir): + + motion_dir = os.path.join(data_dir, 'motion') + if os.path.exists(motion_dir): + os.system('rm -r {0}'.format(motion_dir)) + os.mkdir(motion_dir) + else: + os.mkdir(motion_dir) + + # Split 4D nifti into time phases + os.system('splitvolume ' + '{0}/lvsa_.nii.gz ' + '{1}/lvsa_ ' + '-sequence' + .format(data_dir, motion_dir)) + + n_frame = len(glob.glob('{0}/lvsa_??.nii.gz'.format(motion_dir))) + + # Inter-frame motion estimation + for fr in range(1, n_frame): + target = '{0}/lvsa_{1:02d}.nii.gz'.format(motion_dir, fr-1) + source = '{0}/lvsa_{1:02d}.nii.gz'.format(motion_dir, fr) + par = '{0}/ffd_motion.cfg'.format(par_dir) + dof = '{0}/ffd_{1:02d}_to_{2:02d}.dof.gz'.format(motion_dir, fr-1, fr) + + if not os.path.exists(dof): + os.system('mirtk register ' + '{0} ' + '{1} ' + '-parin {2} ' + '-dofout {3}' + .format(target, source, par, dof)) + + # Compose inter-frame transformation fields + for fr in range(2, n_frame): + dofs = '' + for k in range(1, fr+1): + dof = '{0}/ffd_{1:02d}_to_{2:02d}.dof.gz'.format(motion_dir, k-1, k) + dofs += dof + ' ' + dof_out = '{0}/ffd_comp_00_to_{1:02d}.dof.gz'.format(motion_dir, fr) + if not os.path.exists(dof_out): + os.system('mirtk compose-dofs ' + '{0} ' + '{1}' + .format(dofs, dof_out)) + + # Refine motion fields + # Composition of inter-frame motion fields can lead to accumulative errors. At this step, we refine the motion fields + # by re-registering the n-th frame with the ED frame. + for fr in range(2, n_frame): + target = '{0}/lvsa_00.nii.gz'.format(motion_dir) + source = '{0}/lvsa_{1:02d}.nii.gz'.format(motion_dir, fr) + par = '{0}/ffd_refine.cfg'.format(par_dir) + dofin = '{0}/ffd_comp_00_to_{1:02d}.dof.gz'.format(motion_dir, fr) + dof = '{0}/ffd_00_to_{1:02d}.dof.gz'.format(motion_dir, fr) + if not os.path.exists(dof): + os.system('mirtk register ' + '{0} ' + '{1} ' + '-parin {2} ' + '-dofin {3} ' + '-dofout {4}' + .format(target, source, par, dofin, dof)) + + # Obtain the RV mesh with the same number of points as the template mesh + os.system('srreg ' + '{2}/vtks/RV_ED.vtk ' + '{0}/RVendo_ED.vtk ' + '-dofout {1}/RV_srreg.dof.gz ' + '-symmetric' + .format(tamplate_dir, motion_dir, data_dir)) + + os.system('mirtk transform-points ' + '{0}/RVendo_ED.vtk ' + '{1}/RV_ED_srreg.vtk ' + '-dofin {1}/RV_srreg.dof.gz ' + '-invert' + .format(tamplate_dir, motion_dir)) + + os.system('sareg ' + '{0}/RV_ED_srreg.vtk ' + '{1}/vtks/RV_ED.vtk ' + '-dofout {0}/RV_sareg.dof.gz ' + '-symmetric' + .format(motion_dir, data_dir)) + + os.system('snreg ' + '{0}/RV_ED_srreg.vtk ' + '{1}/vtks/RV_ED.vtk ' + '-dofin {0}/RV_sareg.dof.gz ' + '-dofout {0}/RV_snreg.dof.gz ' + '-ds 20 -symmetric' + .format(motion_dir, data_dir)) + + os.system('mirtk transform-points ' + '{0}/RV_ED_srreg.vtk ' + '{0}/RV_fr00.vtk ' + '-dofin {0}/RV_snreg.dof.gz' + .format(motion_dir)) + + # Transform the mesh + for fr in range(1, n_frame): + os.system('mirtk transform-points ' + '{0}/RV_fr00.vtk ' + '{0}/RV_fr{1:02d}.vtk ' + '-dofin {0}/ffd_00_to_{1:02d}.dof.gz' + .format(motion_dir, fr)) + + # Convert vtks to text files + for fr in range(0, n_frame): + os.system('vtk2txt ' + '{0}/RV_fr{1:02d}.vtk ' + '{0}/RV_fr{1:02d}.txt' + .format(motion_dir, fr)) + + +def apply_PC(subject, data_dir, param_dir, template_dir): + + print(' co-registering {0}'.format(subject)) + + subject_dir = os.path.join(data_dir, subject) + + if os.path.isdir(subject_dir): + + track_cine(subject_dir, param_dir, template_dir) + + print(' finish motion tracking in subject {0}'.format(subject)) + + else: + print(' {0} is not a valid directory, do nothing'.format(subject_dir)) + + +def motionTracking(dir_0, dir_2, dir_3, coreNo, parallel): + + if parallel: + + print('Motion tracking running on {0} cores'.format(coreNo)) + + pool = Pool(processes = coreNo) + + # partial only in Python 2.7+ + pool.map(partial(apply_PC, + data_dir=dir_0, + param_dir=dir_2, + template_dir=dir_3), + sorted(os.listdir(dir_0))) + + else: + + print('Motion tracking from segmentations running subsequently') + + data_dir, param_dir, template_dir = dir_0, dir_2, dir_3 + + for subject in sorted(os.listdir(data_dir)): + + subject_dir = os.path.join(data_dir, subject) + + if not os.path.isdir(subject_dir): + + print(' {0} is not a valid folder, Skip'.format(subject_dir)) + + continue + + print(' motion tracking {0}'.format(subject)) + + track_cine(subject_dir, param_dir, template_dir) + + print(' finish motion tracking in {0}'.format(subject))