|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import nibabel as nib |
|
from nibabel import io_orientation |
|
from batchgenerators.utilities.file_and_folder_operations import * |
|
import numpy as np |
|
import os |
|
from multiprocessing import Pool |
|
import SimpleITK as sitk |
|
|
|
|
|
def print_shapes(folder: str) -> None: |
|
for i in subfiles(folder, suffix='.nii.gz'): |
|
tmp = sitk.ReadImage(i) |
|
print(sitk.GetArrayFromImage(tmp).shape, tmp.GetSpacing()) |
|
|
|
|
|
def reorient_to_ras(image: str) -> None: |
|
""" |
|
Will overwrite image!!! |
|
:param image: |
|
:return: |
|
""" |
|
assert image.endswith('.nii.gz') |
|
origaffine_pkl = image[:-7] + '_originalAffine.pkl' |
|
if not isfile(origaffine_pkl): |
|
img = nib.load(image) |
|
original_affine = img.affine |
|
original_axcode = nib.aff2axcodes(img.affine) |
|
img = img.as_reoriented(io_orientation(img.affine)) |
|
new_axcode = nib.aff2axcodes(img.affine) |
|
print(image.split('/')[-1], 'original axcode', original_axcode, 'now (should be ras)', new_axcode) |
|
nib.save(img, image) |
|
save_pickle((original_affine, original_axcode), origaffine_pkl) |
|
|
|
|
|
def revert_reorientation(image: str) -> None: |
|
assert image.endswith('.nii.gz') |
|
expected_pkl = image[:-7] + '_originalAffine.pkl' |
|
assert isfile(expected_pkl), 'Must have a file with the original affine, as created by ' \ |
|
'reorient_to_ras. Expected filename: %s' % \ |
|
expected_pkl |
|
original_affine, original_axcode = load_pickle(image[:-7] + '_originalAffine.pkl') |
|
img = nib.load(image) |
|
before_revert = nib.aff2axcodes(img.affine) |
|
img = img.as_reoriented(io_orientation(original_affine)) |
|
after_revert = nib.aff2axcodes(img.affine) |
|
print('before revert', before_revert, 'after revert', after_revert) |
|
restored_affine = img.affine |
|
assert np.all(np.isclose(original_affine, restored_affine)), 'restored affine does not match original affine, ' \ |
|
'aborting!' |
|
nib.save(img, image) |
|
os.remove(expected_pkl) |
|
|
|
|
|
def reorient_all_images_in_folder_to_ras(folder: str, num_processes: int = 8): |
|
p = Pool(num_processes) |
|
nii_files = subfiles(folder, suffix='.nii.gz', join=True) |
|
p.map(reorient_to_ras, nii_files) |
|
p.close() |
|
p.join() |
|
|
|
|
|
def revert_orientation_on_all_images_in_folder(folder: str, num_processes: int = 8): |
|
p = Pool(num_processes) |
|
nii_files = subfiles(folder, suffix='.nii.gz', join=True) |
|
p.map(revert_reorientation, nii_files) |
|
p.close() |
|
p.join() |
|
|
|
|
|
if __name__ == '__main__': |
|
"""nib.as_closest_canonical() |
|
test_img = '/home/fabian/data/la_005_0000.nii.gz' |
|
test_img_reorient = test_img[:-7] + '_reorient.nii.gz' |
|
test_img_restored = test_img[:-7] + '_restored.nii.gz' |
|
|
|
img = nib.load(test_img) |
|
print('loaded original') |
|
print('shape', img.shape) |
|
print('affine', img.affine) |
|
original_affine = img.affine |
|
original_axcode = nib.aff2axcodes(img.affine) |
|
print('orientation', nib.aff2axcodes(img.affine)) |
|
|
|
print('reorienting...') |
|
img = img.as_reoriented(io_orientation(img.affine)) |
|
nib.save(img, test_img_reorient) |
|
|
|
print('now loading the reoriented img') |
|
img = nib.load(test_img_reorient) |
|
print('loaded original') |
|
print('shape', img.shape) |
|
print('affine', img.affine) |
|
reorient_affine = img.affine |
|
reorient_axcode = nib.aff2axcodes(img.affine) |
|
print('orientation', nib.aff2axcodes(img.affine)) |
|
|
|
print('restoring original geometry') |
|
img = img.as_reoriented(io_orientation(original_affine)) |
|
restored_affine = img.affine |
|
nib.save(img, test_img_restored) |
|
|
|
print('now loading the restored img') |
|
img = nib.load(test_img_restored) |
|
print('loaded original') |
|
print('shape', img.shape) |
|
print('affine', img.affine) |
|
print('orientation', nib.aff2axcodes(img.affine))""" |
|
|
