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