371 lines
11 KiB
Python
371 lines
11 KiB
Python
|
|
||
|
|
'''
|
||
|
|
metric > -0.27 is bad registration
|
||
|
|
'''
|
||
|
|
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
# import pathlib
|
||
|
|
import shelve
|
||
|
|
import shutil
|
||
|
|
import time
|
||
|
|
|
||
|
|
import ants
|
||
|
|
import filelock
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import numpy as np
|
||
|
|
import SimpleITK as sitk
|
||
|
|
# import skimage
|
||
|
|
|
||
|
|
PATIENTS_ROOT = '/mnt/1220/Public/dataset2/G4'
|
||
|
|
OUT_ROOT = '/mnt/1220/Public/dataset2/H4'
|
||
|
|
SHELVE = os.path.join(OUT_ROOT, '0shelve')
|
||
|
|
|
||
|
|
MAX_Y = 256
|
||
|
|
|
||
|
|
SIZE_X = 249
|
||
|
|
SIZE_Y = 249
|
||
|
|
SIZE_Z = 192
|
||
|
|
# SIZE_Z = 256
|
||
|
|
|
||
|
|
MIN_OVERLAP = 0.50
|
||
|
|
MIN_METRIC = -0.50
|
||
|
|
|
||
|
|
# def resize_with_crop_or_pad(image, tx = SIZE_X, ty = SIZE_Y, tz = SIZE_Z):
|
||
|
|
def resize_with_pad(image, tx = SIZE_X, ty = SIZE_Y, tz = SIZE_Z):
|
||
|
|
sx, sy, sz = image.GetSize()
|
||
|
|
l = [(tx-sx)//2,
|
||
|
|
(ty-sy)//2,
|
||
|
|
(tz-sz)//2,]
|
||
|
|
u = [tx-sx-l[0],
|
||
|
|
ty-sy-l[1],
|
||
|
|
tz-sz-l[2],
|
||
|
|
]
|
||
|
|
# print (l, u)
|
||
|
|
return sitk.ConstantPad(image, l, u)
|
||
|
|
|
||
|
|
|
||
|
|
def draw_sitk(image, d, post):
|
||
|
|
a = sitk.GetArrayFromImage(image)
|
||
|
|
s = a.shape
|
||
|
|
|
||
|
|
fig, axs = plt.subplots(1, 3)
|
||
|
|
# fig.suptitle('%dx%dx%d'%(s[2], s[1], s[0]))
|
||
|
|
axs.flat[0].imshow(a[s[0]//2,:,:], cmap='gray')
|
||
|
|
axs.flat[1].imshow(a[:,s[1]//2,:], cmap='gray')
|
||
|
|
axs.flat[2].imshow(a[:,:,s[2]//2], cmap='gray')
|
||
|
|
axs.flat[0].axis('off')
|
||
|
|
axs.flat[1].axis('off')
|
||
|
|
axs.flat[2].axis('off')
|
||
|
|
axs.flat[1].invert_yaxis()
|
||
|
|
axs.flat[2].invert_yaxis()
|
||
|
|
plt.tight_layout()
|
||
|
|
os.makedirs(d, exist_ok=True)
|
||
|
|
plt.savefig(os.path.join(d, '%dx%dx%d-%s'%(s[2],s[1],s[0],post)))
|
||
|
|
plt.close()
|
||
|
|
# exit()
|
||
|
|
|
||
|
|
def bbox2_3D(img):
|
||
|
|
|
||
|
|
r = np.any(img, axis=(1, 2))
|
||
|
|
c = np.any(img, axis=(0, 2))
|
||
|
|
z = np.any(img, axis=(0, 1))
|
||
|
|
|
||
|
|
if not np.any(r):
|
||
|
|
return -1, -1, -1, -1, -1, -1
|
||
|
|
|
||
|
|
rmin, rmax = np.where(r)[0][[0, -1]]
|
||
|
|
cmin, cmax = np.where(c)[0][[0, -1]]
|
||
|
|
zmin, zmax = np.where(z)[0][[0, -1]]
|
||
|
|
|
||
|
|
return rmin, rmax, cmin, cmax, zmin, zmax
|
||
|
|
|
||
|
|
|
||
|
|
def registration(ct0, ct1, mr):
|
||
|
|
|
||
|
|
ct0n = ct0.numpy()
|
||
|
|
mrn = mr.numpy()
|
||
|
|
|
||
|
|
if np.array_equal(ct0n, mrn):
|
||
|
|
print('EQUAL')
|
||
|
|
return {
|
||
|
|
'fwdtransforms': [],
|
||
|
|
'warpedfixout': ct0,
|
||
|
|
'warpedmovout': ct0,
|
||
|
|
'ct': 0,
|
||
|
|
'type': 'Identity',
|
||
|
|
'metric': -2,
|
||
|
|
'ratio': 1,
|
||
|
|
}
|
||
|
|
|
||
|
|
if ct0n.shape == mrn.shape:
|
||
|
|
print('SAME SHAPE')
|
||
|
|
CTS = (ct0, ct1)
|
||
|
|
TYPES = (
|
||
|
|
# 'Translation',
|
||
|
|
'Rigid',
|
||
|
|
'QuickRigid',
|
||
|
|
'DenseRigid',
|
||
|
|
'BOLDRigid',
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
print('others', mrn.shape)
|
||
|
|
if min(mrn.shape) < 4:
|
||
|
|
print('skip')
|
||
|
|
return None
|
||
|
|
CTS = (ct0, ct1)
|
||
|
|
TYPES = (
|
||
|
|
# 'Translation',
|
||
|
|
'Rigid',
|
||
|
|
'QuickRigid',
|
||
|
|
'DenseRigid',
|
||
|
|
'BOLDRigid',
|
||
|
|
)
|
||
|
|
|
||
|
|
TX = []
|
||
|
|
start = time.time()
|
||
|
|
for m in range(len(CTS)):
|
||
|
|
for typ in TYPES:
|
||
|
|
ct = CTS[m]
|
||
|
|
print(typ)
|
||
|
|
|
||
|
|
mytx = ants.registration(ct, mr, typ)
|
||
|
|
|
||
|
|
ones = np.ones(mr.numpy().shape)
|
||
|
|
mask1 = mr.new_image_like(ones)
|
||
|
|
mytx['mask'] = ants.apply_transforms(ct0, mask1, mytx['fwdtransforms'],
|
||
|
|
interpolator='genericLabel',
|
||
|
|
)
|
||
|
|
|
||
|
|
mytx['metric0M'] = ants.create_ants_metric(ct0,
|
||
|
|
mytx['warpedmovout'],
|
||
|
|
metric_type='MattesMutualInformation',
|
||
|
|
moving_mask = mytx['mask'],
|
||
|
|
).get_value()
|
||
|
|
mytx['metric1M'] = ants.create_ants_metric(ct1,
|
||
|
|
mytx['warpedmovout'],
|
||
|
|
metric_type='MattesMutualInformation',
|
||
|
|
moving_mask = mytx['mask'],
|
||
|
|
).get_value()
|
||
|
|
mytx['metricMM'] = ants.create_ants_metric(mr,
|
||
|
|
mytx['warpedfixout'],
|
||
|
|
metric_type='MattesMutualInformation',
|
||
|
|
).get_value()
|
||
|
|
mytx['metric'] = min([mytx['metric0M'], mytx['metric1M'], mytx['metricMM']])
|
||
|
|
mytx['ct'] = m
|
||
|
|
mytx['type'] = (typ, 'fwd')
|
||
|
|
|
||
|
|
mytx['warpedout'] = mytx['warpedmovout']
|
||
|
|
TX.append(mytx)
|
||
|
|
print(mytx['metric'], mytx['metric0M'], mytx['metric1M'], mytx['metricMM'])
|
||
|
|
|
||
|
|
# '''
|
||
|
|
|
||
|
|
mytx = ants.registration(mr, ct, typ)
|
||
|
|
|
||
|
|
ones = np.ones(mr.numpy().shape)
|
||
|
|
mask1 = mr.new_image_like(ones)
|
||
|
|
mytx['mask'] = ants.apply_transforms(ct0, mask1, mytx['fwdtransforms'],
|
||
|
|
interpolator='genericLabel',
|
||
|
|
whichtoinvert=[True],
|
||
|
|
)
|
||
|
|
|
||
|
|
mytx['metricMM'] = ants.create_ants_metric(mr,
|
||
|
|
mytx['warpedmovout'],
|
||
|
|
metric_type='MattesMutualInformation',
|
||
|
|
).get_value()
|
||
|
|
mytx['metric0M'] = ants.create_ants_metric(ct0,
|
||
|
|
mytx['warpedfixout'],
|
||
|
|
metric_type='MattesMutualInformation',
|
||
|
|
fixed_mask = mytx['mask'],
|
||
|
|
).get_value()
|
||
|
|
mytx['metric1M'] = ants.create_ants_metric(ct1,
|
||
|
|
mytx['warpedfixout'],
|
||
|
|
metric_type='MattesMutualInformation',
|
||
|
|
fixed_mask = mytx['mask'],
|
||
|
|
).get_value()
|
||
|
|
mytx['metric'] = min([mytx['metricMM'], mytx['metric0M'], mytx['metric1M']])
|
||
|
|
mytx['ct'] = m
|
||
|
|
mytx['type'] = (typ, 'inv')
|
||
|
|
|
||
|
|
mytx['warpedout'] = mytx['warpedfixout']
|
||
|
|
TX.append(mytx)
|
||
|
|
print(mytx['metric'], mytx['metricMM'], mytx['metric0M'], mytx['metric1M'])
|
||
|
|
|
||
|
|
if min([t['metric'] for t in TX]) < MIN_METRIC:
|
||
|
|
break
|
||
|
|
|
||
|
|
if min([t['metric'] for t in TX]) < MIN_METRIC:
|
||
|
|
break
|
||
|
|
|
||
|
|
print(time.time()-start, 'seconds')
|
||
|
|
|
||
|
|
# exit()
|
||
|
|
|
||
|
|
tx = {
|
||
|
|
'metric': 0,
|
||
|
|
}
|
||
|
|
|
||
|
|
for t in TX:
|
||
|
|
if t['metric'] < tx['metric']:
|
||
|
|
tx = t
|
||
|
|
|
||
|
|
tx['ratio'] = tx['mask'].numpy().sum() / np.prod(ct0.numpy().shape)
|
||
|
|
|
||
|
|
return tx
|
||
|
|
|
||
|
|
def check(epath):
|
||
|
|
registered = 0
|
||
|
|
for root, dirs, files in os.walk(epath):
|
||
|
|
dirs.sort()
|
||
|
|
|
||
|
|
RT_DIR = os.path.join(root, 'RT')
|
||
|
|
|
||
|
|
ORGAN_DIR = os.path.join(RT_DIR, 'ORGAN')
|
||
|
|
if not os.path.isdir(ORGAN_DIR):
|
||
|
|
continue
|
||
|
|
|
||
|
|
# if there is no eye, it's no a brain image
|
||
|
|
eye = None
|
||
|
|
organs = sorted(os.scandir(ORGAN_DIR), key=lambda e: e.name)
|
||
|
|
for o in organs:
|
||
|
|
if 'eye' in o.name.lower():
|
||
|
|
eye = o
|
||
|
|
if eye is None:
|
||
|
|
print('no eye... skip', root)
|
||
|
|
# exit()
|
||
|
|
return None
|
||
|
|
|
||
|
|
ct_image = os.path.join(RT_DIR, 'ct_image.nii.gz')
|
||
|
|
ct0 = sitk.ReadImage(ct_image)
|
||
|
|
ct1 = sitk.Clamp(ct0, sitk.sitkUInt8, 0, 80)
|
||
|
|
print(ct_image, ct0.GetSize())
|
||
|
|
|
||
|
|
outdir = os.path.join(OUT_ROOT, os.path.relpath(root, PATIENTS_ROOT))
|
||
|
|
print(outdir)
|
||
|
|
|
||
|
|
os.makedirs(outdir, exist_ok=True)
|
||
|
|
ct0_nii = os.path.join(outdir, 'ct0.nii.gz')
|
||
|
|
ct1_nii = os.path.join(outdir, 'ct1.nii.gz')
|
||
|
|
shutil.copy(ct_image, ct0_nii)
|
||
|
|
sitk.WriteImage(ct1, ct1_nii)
|
||
|
|
# sitk.WriteImage(sitk.DICOMOrient(ct0), ct0_nii)
|
||
|
|
# sitk.WriteImage(sitk.DICOMOrient(ct1), ct1_nii)
|
||
|
|
|
||
|
|
ct0 = ants.image_read(ct0_nii, reorient=True)
|
||
|
|
ct1 = ants.image_read(ct1_nii, reorient=True)
|
||
|
|
|
||
|
|
|
||
|
|
for root2, dirs2, files2 in os.walk(root):
|
||
|
|
dirs2.sort()
|
||
|
|
skip = (root2==root) or ('RT' in root2.split('/'))
|
||
|
|
if skip:
|
||
|
|
continue
|
||
|
|
if root2.endswith('CT'):
|
||
|
|
modality = 'CT'
|
||
|
|
# continue
|
||
|
|
else:
|
||
|
|
modality = 'other'
|
||
|
|
print(skip, root2, modality)
|
||
|
|
outdir = os.path.join(OUT_ROOT, os.path.relpath(root2, PATIENTS_ROOT))
|
||
|
|
os.makedirs(outdir, exist_ok=True)
|
||
|
|
for e in sorted(os.scandir(root2), key=lambda e: e.name):
|
||
|
|
if not e.name.endswith('.nii.gz'):
|
||
|
|
continue
|
||
|
|
if '_RTDOSE_' in e.name:
|
||
|
|
continue
|
||
|
|
if '_DTI_' in e.name:
|
||
|
|
continue
|
||
|
|
if '_ROI1.' in e.name:
|
||
|
|
continue
|
||
|
|
|
||
|
|
OUT_IMG = os.path.join(outdir, e.name)
|
||
|
|
if os.path.isfile(OUT_IMG):
|
||
|
|
print('skip', OUT_IMG)
|
||
|
|
continue
|
||
|
|
|
||
|
|
print(e.name, e.path)
|
||
|
|
|
||
|
|
fix = ants.image_read(e.path, reorient=True)
|
||
|
|
mytx = registration(ct0, ct1, fix)
|
||
|
|
if mytx is None:
|
||
|
|
continue
|
||
|
|
print(mytx['ratio'], mytx['metric'])
|
||
|
|
|
||
|
|
registered += 1
|
||
|
|
|
||
|
|
os.makedirs(outdir, exist_ok=True)
|
||
|
|
OUT_WARP = OUT_IMG.replace('.nii.gz', '.warp.nii.gz')
|
||
|
|
OUT_MSK = OUT_IMG.replace('.nii.gz', '.mask.nii.gz')
|
||
|
|
OUT_JSON = OUT_IMG.replace('.nii.gz', '.json')
|
||
|
|
OUT_MAT = OUT_IMG.replace('.nii.gz', '.mat')
|
||
|
|
|
||
|
|
jj = {
|
||
|
|
'ct' : mytx['ct'],
|
||
|
|
'type' : mytx['type'],
|
||
|
|
'metric': mytx['metric'],
|
||
|
|
'ratio' : mytx['ratio'],
|
||
|
|
}
|
||
|
|
|
||
|
|
with open(OUT_JSON, 'w') as f:
|
||
|
|
json.dump(jj, f, indent=1)
|
||
|
|
|
||
|
|
if mytx['fwdtransforms']:
|
||
|
|
shutil.copy(mytx['fwdtransforms'][0], OUT_MAT)
|
||
|
|
ants.image_write(mytx['mask'], OUT_MSK)
|
||
|
|
ants.image_write(mytx['warpedout'], OUT_WARP)
|
||
|
|
|
||
|
|
metric = mytx['metric']
|
||
|
|
metric_dir = os.path.join(OUT_ROOT, "1")
|
||
|
|
OUT_TXT = os.path.join(metric_dir, '%f-%s'%(metric, e.name.replace('.nii.gz', '.txt')))
|
||
|
|
jj['dir'] = outdir
|
||
|
|
os.makedirs(metric_dir, exist_ok=True)
|
||
|
|
with open(OUT_TXT, 'w') as f:
|
||
|
|
json.dump(jj, f, indent=1)
|
||
|
|
|
||
|
|
shutil.copy(e.path, OUT_IMG)
|
||
|
|
|
||
|
|
|
||
|
|
return registered
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
|
||
|
|
# check('/mnt/1220/Public/dataset2/G4/3L6LOEER') # bad registration
|
||
|
|
# exit()
|
||
|
|
|
||
|
|
EXCLUDE = (
|
||
|
|
# 'LLUQJUY4', #cervical
|
||
|
|
)
|
||
|
|
|
||
|
|
os.makedirs(OUT_ROOT, exist_ok=True)
|
||
|
|
|
||
|
|
LOCK_DIR = os.path.join(OUT_ROOT, '0lock')
|
||
|
|
os.makedirs(LOCK_DIR, exist_ok=True)
|
||
|
|
for e in sorted(os.scandir(PATIENTS_ROOT), key=lambda e: e.name):
|
||
|
|
if e.is_dir():
|
||
|
|
d = shelve.open(SHELVE)
|
||
|
|
if e.name in d or e.name in EXCLUDE:
|
||
|
|
print('skip', e.name)
|
||
|
|
d.close()
|
||
|
|
continue
|
||
|
|
d.close()
|
||
|
|
lock_path = os.path.join(LOCK_DIR, '%s.lock'%e.name)
|
||
|
|
lock = filelock.FileLock(lock_path, timeout=1)
|
||
|
|
try:
|
||
|
|
lock.acquire()
|
||
|
|
except:
|
||
|
|
print(lock_path, 'locked')
|
||
|
|
continue
|
||
|
|
ret = check(e.path)
|
||
|
|
lock.release()
|
||
|
|
# exit()
|
||
|
|
d = shelve.open(SHELVE)
|
||
|
|
d[e.name] = ret
|
||
|
|
d.close()
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|