Skip to content

Commit

Permalink
Merge pull request #6 from DCAN-Labs/develop
Browse files Browse the repository at this point in the history 
Fixed postBIBSnet for now; it and preBIBSnet work
  • Loading branch information
@GregConan
GregConan authored May 9, 2022
2 parents ca28a5f + baac93c commit 3c4660e
Show file tree
Hide file tree
Showing 2 changed files with 160 additions and 64 deletions.
91 changes: 46 additions & 45 deletions run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
Connectome ABCD-XCP niBabies Imaging nnu-NET (CABINET)
Greg Conan: [email protected]
Created: 2021-11-12
Updated: 2022-04-29
Updated: 2022-05-01
"""

# Import standard libraries
Expand All @@ -17,6 +17,7 @@
import os
import pandas as pd
import pdb
import shutil
import subprocess
import sys

Expand All @@ -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
)


Expand Down Expand Up @@ -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"],
Expand Down Expand Up @@ -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

Expand All @@ -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(
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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):
Expand All @@ -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)
)))

Expand Down
133 changes: 114 additions & 19 deletions src/utilities.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
Common source for utility functions used by CABINET :)
Greg Conan: [email protected]
Created: 2021-11-12
Updated: 2022-04-29
Updated: 2022-05-01
"""
# Import standard libraries
import argparse
Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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))
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
)
Expand Down Expand Up @@ -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]
Expand Down

0 comments on commit 3c4660e

Please sign in to comment.