diff --git a/run.py b/run.py index 6017617..e1cf418 100755 --- a/run.py +++ b/run.py @@ -5,7 +5,7 @@ Connectome ABCD-XCP niBabies Imaging nnu-NET (CABINET) Greg Conan: gconan@umn.edu Created: 2021-11-12 -Updated: 2022-04-29 +Updated: 2022-05-01 """ # Import standard libraries @@ -17,6 +17,7 @@ import os import pandas as pd import pdb +import shutil import subprocess import sys @@ -43,12 +44,12 @@ def find_myself(flg): # Custom local imports from src.utilities import ( - as_cli_attr, as_cli_arg, copy_and_rename_file, correct_chirality, - create_anatomical_average, crop_image, dict_has, ensure_dict_has, - exit_with_time_info, extract_from_json, get_and_make_preBIBSnet_work_dirs, - get_optional_args_in, get_stage_name, get_subj_ID_and_session, - resize_images, run_all_stages, valid_readable_json, - validate_parameter_types, valid_readable_dir, warn_user_of_conditions + as_cli_attr, as_cli_arg, correct_chirality, + create_anatomical_average, crop_image, dict_has, dilate_LR_mask, + ensure_dict_has, exit_with_time_info, extract_from_json, + get_and_make_preBIBSnet_work_dirs, get_optional_args_in, get_stage_name, + get_subj_ID_and_session, resize_images, run_all_stages, + valid_readable_json, validate_parameter_types, valid_readable_dir ) @@ -287,7 +288,7 @@ def run_preBIBSnet(j_args, logger): # TODO Unredundantify the blocks below os.makedirs(os.path.dirname(out_nii_fpath), exist_ok=True) if not os.path.exists(out_nii_fpath): # j_args["common"]["overwrite"] or - os.symlink(transformed_images[tw], out_nii_fpath) + shutil.copy2(transformed_images[tw], out_nii_fpath) # os.symlink( """ concat_mat = transformed_images["T1w_concat_mat"] out_mat_fpath = os.path.join(j_args["optional_out_dirs"]["postBIBSnet"], @@ -359,25 +360,35 @@ def run_postBIBSnet(j_args, logger): # Run left/right registration script and chirality correction left_right_mask_nifti_fpath = run_left_right_registration( j_args, sub_ses, age_months, 2 if int(age_months) < 22 else 1, logger # NOTE 22 cutoff might change - ) + ) # TODO Add Luci's function to dilate the L/R mask after running this, and feed its output into chirality correction logger.info("Left/right image registration completed") - chiral_out_dir = run_chirality_correction(left_right_mask_nifti_fpath, # TODO Rename to mention that this also does registration? - j_args, logger) + dilated_LRmask_fpath = dilate_LR_mask( + os.path.join(j_args["optional_out_dirs"]["postBIBSnet"], *sub_ses), + left_right_mask_nifti_fpath + ) + logger.info("Finished dilating left/right segmentation mask") + + nii_outfpath = run_chirality_correction(dilated_LRmask_fpath, # left_right_mask_nifti_fpath, # TODO Rename to mention that this also does registration? + j_args, logger) + chiral_out_dir = os.path.dirname(nii_outfpath) logger.info("The BIBSnet segmentation has had its chirality checked and " - "registered if needed. Now making aseg-") + "registered if needed. Now making aseg-derived mask.") - masks = make_asegderived_mask(j_args, chiral_out_dir) # NOTE Mask must be in native T1 space too + # TODO Skip mask creation if outputs already exist and not j_args[common][overwrite] + aseg_mask = make_asegderived_mask(j_args, chiral_out_dir, nii_outfpath) # NOTE Mask must be in native T1 space too logger.info("A mask of the BIBSnet segmentation has been produced") # Make nibabies input dirs derivs_dir = os.path.join(j_args["optional_out_dirs"]["derivatives"], "precomputed", *sub_ses, "anat") os.makedirs(derivs_dir, exist_ok=True) + copy_to_derivatives_dir(nii_outfpath, derivs_dir, sub_ses, "aseg_dseg") + """ for eachfile in os.scandir(chiral_out_dir): - if "final" in os.path.basename(eachfile): + if "native" in os.path.basename(eachfile): copy_to_derivatives_dir(eachfile, derivs_dir, sub_ses, "aseg_dseg") # TODO Can these be symlinks? - for each_mask in masks: # There should only be 1, for the record - copy_to_derivatives_dir(each_mask, derivs_dir, sub_ses, "brain_mask") # TODO Can these be symlinks? + """ + copy_to_derivatives_dir(aseg_mask, derivs_dir, sub_ses, "brain_mask") # TODO Get dataset_description.json and put it in derivs_dir @@ -396,8 +407,7 @@ def run_left_right_registration(j_args, sub_ses, age_months, t1or2, logger): # Paths for left & right registration chiral_in_dir = os.path.join(SCRIPT_DIR, "data", "chirality_masks") tmpl_head = os.path.join(chiral_in_dir, "{}mo_T{}w_acpc_dc_restore.nii.gz") - tmpl_mask = os.path.join(chiral_in_dir, "brainmasks", - "{}mo_template_brainmask.nii.gz") + tmpl_mask = os.path.join(chiral_in_dir, "{}mo_template_LRmask.nii.gz") # "brainmasks", {}mo_template_brainmask.nii.gz") # Grab the first resized T?w from preBIBSnet to use for L/R registration first_subject_head = glob(os.path.join( @@ -436,7 +446,7 @@ def run_left_right_registration(j_args, sub_ses, age_months, t1or2, logger): logger.info(msg.format("Skipping", "{} because output already exists at {}".format( first_subject_head, left_right_mask_nifti_fpath ))) - logger.info(msg.format("Finished", first_subject_head)) + logger.info(msg.format("Finished", first_subject_head)) # TODO Only print this message if not skipped (and do the same for all other stages) return left_right_mask_nifti_fpath @@ -473,40 +483,31 @@ def run_chirality_correction(l_r_mask_nifti_fpath, j_args, logger): *sub_ses, "anat", "*_T1w.nii.gz"))[0] # Run chirality correction script - correct_chirality(seg_BIBSnet_outfiles[0], segment_lookup_table_path, - l_r_mask_nifti_fpath, chiral_out_dir, - path_T1w, j_args, logger) - return chiral_out_dir + nii_outfpath = correct_chirality(seg_BIBSnet_outfiles[0], segment_lookup_table_path, + l_r_mask_nifti_fpath, chiral_out_dir, + path_T1w, j_args, logger) + return nii_outfpath # chiral_out_dir -def make_asegderived_mask(j_args, aseg_dir): +def make_asegderived_mask(j_args, aseg_dir, nii_outfpath): """ Create mask file(s) derived from aseg file(s) in aseg_dir :param aseg_dir: String, valid path to existing directory with output files from chirality correction :return: List of strings; each is a valid path to an aseg mask file """ - # Only make masks from chirality-corrected nifti files - sub_ses = get_subj_ID_and_session(j_args) - aseg = glob(os.path.join(j_args["optional_out_dirs"]["postBIBSnet"], - *sub_ses, "chirality_correction", "final*.nii.gz")) # TODO Does this glob need to be more specific? - # aseg = glob(os.path.join(aseg_dir, "sub-*_aseg.nii.gz")) - # aseg.sort() - # binarize, fillh, and erode aseg to make mask: - mask_out_files = list() - for aseg_file in aseg: - mask_out_files.append(os.path.join( - aseg_dir, "{}_mask.nii.gz".format(aseg_file.split(".nii.gz")[0]) - )) - if (j_args["common"]["overwrite"] or not - os.path.exists(mask_out_files[-1])): - maths = fsl.ImageMaths(in_file=aseg_file, # (anat file) - op_string=("-bin -dilM -dilM -dilM -dilM " - "-fillh -ero -ero -ero -ero"), - out_file=mask_out_files[-1]) - maths.run() - return mask_out_files + output_mask_fpath = os.path.join( + aseg_dir, "{}_mask.nii.gz".format(nii_outfpath.split(".nii.gz")[0]) + ) + if (j_args["common"]["overwrite"] or not + os.path.exists(output_mask_fpath)): + maths = fsl.ImageMaths(in_file=nii_outfpath, # (anat file) + op_string=("-bin -dilM -dilM -dilM -dilM " + "-fillh -ero -ero -ero -ero"), + out_file=output_mask_fpath) + maths.run() + return output_mask_fpath def copy_to_derivatives_dir(file_to_copy, derivs_dir, sub_ses, new_fname_pt): @@ -517,7 +518,7 @@ def copy_to_derivatives_dir(file_to_copy, derivs_dir, sub_ses, new_fname_pt): :param sub_ses: List with either only the subject ID str or the session too :param new_fname_pt: String to add to the end of the new filename """ - copy_and_rename_file(file_to_copy, os.path.join(derivs_dir, ( + shutil.copy2(file_to_copy, os.path.join(derivs_dir, ( "{}_space-orig_desc-{}.nii.gz".format("_".join(sub_ses), new_fname_pt) ))) diff --git a/src/utilities.py b/src/utilities.py index 2524d57..dc576fc 100755 --- a/src/utilities.py +++ b/src/utilities.py @@ -5,7 +5,7 @@ Common source for utility functions used by CABINET :) Greg Conan: gconan@umn.edu Created: 2021-11-12 -Updated: 2022-04-29 +Updated: 2022-05-01 """ # Import standard libraries import argparse @@ -14,7 +14,7 @@ from nipype.interfaces import fsl import numpy as np import os -import pdb # TODO Remove this line, which is used for debugging by adding pdb.set_trace() before a problematic line +import pdb # TODO Remove this line(?), which is used for debugging by adding pdb.set_trace() before a problematic line import shutil import subprocess import sys @@ -180,17 +180,6 @@ def check_and_correct_region(should_be_left, region, segment_name_to_number, new_data[chirality][floor_ceiling][scanner_bore] = flipped_id -def copy_and_rename_file(old_file, new_file): - """ - Rename a file and copy it to a new location - :param old_file: String, valid path to an existing file to copy - :param new_file: String, valid path to what will be a copy of old_file - """ - # shutil.move() - # os.rename(shutil.copy2(old_file, os.path.dirname(new_file)), new_file) - shutil.copy2(old_file, new_file) - - def correct_chirality(nifti_input_file_path, segment_lookup_table, left_right_mask_nifti_file, chiral_out_dir, t1w_path, j_args, logger): @@ -210,7 +199,7 @@ def correct_chirality(nifti_input_file_path, segment_lookup_table, chiral_out_dir, "corrected_" + os.path.basename(nifti_input_file_path) )# j_args["BIBSnet"]["aseg_outfile"]) nifti_output_file_path = os.path.join( - chiral_out_dir, "final_" + os.path.basename(nifti_input_file_path) + chiral_out_dir, "native_" + os.path.basename(nifti_input_file_path) )# j_args["BIBSnet"]["aseg_outfile"]) logger.info(msg.format("Running", nifti_input_file_path)) @@ -243,11 +232,13 @@ def correct_chirality(nifti_input_file_path, segment_lookup_table, fixed_img = nib.Nifti1Image(new_data, img.affine, img.header) nib.save(fixed_img, nifti_corrected_file_path) + # TODO Make everything below its own function called "reverse_registration" or "revert_to_native" or something + # Undo resizing right here (do inverse transform) using RobustFOV so padding isn't necessary; revert aseg to native space dummy_copy = "_dummy".join(split_2_exts(nifti_corrected_file_path)) - copy_and_rename_file(nifti_corrected_file_path, dummy_copy) - concat_preBIBSnet_xfms = os.path.join(chiral_out_dir, "{}_concatenated.mat" # TODO Give this a more descriptive name? - .format("_".join(sub_ses))) + shutil.copy2(nifti_corrected_file_path, dummy_copy) + # concat_preBIBSnet_xfms = os.path.join(chiral_out_dir, "{}_concatenated.mat".format("_".join(sub_ses))) # TODO Give this a more descriptive name? Or just delete it? + seg_to_T1w_nat = os.path.join(chiral_out_dir, "seg_reg_to_T1w_native.mat") preBIBSnet_mat = os.path.join(j_args["optional_out_dirs"]["postBIBSnet"], *sub_ses, "preBIBSnet_crop_T1w_to_BIBS_template.mat") # "preBIBSnet_T1w_final.mat") crop_T{}w_to_BIBS_template.mat @@ -256,8 +247,9 @@ def correct_chirality(nifti_input_file_path, segment_lookup_table, run_FSL_sh_script(j_args, logger, "flirt", "-applyxfm", "-ref", t1w_path, "-in", dummy_copy, "-init", seg_to_T1w_nat, # concat_preBIBSnet_xfms, # TODO -applyxfm might need to be changed to -applyisoxfm with resolution - "-o", nifti_output_file_path) + "-o", nifti_output_file_path, "-interp", "nearestneighbour") logger.info(msg.format("Finished", nifti_input_file_path)) + return nifti_output_file_path def create_anatomical_average(avg_params): @@ -326,6 +318,107 @@ def dict_has(a_dict, a_key): return a_key in a_dict and a_dict[a_key] +def dilate_LR_mask(sub_LRmask_dir, anatfile): + """ + Taken from https://github.com/DCAN-Labs/SynthSeg/blob/master/SynthSeg/dcan/img_processing/chirality_correction/dilate_LRmask.py + :param sub_LRmask_dir: String, valid path to placeholder + :param anatfile: String, valid path to placeholder + """ + os.chdir(sub_LRmask_dir) + if not os.path.exists('lrmask_dil_wd'): + os.mkdir('lrmask_dil_wd') + + # anatfile = 'LRmask.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile, op_string='-thr 1 -uthr 1', + out_file='lrmask_dil_wd/Lmask.nii.gz') + maths.run() + + maths = fsl.ImageMaths(in_file=anatfile, op_string='-thr 2 -uthr 2', + out_file='lrmask_dil_wd/Rmask.nii.gz') + maths.run() + + maths.run() + maths = fsl.ImageMaths(in_file=anatfile, op_string='-thr 3 -uthr 3', + out_file='lrmask_dil_wd/Mmask.nii.gz') + maths.run() + + # dilate, fill, and erode each mask in order to get rid of holes + # (also binarize L and M images in order to perform binary operations) + anatfile = 'lrmask_dil_wd/Lmask.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile, op_string='-dilM -dilM -dilM -fillh -ero', + out_file='lrmask_dil_wd/L_mask_holes_filled.nii.gz') + maths.run() + + anatfile = 'lrmask_dil_wd/Rmask.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile, op_string='-bin -dilM -dilM -dilM -fillh -ero', + out_file='lrmask_dil_wd/R_mask_holes_filled.nii.gz') + maths.run() + + anatfile = 'lrmask_dil_wd/Mmask.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile, op_string='-bin -dilM -dilM -dilM -fillh -ero', + out_file='lrmask_dil_wd/M_mask_holes_filled.nii.gz') + maths.run() + + # Reassign values of 2 and 3 to R and M masks (L mask already a value of 1) + anatfile = 'lrmask_dil_wd/R_mask_holes_filled.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile, op_string='-mul 2', + out_file='lrmask_dil_wd/R_mask_holes_filled_label2.nii.gz') + maths.run() + + anatfile = 'lrmask_dil_wd/M_mask_holes_filled.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile, op_string='-mul 3', + out_file='lrmask_dil_wd/M_mask_holes_filled_label3.nii.gz') + maths.run() + + # recombine new L, R, and M mask files + anatfile_left = 'lrmask_dil_wd/L_mask_holes_filled.nii.gz' + anatfile_right = 'lrmask_dil_wd/R_mask_holes_filled_label2.nii.gz' + anatfile_mid = 'lrmask_dil_wd/M_mask_holes_filled_label3.nii.gz' + maths = fsl.ImageMaths(in_file=anatfile_left, op_string='-add {}'.format(anatfile_right), + out_file='lrmask_dil_wd/recombined_mask_LR.nii.gz') + maths.run() + + maths = fsl.ImageMaths(in_file=anatfile_mid, op_string='-add lrmask_dil_wd/recombined_mask_LR.nii.gz', + out_file='lrmask_dil_wd/dilated_LRmask.nii.gz') + maths.run() + + ## Fix incorrect values resulting from recombining dilated components + orig_LRmask_img = nib.load('LRmask.nii.gz') + orig_LRmask_data = orig_LRmask_img.get_fdata() + + fill_LRmask_img = nib.load('lrmask_dil_wd/dilated_LRmask.nii.gz') + fill_LRmask_data = fill_LRmask_img.get_fdata() + + # Flatten numpy arrays + orig_LRmask_data_2D = orig_LRmask_data.reshape((182, 39676), order='C') + orig_LRmask_data_1D = orig_LRmask_data_2D.reshape(7221032, order='C') + + fill_LRmask_data_2D = fill_LRmask_data.reshape((182, 39676), order='C') + fill_LRmask_data_1D = fill_LRmask_data_2D.reshape(7221032, order='C') + + # grab index values of voxels with a value greater than 2.0 in filled L/R mask + voxel_check = np.where(fill_LRmask_data_1D > 2.0) + + # Replace possible overlapping label values with corresponding label values from initial mask + for i in voxel_check[:]: + fill_LRmask_data_1D[i] = orig_LRmask_data_1D[i] + + # reshape numpy array + fill_LRmask_data_2D = fill_LRmask_data_1D.reshape((182, 39676), order='C') + fill_LRmask_data_3D = fill_LRmask_data_2D.reshape((182, 218, 182), order='C') + + # save new numpy array as image + empty_header = nib.Nifti1Header() + out_img = nib.Nifti1Image(fill_LRmask_data_3D, orig_LRmask_img.affine, empty_header) + out_fpath = os.path.join(sub_LRmask_dir, 'LRmask_dil.nii.gz') + nib.save(out_img, out_fpath) + + #remove working directory with intermediate outputs + #shutil.rmtree('lrmask_dil_wd') + + return out_fpath + + def ensure_dict_has(a_dict, a_key, new_value): """ :param a_dict: Dictionary (any) @@ -568,7 +661,7 @@ def optimal_realigned_imgs(xfm_imgs_non_ACPC, xfm_imgs_ACPC_and_reg, j_args, log # Create symlinks with the same name regardless of which is chosen, so # postBIBSnet can use the correct/chosen .mat file concat_mat = optimal_resize["T1w_crop2BIBS_mat"] - out_mat_fpath = os.path.join( + out_mat_fpath = os.path.join( # TODO Pass this in (or out) from the beginning so we don't have to build the path twice (once here and once in postBIBSnet) j_args["optional_out_dirs"]["postBIBSnet"], *sub_ses, "preBIBSnet_" + os.path.basename(concat_mat) ) @@ -857,6 +950,8 @@ def run_FSL_sh_script(j_args, logger, fsl_fn_name, *fsl_args): :param fsl_fn_name: String naming the FSL function which is an executable file in j_args[common][fsl_bin_path] """ + # TODO Run FSL commands using the Python fsl.ImageMaths /etc functions instead of subprocess + # FSL command to (maybe) run in a subprocess to_run = [os.path.join(j_args["common"]["fsl_bin_path"], fsl_fn_name) ] + [str(f) for f in fsl_args]