Diff of /code/motionEstimation.py [000000] .. [32d261]

Switch to side-by-side view

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