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

Switch to side-by-side view

--- a
+++ b/code/p1processing.py
@@ -0,0 +1,161 @@
+import os
+import numpy as np
+from multiprocessing import Pool
+from functools import partial
+from image_utils import * 
+    
+
+def output2DRefinement(atlases, DLSeg, param_dir, tmps_dir, dofs_dir, subject_dir, savedInd, fr, mirtk):
+   
+    segstring = ''
+    ind = 0
+    atlasUsedNo = len(atlases)
+        
+    for i in range(atlasUsedNo):
+         
+        if mirtk:
+            os.system('mirtk register '
+                      '{0} '
+                      '{1} '
+                      '-parin {2}/ffd_label_.cfg ' 
+                      '-dofin {3}/shapelandmarks_{4}.dof.gz '
+                      '-dofout {3}/shapeffd_{5}_{6}.dof.gz' 
+                      .format(DLSeg, atlases[i], param_dir, dofs_dir, savedInd[i], i, fr))
+            
+            os.system('mirtk transform-image '
+                      '{0} '
+                      '{1}/seg_lvsa_{4}_{5}.nii.gz ' 
+                      '-dofin {2}/shapeffd_{4}_{5}.dof.gz '
+                      '-target {3}/lvsa_{5}.nii.gz -interp NN'  
+                      .format(atlases[i], tmps_dir, dofs_dir, subject_dir, i, fr)) 
+
+        else:
+            os.system('nreg '
+                      '{0} '
+                      '{1} '
+                      '-parin {2}/segreg.txt ' 
+                      '-dofin {3}/shapelandmarks_{4}.dof.gz '
+                      '-dofout {3}/shapeffd_{5}_{6}.dof.gz' 
+                      .format(DLSeg, atlases[i], param_dir, dofs_dir, savedInd[i], i, fr))
+                          
+            os.system('transformation '
+                      '{0} '
+                      '{1}/seg_lvsa_{4}_{5}.nii.gz ' 
+                      '-dofin {2}/shapeffd_{4}_{5}.dof.gz '
+                      '-target {3}/lvsa_{5}.nii.gz -nn' 
+                      .format(atlases[i], tmps_dir, dofs_dir, subject_dir, i, fr))    
+
+        segstring += '{0}/seg_lvsa_{1}_{2}.nii.gz '.format(tmps_dir, i, fr)
+                
+        ind += 1
+        
+    # apply label fusion 
+    os.system('combineLabels {0}/seg_lvsa_{1}.nii.gz {2} {3}'.format(subject_dir, fr, ind, segstring))
+
+    
+def apply_PC(subject, data_dir, param_dir, atlases_list, landmarks_list, mirtk):
+    
+    print('  registering {0}'.format(subject))    
+    
+    subject_dir = os.path.join(data_dir, subject)
+        
+    if os.path.isdir(subject_dir):
+        
+        tmps_dir = '{0}/tmps'.format(subject_dir)
+
+        dofs_dir = '{0}/dofs'.format(subject_dir)
+    
+        segs_dir = '{0}/segs'.format(subject_dir)
+        
+        subject_landmarks = '{0}/landmarks.vtk'.format(subject_dir)
+                                  
+        for fr in ['ED', 'ES']:
+                
+            DLSeg = '{0}/seg_lvsa_{1}.nii.gz'.format(segs_dir, fr)
+            
+            if not os.path.exists(DLSeg):
+                
+                print(' segmentation {0} does not exist. Skip.'.format(DLSeg))
+                
+                continue
+                 
+            topSimilarAtlases_list, savedInd = topSimilarAtlasShapeSelection(atlases_list[fr], landmarks_list[fr], 
+                                                subject_landmarks, tmps_dir, dofs_dir, DLSeg, param_dir, 3) 
+                                      
+            output2DRefinement(topSimilarAtlases_list, DLSeg, param_dir, tmps_dir, dofs_dir, subject_dir, savedInd, fr, mirtk)
+                
+            refineFusionResults(subject_dir, 'seg_lvsa_{0}.nii.gz'.format(fr), 2)       
+            
+        print('  finish 2D nonrigid-registering one subject {}'.format(subject))
+    
+    else:  
+        print('  {0} is not a valid directory, do nothing'.format(subject_dir))
+        
+       
+def multiatlasreg2D(dir_0, dir_1, dir_2, coreNo, parallel, mirtk):
+               
+    print('Select all the shape atlases for 2D multi-atlas registration')
+    
+    atlases_list, landmarks_list = allAtlasShapeSelection(dir_1)
+
+    if parallel:
+    
+        print('Start 2D multi-atlas registration and program 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, 
+                         atlases_list=atlases_list, 
+                         landmarks_list=landmarks_list,
+                         mirtk=mirtk), 
+                         sorted(os.listdir(dir_0)))       
+                
+    else:
+        
+        print('Start 2D multi-atlas registration and program running subsequently')
+                
+        data_dir, param_dir = dir_0, dir_2
+                                     
+        for subject in sorted(os.listdir(data_dir)):
+            
+            print('  registering {0}'.format(subject))   
+            
+            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 
+          
+            tmps_dir = '{0}/tmps'.format(subject_dir)
+
+            dofs_dir = '{0}/dofs'.format(subject_dir)
+    
+            segs_dir = '{0}/segs'.format(subject_dir)
+            
+            subject_landmarks = '{0}/landmarks.vtk'.format(subject_dir)
+            
+            for fr in ['ED', 'ES']:
+                    
+                DLSeg = '{0}/seg_lvsa_{1}.nii.gz'.format(segs_dir, fr)
+                
+                if not os.path.exists(DLSeg):
+                
+                    print(' segmentation {0} does not exist. Skip.'.format(DLSeg))
+                
+                    continue
+                 
+                topSimilarAtlases_list, savedInd = topSimilarAtlasShapeSelection(atlases_list[fr], landmarks_list[fr], 
+                                                   subject_landmarks, tmps_dir, dofs_dir, DLSeg, param_dir, 3) 
+                                      
+                output2DRefinement(topSimilarAtlases_list, DLSeg, param_dir, tmps_dir, dofs_dir, subject_dir, savedInd, fr, mirtk)
+                
+                refineFusionResults(subject_dir, 'seg_lvsa_{0}.nii.gz'.format(fr), 2)       
+                              
+ 
+            print('  finish 2D nonrigid-registering one subject {}'.format(subject))
+                
\ No newline at end of file