|
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'])) |
|
|
|
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() |
