add automatic registration input argument

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,

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:

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