diff --git a/manual_correction.py b/manual_correction.py index 894133a..74abaec 100644 --- a/manual_correction.py +++ b/manual_correction.py @@ -258,6 +258,21 @@ def get_parser(): }\n """), ) + parser.add_argument( + '-change-orient', + type=str, + help= + " R| Orientation to show the image in the viewer. If provided, the image and label will be reoriented before " + "opening the viewer. After manual correction, the image and label will be reoriented back to the original " + "orientation. " + "Possible options: LAS,LAI,LPS,LPI,LSA,LSP,LIA,LIP,RAS,RAI,RPS,RPI,RSA,RSP,RIA,RIP,ALS,ALI,ARS,ARI,ASL,ASR," + "AIL,AIR,PLS,PLI,PRS,PRI,PSL,PSR,PIL,PIR,SLA,SLP,SRA,SRP,SAL,SAR,SPL,SPR,ILA,ILP,IRA,IRP,IAL,IAR,IPL,IPR}", + choices=['LAS', 'LAI', 'LPS', 'LPI', 'LSA', 'LSP', 'LIA', 'LIP', 'RAS', 'RAI', 'RPS', 'RPI', 'RSA', 'RSP', + 'RIA', 'RIP', 'ALS', 'ALI', 'ARS', 'ARI', 'ASL', 'ASR', 'AIL', 'AIR', 'PLS', 'PLI', 'PRS', 'PRI', + 'PSL', 'PSR', 'PIL', 'PIR', 'SLA', 'SLP', 'SRA', 'SRP', 'SAL', 'SAR', 'SPL', 'SPR', 'ILA', 'ILP', + 'IRA', 'IRP', 'IAL', 'IAR', 'IPL', 'IPR'], + default='' + ) parser.add_argument( '-v', '--verbose', help="Full verbose (for debugging)", @@ -862,6 +877,16 @@ def main(): # For example: '/Users/user/dataset/derivatives/labels/sub-001/anat/sub-001_T2w_seg.nii.gz' # The information regarding the modified data will be stored within the sidecar .json file fname_out = utils.add_suffix(os.path.join(path_out, subject, ses, contrast, filename), suffix_dict[task]) + + # Change orientation of the input image (if different from the original orientation) + if args.change_orient: + # Get image and label orientation + image_orig_orient = utils.get_orientation(fname) + label_orig_orient = utils.get_orientation(fname_label) + # Change orientation of the input image for better visualization + if image_orig_orient != args.change_orient or label_orig_orient != args.change_orient: + utils.change_orientation(fname, args.change_orient) + utils.change_orientation(fname_label, args.change_orient) # Create subject folder in output if they do not exist os.makedirs(os.path.join(path_out, subject, ses, contrast), exist_ok=True) @@ -924,6 +949,14 @@ def main(): # Keep track of corrected files in YAML. dict_yml = utils.track_corrections(files_dict=dict_yml.copy(), config_path=args.config, file_path=fname, task=task) + # Change orientation of the input image back to the original orientation + if args.change_orient: + image_current_orientation = utils.get_orientation(fname) + label_current_orientation = utils.get_orientation(fname_label) + if image_current_orientation != image_orig_orient or label_current_orientation != label_orig_orient: + utils.change_orientation(fname, image_orig_orient) + utils.change_orientation(fname_out, label_orig_orient) + else: sys.exit("ERROR: The list of files to correct is empty. \nMaybe, you have already corrected all the " "files? Please, check the YAML file: {}".format(args.config)) diff --git a/utils.py b/utils.py index f838ac5..3e10249 100644 --- a/utils.py +++ b/utils.py @@ -332,3 +332,43 @@ def track_corrections(files_dict, config_path, file_path, task): return files_dict +def get_orientation(file_path): + """ + Get the orientation of the input nifti file + :param file_path: path to the nifti file + :return: actual orientation of the nifti file, e.g., 'RPI' + """ + + def _parse_orientation(output_bytes: bytes) -> str: + """ + Parse the image orientation from the provided output bytes. + Args: output_bytes (bytes): The input bytes containing the output from sct_image command. + Returns: + str: The parsed image orientation. + """ + + output_string = output_bytes.decode('utf-8') + lines = output_string.strip().split('\n') + # Get only the string containing the orientation, e.g., 'RPI' + orientation = lines[-1].strip() + return orientation + + # Note: we use bash command 'sct_image' to get the orientation instead of SCT's Image class because this would + # introduce dependency on SCT conda environment + output = subprocess.run(['sct_image', '-i', file_path, '-getorient'], capture_output=True) + orientation = _parse_orientation(output.stdout) # e.g., 'RPI' + + return orientation + + +def change_orientation(file_path, orientation): + """ + Change the orientation of the input nifti file + :param file_path: path to the nifti file + :param orientation: desired orientation of the nifti file, e.g., 'RPI' + """ + + # Change the orientation of the input nifti file + # Note: The image is currently being overwritten + subprocess.run(['sct_image', '-i', file_path, '-setorient', orientation, '-o', file_path], capture_output=True) + print(f"Orientation of {file_path} has been changed to {orientation}")