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