fix ants
This commit is contained in:
parent
7b0222946c
commit
741b07fad1
2 changed files with 368 additions and 1 deletions
348
src/g4ants.py
Normal file
348
src/g4ants.py
Normal file
|
|
@ -0,0 +1,348 @@
|
|||
|
||||
'''
|
||||
Use ants to register M6 images
|
||||
|
||||
https://download-directory.github.io/
|
||||
https://github.com/freesurfer/freesurfer/tree/dev/mri_synthmorph
|
||||
|
||||
|
||||
CUDA_VISIBLE_DEVICES=3 python m6synthmorph.py
|
||||
|
||||
|
||||
|
||||
|
||||
'''
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import logging
|
||||
import json
|
||||
import os
|
||||
import shelve
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
from skimage.metrics import normalized_mutual_information
|
||||
|
||||
import filelock
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import SimpleITK as sitk
|
||||
|
||||
import ants
|
||||
|
||||
### Need NFS for lock
|
||||
PATIENTS_ROOT = '/mnt/1220/Public/dataset2/G4'
|
||||
OUT_ROOT = '/mnt/1248/Public/dataset2/G4-ants'
|
||||
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
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||||
handlers=[
|
||||
logging.StreamHandler(),
|
||||
logging.FileHandler('g4synthmorph.log')
|
||||
]
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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
|
||||
|
||||
'''
|
||||
Namespace(command='register', moving='/nn/7295866/20250127/nii/7_3D_SAG_T1_MPRAGE_+C_20250127132612_100.nii.gz', fixed='/123/onlylian/0/tmpgp96622o/clipped.nii.gz',
|
||||
model='joint', out_moving='/123/onlylian/0/tmpgp96622o/joint.nii.gz', out_fixed='/123/onlylian/0/tmpgp96622o/out_fixed-joint.nii.gz',
|
||||
header_only=False, trans='/123/onlylian/0/tmpgp96622o/moving_to_fixed-joint.nii.gz', inverse='/123/onlylian/0/tmpgp96622o/fixed_to_moving-joint.nii.gz',
|
||||
init=None, mid_space=False, threads=None, gpu=True, hyper=0.5, steps=7, extent=256, weights=None, verbose=False, out_dir=None)
|
||||
'''
|
||||
def register(ct0, ct1, moving, out_root):
|
||||
FREESURFER_HOME = '/mnt/1218/Public/packages/freesurfer-8.0.0-beta/'
|
||||
# out_root = Path(ct0).resolve().parent/os.path.basename(mr).replace('.nii.gz','')
|
||||
|
||||
# print(out_root)
|
||||
modality = os.path.basename(out_root)
|
||||
# exit()
|
||||
|
||||
out_root = Path(out_root)/os.path.basename(moving).replace('.nii.gz','')
|
||||
out_root.mkdir(exist_ok=True)
|
||||
|
||||
logger.info(' '.join((modality, ct0, ct1, moving, str(out_root))))
|
||||
|
||||
base = ants.image_read(ct0)
|
||||
base1 = ants.image_read(ct1)
|
||||
orig = ants.image_read(moving)
|
||||
|
||||
ct0_numpy = base.numpy()
|
||||
ct1_numpy = base1.numpy()
|
||||
mov_numpy = orig.numpy()
|
||||
|
||||
# logger.info(str(mov_numpy.shape))
|
||||
# if min(mov_numpy.shape) < 3:
|
||||
# logger.error('The image is too small... skip')
|
||||
# return None
|
||||
|
||||
mov_min = orig.min()
|
||||
mov_max = orig.max()
|
||||
if mov_min == mov_max:
|
||||
logger.error('Total Mass of the image was zero. Aborting here to prevent division by zero later on.... skip')
|
||||
return None
|
||||
|
||||
|
||||
logger.info(str(mov_numpy.shape))
|
||||
if min(mov_numpy.shape) < 3:
|
||||
logger.error('The image is too small... skip')
|
||||
return None
|
||||
|
||||
|
||||
|
||||
if modality == 'XA':
|
||||
exit()
|
||||
|
||||
if modality == 'CT':
|
||||
# clipped = out_root/'clipped.nii.gz'
|
||||
|
||||
# cl = sitk.Clamp(orig, sitk.sitkUInt8, 0, 80)
|
||||
# sitk.WriteImage(cl, clipped)
|
||||
|
||||
MODELS = [
|
||||
'Rigid',
|
||||
# 'TRSAA',
|
||||
]
|
||||
|
||||
else:
|
||||
clipped = moving
|
||||
MODELS = [
|
||||
'Rigid',
|
||||
# 'Affine',
|
||||
'TRSAA',
|
||||
# 'SyN',
|
||||
'SyNRA',
|
||||
# 'antsRegistrationSyNsr',
|
||||
# 'antsRegistrationSyNQuicksr',
|
||||
# 'joint',
|
||||
]
|
||||
|
||||
# anst_fixed = ants.image_read(ct0)
|
||||
# anst_moving = ants.image_read(moving)
|
||||
|
||||
METRICS0 = {}
|
||||
METRICS1 = {}
|
||||
|
||||
for m in MODELS:
|
||||
|
||||
os.makedirs(m, exist_ok=True)
|
||||
logger.info('%s registerinig'%m)
|
||||
|
||||
ret = ants.registration(base, orig,
|
||||
type_of_transform = m,
|
||||
outprefix = '%s/'%m,
|
||||
)
|
||||
logger.info('%s registered'%m)
|
||||
ret['warpedmovout'].image_write('%s/warpedmovout.nii.gz'%m)
|
||||
ret['warpedfixout'].image_write('%s/warpedfixout.nii.gz'%m)
|
||||
# print(ret)
|
||||
|
||||
warpedmovout_numpy = ret['warpedmovout'].numpy()
|
||||
fix0_mov = normalized_mutual_information(ct0_numpy, warpedmovout_numpy)
|
||||
fix1_mov = normalized_mutual_information(ct1_numpy, warpedmovout_numpy)
|
||||
mov_fix0 = normalized_mutual_information(mov_numpy, ret['warpedfixout'].numpy())
|
||||
|
||||
METRICS0[m] = (fix0_mov, fix1_mov, mov_fix0)
|
||||
METRICS1[m] = max(fix0_mov, fix1_mov, mov_fix0)
|
||||
|
||||
print(METRICS1)
|
||||
|
||||
with open('%s/fwdtransforms.json'%m, 'w') as transforms:
|
||||
json.dump(ret['fwdtransforms'], transforms, indent=1)
|
||||
|
||||
with open('%s/invtransforms.json'%m, 'w') as transforms:
|
||||
json.dump(ret['invtransforms'], transforms, indent=1)
|
||||
|
||||
outdir = '%s/%s'%(out_root,m)
|
||||
|
||||
logger.info('moving output to %s'%outdir)
|
||||
if os.path.exists(outdir):
|
||||
shutil.rmtree(outdir)
|
||||
shutil.move(m, out_root)
|
||||
# shutil.copytree(src, dst, symlinks=False, ignore=None, copy_function=copy2, ignore_dangling_symlinks=False, dirs_exist_ok=False)¶
|
||||
|
||||
# exit()
|
||||
|
||||
with open(out_root/'metrics0.json', 'w') as f_metrics:
|
||||
json.dump(METRICS0, f_metrics, indent=1)
|
||||
|
||||
with open(out_root/'metrics1.json', 'w') as f_metrics:
|
||||
json.dump(METRICS1, f_metrics, indent=1)
|
||||
|
||||
print(METRICS0)
|
||||
print(METRICS1)
|
||||
# exit()
|
||||
|
||||
return out_root
|
||||
|
||||
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:
|
||||
logger.info('no eye... skip ' + root)
|
||||
# exit()
|
||||
return None
|
||||
|
||||
ct_image = os.path.join(RT_DIR, 'ct_image.nii.gz')
|
||||
|
||||
outdir = os.path.join(OUT_ROOT, os.path.relpath(root, PATIENTS_ROOT))
|
||||
logger.info(outdir)
|
||||
|
||||
os.makedirs(outdir, exist_ok=True)
|
||||
# ct0_nii = os.path.join(outdir, 'ct0.nii.gz')
|
||||
ct1_nii = os.path.join(outdir, 'clipped.nii.gz')
|
||||
# shutil.copy(ct_image, ct0_nii)
|
||||
|
||||
ct = sitk.ReadImage(ct_image)
|
||||
clipped = sitk.Clamp(ct, sitk.sitkUInt8, 0, 80)
|
||||
sitk.WriteImage(clipped, ct1_nii)
|
||||
|
||||
for root2, dirs2, files2 in os.walk(root):
|
||||
dirs2.sort()
|
||||
outdir = os.path.join(OUT_ROOT, os.path.relpath(root2, PATIENTS_ROOT))
|
||||
if root2.endswith('RT'):
|
||||
modality = 'RT'
|
||||
logger.info('copying %s %s' %(root2, outdir))
|
||||
shutil.copytree(root2, outdir, dirs_exist_ok=True)
|
||||
# exit()
|
||||
continue
|
||||
skip = (root2==root) or ('RT' in root2.split('/'))
|
||||
if skip:
|
||||
continue
|
||||
if root2.endswith('CT'):
|
||||
modality = 'CT'
|
||||
# continue
|
||||
else:
|
||||
modality = 'other'
|
||||
logger.info(' '.join([str(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):
|
||||
logger.info('skip '+ OUT_IMG)
|
||||
continue
|
||||
|
||||
logger.info(' '.join([e.name, e.path]))
|
||||
|
||||
moving = e.path
|
||||
ret = register(ct_image, ct1_nii, moving, outdir)
|
||||
if ret is None:
|
||||
return None
|
||||
registered += 1
|
||||
# exit()
|
||||
|
||||
|
||||
|
||||
# exit()
|
||||
return registered
|
||||
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
# check('/mnt/1218/Public/dataset2/M6/22GD5VL2') # first case
|
||||
# exit()
|
||||
|
||||
EXCLUDE = (
|
||||
# 'LLUQJUY4', #cervical
|
||||
|
||||
# I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Input is not invertible.
|
||||
# '2XYU7UHB', '4YADZMJP', '7NVOP5OK',
|
||||
# 'BQ773ST6', 'DHCW26GK', 'DS5IIRYB',
|
||||
)
|
||||
|
||||
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:
|
||||
logger.info('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:
|
||||
logger.info(lock_path + ' locked')
|
||||
continue
|
||||
ret = check(e.path)
|
||||
lock.release()
|
||||
if ret is None:
|
||||
continue
|
||||
# exit()
|
||||
d = shelve.open(SHELVE)
|
||||
d[e.name] = ret
|
||||
d.close()
|
||||
|
||||
if __name__ == '__main__':
|
||||
# dn = os.path.dirname(os.path.realpath(__file__))
|
||||
# prefix='%s/tmp/'%dn
|
||||
# os.makedirs(prefix, exist_ok=True)
|
||||
|
||||
home = Path.home()
|
||||
prefix = str(home / '0')
|
||||
print(prefix)
|
||||
os.makedirs(prefix, exist_ok=True)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix=prefix) as tmpdirname:
|
||||
# print('created temporary directory', tmpdirname)
|
||||
os.chdir(tmpdirname)
|
||||
main()
|
||||
|
|
@ -32,7 +32,7 @@ import SimpleITK as sitk
|
|||
import ants
|
||||
|
||||
### Need NFS for lock
|
||||
PATIENTS_ROOT = '/mnt/1218/Public/dataset2/M6'
|
||||
PATIENTS_ROOT = '/mnt/1220/Public/dataset2/M6'
|
||||
OUT_ROOT = '/mnt/1248/Public/dataset2/M6-ants'
|
||||
SHELVE = os.path.join(OUT_ROOT, '0shelve')
|
||||
|
||||
|
|
@ -100,6 +100,25 @@ def register(ct0, ct1, moving, out_root):
|
|||
ct1_numpy = base1.numpy()
|
||||
mov_numpy = orig.numpy()
|
||||
|
||||
# logger.info(str(mov_numpy.shape))
|
||||
# if min(mov_numpy.shape) < 3:
|
||||
# logger.error('The image is too small... skip')
|
||||
# return None
|
||||
|
||||
mov_min = orig.min()
|
||||
mov_max = orig.max()
|
||||
if mov_min == mov_max:
|
||||
logger.error('Total Mass of the image was zero. Aborting here to prevent division by zero later on.... skip')
|
||||
return None
|
||||
|
||||
|
||||
logger.info(str(mov_numpy.shape))
|
||||
if min(mov_numpy.shape) < 3:
|
||||
logger.error('The image is too small... skip')
|
||||
return None
|
||||
|
||||
|
||||
|
||||
if modality == 'XA':
|
||||
exit()
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue