diff --git a/fsub_extractor/cli_starters/extractor_start.py b/fsub_extractor/cli_starters/extractor_start.py index 12c5854..acde8a5 100644 --- a/fsub_extractor/cli_starters/extractor_start.py +++ b/fsub_extractor/cli_starters/extractor_start.py @@ -4,9 +4,9 @@ from pathlib import Path from fsub_extractor.functions.extractor import extractor + # Add input arguments def get_parser(): - parser = argparse.ArgumentParser( description="Functionally segments a tract file based on intersections with prespecified ROI(s) in gray matter." ) @@ -172,9 +172,16 @@ def get_parser(): # Registration arguments reg_group = parser.add_argument_group("Options for Registration") reg_dir_group = reg_group.add_mutually_exclusive_group() + reg_dir_group.add_argument( + "--reg-anat", + "--reg_anat", + help="Paths to T1 image and brain mask (NIfTIs) in DWI space, separated by a comma (no spaces). These are used to calculate a registration if DWI space is not aligned to FreeSurfer (e.g., a QSIPrep T1w image in ACPC space). This overrides --fs2dwi, --dwi2fs, and --reg-type.", + type=validate_file, + metavar=("/PATH/TO/T1.nii.gz,/PATH/TO/mask.nii.gz"), + ) reg_dir_group.add_argument( "--fs2dwi", - help="Path to MRTrix-ready or ANTs/ITK-generated registration for mapping FreeSurfer-to-DWI space. Mutually exclusive with --dwi2fs.", + help="Path to MRTrix-ready or ANTs/ITK-generated registration for mapping FreeSurfer-to-DWI space. This can be generated by first running with '--reg-anat'. Mutually exclusive with --dwi2fs.", type=validate_file, metavar=("/PATH/TO/FS2DWI-REG.txt"), action=CheckExt({".txt"}), @@ -399,7 +406,6 @@ def check_positive_int(value): def main(): - # Parse arguments and run the main code parser = get_parser() args = parser.parse_args() @@ -427,6 +433,7 @@ def main(): exclude_mask=args.exclude_mask, include_mask=args.include_mask, streamline_mask=args.streamline_mask, + reg_anat=args.reg_anat, fs2dwi=args.fs2dwi, dwi2fs=args.dwi2fs, reg_type=args.reg_type, diff --git a/fsub_extractor/functions/extractor.py b/fsub_extractor/functions/extractor.py index 458039e..7ec5b28 100644 --- a/fsub_extractor/functions/extractor.py +++ b/fsub_extractor/functions/extractor.py @@ -30,6 +30,7 @@ def extractor( exclude_mask, include_mask, streamline_mask, + reg_anat, fs2dwi, dwi2fs, reg_type, @@ -129,8 +130,31 @@ def extractor( "The 'delta' paramater of projfrac-params must be positive to iterate correctly." ) + ### Pre-checks are over, begin the processing! + + # Make output folders if they do not exist, and define the naming convention + anat_out_dir = op.join(out_dir, subject, "anat") + dwi_out_dir = op.join(out_dir, subject, "dwi") + func_out_dir = op.join(out_dir, subject, "func") + os.makedirs(anat_out_dir, exist_ok=True) + os.makedirs(dwi_out_dir, exist_ok=True) + os.makedirs(func_out_dir, exist_ok=True) + # Define and prepare registration - if fs2dwi == None and dwi2fs == None: + if reg_anat != None: + t1 = reg_anat.split(",")[0] + t1_mask = reg_anat.split(",")[1] + reg_invert = False + reg = calculate_fs2dwi_reg( + subject, + anat_out_dir, + anat_out_dir, + fs_dir, + t1, + t1_mask, + overwrite=overwrite, + ) + elif fs2dwi == None and dwi2fs == None: reg = None reg_invert = None reg_type = None @@ -155,16 +179,6 @@ def extractor( # XX. Make sure FS license is valid [TODO: HOW??] - ### Pre-checks are over, begin the processing! - - # Make output folders if they do not exist, and define the naming convention - anat_out_dir = op.join(out_dir, subject, "anat") - dwi_out_dir = op.join(out_dir, subject, "dwi") - func_out_dir = op.join(out_dir, subject, "func") - os.makedirs(anat_out_dir, exist_ok=True) - os.makedirs(dwi_out_dir, exist_ok=True) - os.makedirs(func_out_dir, exist_ok=True) - # Prepare registration, if needed if reg != None and reg_type != "mrtrix": if reg_invert: diff --git a/fsub_extractor/utils/anat_utils.py b/fsub_extractor/utils/anat_utils.py index f53cdcc..5dc16d4 100644 --- a/fsub_extractor/utils/anat_utils.py +++ b/fsub_extractor/utils/anat_utils.py @@ -1,5 +1,6 @@ import os.path as op import os +import shlex from fsub_extractor.utils.system_utils import * @@ -144,7 +145,6 @@ def get_pial_surf( anat_out_dir=os.getcwd(), overwrite=True, ): - """Returns volumetric mask of pial surface Parameters @@ -254,3 +254,121 @@ def convert_to_mrtrix_reg(reg_in, mrtrix_reg_out, reg_in_type="itk", overwrite=T run_command(cmd_transformconvert) return mrtrix_reg_out + + +def calculate_fs2dwi_reg( + subject, outdir, scratch_dir, fs_dir, t1, t1_mask, overwrite=True +): + """Makes an LTA convert file for mapping between FreeSurfer and DWI space + + Parameters + ========== + subject: str + Subject name as found in FreeSurfer subjects directory. + outdir: str + Path to output directory. + scratch_dir: str + Path to scratch directory. + fs_dir: str + Path to FreeSurfer subjects directory. + t1: str + Path to T1 image in DWI space (NIfTI). + t1_mask: str + Path to T1 brain mask image in DWI space (NIfTI). + overwrite: bool + Whether to allow overwriting outputs + + Outputs + ======= + Function creates and returns path to MRTrix-readable registration file for mapping between FreeSurfer and DWI space + """ + + # Start by converting FS brain from .mgz to .nii + mrconvert = find_program("mrconvert") + fs_brain_path = op.join(fs_dir, subject, "mri", "brain.mgz") # Input + fs_brain_nii_out = op.join(scratch_dir, f"{subject}_fs_brain.nii") # Output + cmd_mrconvert = [mrconvert, "-stride", "-1,-2,3", fs_brain_path, fs_brain_nii_out] + + if overwrite: + cmd_mrconvert += ["-force"] + else: + overwrite_check(fs_brain_nii_out) + + run_command(cmd_mrconvert) + + # Perform the registration + antsRegistration = find_program("antsRegistration") + cmd_antsRegistration = [ + antsRegistration, + "--collapse-output-transforms", + "1", + "--dimensionality", + "3", + "--float", + "0", + "--initial-moving-transform", + f"[{t1},{fs_brain_nii_out},1]", + "--initialize-transforms-per-stage", + "0", + "--interpolation", + "BSpline", + "--output", + f"[{scratch_dir}/transform,{scratch_dir}/transform_Warped.nii.gz]", + "--transform", + "Rigid[0.1]", + "--metric", + f"Mattes[{t1},{fs_brain_nii_out},1,32,Random,0.25]", + "--convergence", + "[1000x500x250x100,1e-06,10]", + "--smoothing-sigmas", + "3.0x2.0x1.0x0.0mm", + "--shrink-factors", + "8x4x2x1", + "--use-histogram-matching", + "0", + "--masks", + f"[{t1_mask},NULL]", + "--winsorize-image-intensities", + "[0.002,0.998]", + "--write-composite-transform", + "0", + ] + + if overwrite == False: + overwrite_check(f"{scratch_dir}/transform_Warped.nii.gz") + + run_command(cmd_antsRegistration) + + # Convert ANTs .mat transform to .txt, and rename it + ConvertTransformFile = find_program("ConvertTransformFile") + reg_txt = f"{scratch_dir}/{subject}_from-FS_to-T1wACPC_mode-image_xfm.txt" + cmd_ConvertTransformFile = [ + ConvertTransformFile, + "3", + f"{scratch_dir}/transform0GenericAffine.mat", + reg_txt, + ] + + if overwrite == False: + overwrite_check(reg_txt) + + run_command(cmd_ConvertTransformFile) + + # Convert ANTs transform to MRTrix compatible transform, save out + transformconvert = find_program("transformconvert") + mrtrix_reg_out = f"{outdir}/{subject}_from-FS_to-DWI_mode-image_xfm.txt" + cmd_transformconvert = [ + transformconvert, + reg_txt, + "itk_import", + mrtrix_reg_out, + ] + + if overwrite: + cmd_transformconvert += ["-force"] + else: + overwrite_check(mrtrix_reg_out) + + run_command(cmd_transformconvert) + + return mrtrix_reg_out