File size: 4,327 Bytes
c642393 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import numpy as np
import glob
import ants
import nibabel as nib
import os
import argparse
import sys
from pathlib import Path
def parse_command_line():
parser = argparse.ArgumentParser(
description='pipeline for data preprocessing')
parser.add_argument('-bp', metavar='base path', type=str,
help="absolute path of the base directory")
parser.add_argument('-ip', metavar='image path', type=str,
help="relative path of the image directory")
parser.add_argument('-sp', metavar='segmentation path', type=str,
help="relative path of the image directory")
parser.add_argument('-op', metavar='preprocessing result output path', type=str, default='output',
help='relative path of the preprocessing result directory')
argv = parser.parse_args()
return argv
def flip(nib_img, nib_seg, ants_img, ants_seg, seg_fomat):
img = nib_img.get_fdata()
if seg_fomat == 'nii.gz' or seg_fomat == 'nii':
seg = nib_seg.get_fdata()
else:
seg = nib_seg[0]
gem = ants.label_geometry_measures(ants_seg, ants_img)
low_x = min(list(gem.loc[:, 'BoundingBoxLower_x']))
upp_x = max(list(gem.loc[:, 'BoundingBoxUpper_x']))
low_y = min(list(gem.loc[:, 'BoundingBoxLower_y']))
upp_y = max(list(gem.loc[:, 'BoundingBoxUpper_y']))
low_z = min(list(gem.loc[:, 'BoundingBoxLower_z']))
upp_z = max(list(gem.loc[:, 'BoundingBoxUpper_z']))
# Compute mid point
mid_x = int((low_x + upp_x) / 2)
left_seg = seg[:mid_x, :, :]
left_img = img[:mid_x, :, :]
right_seg = seg[mid_x:, :, :]
right_img = img[mid_x:, :, :]
flipped_right_seg = np.flip(right_seg, axis=0)
flipped_right_img = np.flip(right_img, axis=0)
print("finish flip")
return left_img, left_seg, flipped_right_img, flipped_right_seg
def load_data(img_path, seg_path):
nib_seg = nib.load(seg_path)
nib_img = nib.load(img_path)
ants_seg = ants.image_read(seg_path)
ants_img = ants.image_read(img_path)
return nib_img, nib_seg, ants_img, ants_seg
def crop_flip_save_file(left_img, left_seg, flipped_right_img, flipped_right_seg, nib_img, nib_seg, output_img, output_seg, scan_id):
left_img_nii = nib.Nifti1Image(
left_img, affine=nib_img.affine, header=nib_img.header)
left_seg_nii = nib.Nifti1Image(
left_seg, affine=nib_seg.affine, header=nib_seg.header)
right_img_nii = nib.Nifti1Image(
flipped_right_img, affine=nib_img.affine, header=nib_img.header)
right_seg_nii = nib.Nifti1Image(
flipped_right_seg, affine=nib_seg.affine, header=nib_seg.header)
left_img_nii.to_filename(os.path.join(
output_img, scan_id + '1.nii.gz'))
left_seg_nii.to_filename(os.path.join(
output_seg, scan_id + '1.nii.gz'))
right_img_nii.to_filename(os.path.join(
output_img, scan_id + '0.nii.gz'))
right_seg_nii.to_filename(os.path.join(
output_seg, scan_id + '0.nii.gz'))
def main():
args = parse_command_line()
base_path = args.bp
image_path = os.path.join(base_path, args.ip)
seg_path = os.path.join(base_path, args.sp)
output_path = os.path.join(base_path, args.op)
output_img = os.path.join(output_path, 'images')
output_seg = os.path.join(output_path, 'labels')
try:
os.mkdir(output_path)
except:
print(f'{output_path} is already existed')
try:
os.mkdir(output_img)
except:
print(f'{output_img} is already existed')
try:
os.mkdir(output_seg)
except:
print(f'{output_seg} is already existed')
for i in sorted(glob.glob(image_path + '/*nii.gz')):
id = os.path.basename(i).split('.')[0]
label_path = os.path.join(seg_path, id + '.nii.gz')
nib_img, nib_seg, ants_img, ants_seg = load_data(i, label_path)
left_img, left_seg, flipped_right_img, flipped_right_seg = flip(nib_img, nib_seg, ants_img, ants_seg, 'nii.gz')
print('Scan ID: ' + id + f', img & seg before cropping: {nib_img.get_fdata().shape}, after flipping: {left_img.shape} and {flipped_right_img.shape}')
crop_flip_save_file(left_img, left_seg, flipped_right_img, flipped_right_seg, nib_img, nib_seg, output_img, output_seg, id)
if __name__ == '__main__':
main() |