From 7b0222946c2fb4c89f601f1827aa10e295d5c474 Mon Sep 17 00:00:00 2001 From: Xiao Furen Date: Sat, 8 Feb 2025 08:39:18 +0800 Subject: [PATCH] m6ants creates temp dir in home --- src/g4synthmorph.py | 4 +- src/m6ants.py | 329 +++++++++++++++++++++++++++++++++++++++++++ src/m6synthmorph.py | 20 ++- src/requirements.txt | 1 + zz/fireants-test.py | 78 ++++++++++ zz/greedy-test.py | 42 ++++++ 6 files changed, 469 insertions(+), 5 deletions(-) create mode 100644 src/m6ants.py create mode 100644 zz/fireants-test.py create mode 100644 zz/greedy-test.py diff --git a/src/g4synthmorph.py b/src/g4synthmorph.py index 92c9e60..42ee04b 100644 --- a/src/g4synthmorph.py +++ b/src/g4synthmorph.py @@ -350,7 +350,9 @@ def main(): # exit() EXCLUDE = ( - # 'LLUQJUY4', #cervical + '6HKU7ZLF', # W external/local_tsl/tsl/framework/bfc_allocator.cc:482] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.00GiB (rounded to 2145976320)requested by op AddV2 + 'ED4NQNMK', # I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Input is not invertible. + 'NJJS6NUS', 'PNDQBHFH', # W external/local_tsl/tsl/framework/bfc_allocator.cc:482] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.00GiB (rounded to 2145976320)requested by op AddV2 ) os.makedirs(OUT_ROOT, exist_ok=True) diff --git a/src/m6ants.py b/src/m6ants.py new file mode 100644 index 0000000..fc2b183 --- /dev/null +++ b/src/m6ants.py @@ -0,0 +1,329 @@ + +''' +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/1218/Public/dataset2/M6' +OUT_ROOT = '/mnt/1248/Public/dataset2/M6-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() + + 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() diff --git a/src/m6synthmorph.py b/src/m6synthmorph.py index f723d4a..9376824 100644 --- a/src/m6synthmorph.py +++ b/src/m6synthmorph.py @@ -214,7 +214,12 @@ def register(ct0, ct1, moving, out_root): # XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda logger.info('registering %s'%m) - registration.register(arg) + try: + registration.register(arg) + except Exception as e: + logger.error('failed to register %s %s'% (moving, m)) + logger.exception(e) + return None logger.info('registered %s'%m) @@ -336,7 +341,9 @@ def check(epath): logger.info(' '.join([e.name, e.path])) moving = e.path - register(ct_image, ct1_nii, moving, outdir) + ret = register(ct_image, ct1_nii, moving, outdir) + if ret is None: + return None registered += 1 # exit() @@ -356,7 +363,10 @@ def main(): EXCLUDE = ( # 'LLUQJUY4', #cervical - '2XYU7UHB', # I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Input is not invertible. + + # 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) @@ -378,8 +388,10 @@ def main(): except: logger.info(lock_path + ' locked') continue - ret = check(e.path) + ret = check(e.path) lock.release() + if ret is None: + continue # exit() d = shelve.open(SHELVE) d[e.name] = ret diff --git a/src/requirements.txt b/src/requirements.txt index 94c5307..afa9afa 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -1,4 +1,5 @@ +antspyx filelock git+https://github.com/adalca/neurite.git diff --git a/zz/fireants-test.py b/zz/fireants-test.py new file mode 100644 index 0000000..45205c9 --- /dev/null +++ b/zz/fireants-test.py @@ -0,0 +1,78 @@ +''' + + + +conda deactivate +conda create -y -n fireants -c conda-forge 'numpy<2' simpleitk=2.2.1 +conda activate fireants +pip install fireants + + +CUDA out of memory. ??? + +''' + + +from fireants.io import Image, BatchedImages +from fireants.registration import AffineRegistration, GreedyRegistration +import matplotlib.pyplot as plt +import SimpleITK as sitk +from time import time + +# load the images +# image1 = Image.load_file("atlas_2mm_1000_3.nii.gz") +# image2 = Image.load_file("atlas_2mm_1001_3.nii.gz") +image1 = Image.load_file("/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz") +image2 = Image.load_file("/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz") + +# batchify them (we only have a single image per batch, but we can pass multiple images) +batch1 = BatchedImages([image1]) +batch2 = BatchedImages([image2]) + +# check device name +print(batch1().device) + +# specify some values +scales = [4, 2, 1] # scales at which to perform registration +iterations = [200, 100, 50] +optim = 'Adam' +lr = 3e-3 + +# create affine registration object +affine = AffineRegistration(scales, iterations, batch1, batch2, optimizer=optim, optimizer_lr=lr, + loss_type = 'mi', + cc_kernel_size=5) + +# run registration +start = time() +transformed_images = affine.optimize(save_transformed=True) +end = time() + +print("Runtime", end - start, "seconds") + +reg = GreedyRegistration(scales=[4, 2, 1], iterations=[200, 100, 25], + fixed_images=batch1, moving_images=batch2, + cc_kernel_size=5, deformation_type='compositive', + smooth_grad_sigma=1, + loss_type = 'mi', + optimizer='adam', optimizer_lr=0.5, init_affine=affine.get_affine_matrix().detach()) + +start = time() +reg.optimize(save_transformed=False) +end = time() + +print("Runtime", end - start, "seconds") + +moved = reg.evaluate(batch1, batch2) + +# reference_img = sitk.ReadImage("atlas_2mm_1000_3.nii.gz") +reference_img = sitk.ReadImage("/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz") + +# Preparing the moving image to be written out +moved_image_np = moved[0, 0].detach().cpu().numpy() # volumes are typically stored in tensors with dimensions [Batch, Channels, Depth, Height, Width], so extracting the latter 3 for nifti +moved_sitk_image = sitk.GetImageFromArray(moved_image_np) +moved_sitk_image.SetOrigin(reference_img.GetOrigin()) +moved_sitk_image.SetSpacing(reference_img.GetSpacing()) +moved_sitk_image.SetDirection(reference_img.GetDirection()) +sitk.WriteImage(moved_sitk_image, 'reslice_deform_atlas_2mm_1000_3.nii.gz') + diff --git a/zz/greedy-test.py b/zz/greedy-test.py new file mode 100644 index 0000000..053172c --- /dev/null +++ b/zz/greedy-test.py @@ -0,0 +1,42 @@ +''' +conda deactivate +conda create -y -n greedy -c conda-forge python +conda activate greedy + +cmake -DCMAKE_INSTALL_PREFIX=/home/xfr/.conda/envs/greedy .. + + +greedy -d 3 -i \ +"/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz" \ +"/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz" \ +-o test \ +-m MI \ + + +time greedy -d 3 -a \ + -m MI \ + -i \ + "/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz" \ + "/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz" \ + -o affine.mat \ + -dof 6 \ + -ia-image-centers -n 100x50x10 + + +time greedy -d 3 \ + -m MI \ + -i \ + "/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz" \ + "/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz" \ + -it affine.mat \ + -o warp.nii.gz \ + -oinv inverse_warp.nii.gz \ + -sv -n 100x50x10 + +greedy -d 3 \ + -rf "/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz" \ + -rm "/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz" resliced.nii.gz \ + -r warp.nii.gz affine.mat + + +''' \ No newline at end of file