Merge pull request #281 from khanlab/templateseg2

- Template-based segmentation, and new templates (marmoset, macaque) - Downloads external resources (models, templates, atlases) instead of relying on local resources folder
khanlab · Apr 22, 2024 · 234c7f7 · 234c7f7
2 parents 48f96b9 + e063a06
commit 234c7f7
Show file tree

Hide file tree

Showing 25 changed files with 492 additions and 71 deletions.
diff --git a/.dryrun_test_all.sh b/.dryrun_test_all.sh
@@ -10,7 +10,8 @@ hippunfold test_data/bids_hippb500 test_out participant -np --modality hippb500
 hippunfold test_data/bids_T1w_longitudinal test_out participant -np --modality T1w
 hippunfold test_data/bids_singleT2w_longitudinal test_out participant -np --modality T2w
 hippunfold test_data/bids_segT2w test_out participant -np --modality segT2w
-hippunfold - test_out participant -np --modality cropseg --path_cropseg test_data/data_cropseg/sub-{subject}_hemi-{hemi}_dseg.nii.gz
-hippunfold - test_out participant -np --modality cropseg --path_cropseg test_data/data_cropseg_1hemi/sub-{subject}_hemi-{hemi}_dseg.nii.gz --hemi L
+hippunfold . test_out participant -np --modality cropseg --path_cropseg test_data/data_cropseg/sub-{subject}_hemi-{hemi}_dseg.nii.gz
+hippunfold . test_out participant -np --modality cropseg --path_cropseg test_data/data_cropseg_1hemi/sub-{subject}_hemi-{hemi}_dseg.nii.gz --hemi L
 hippunfold test_data/bids_singleT2w test_out participant -np --modality T2w --t1_reg_template
 hippunfold test_data/bids_singleT2w test_out participant -np --modality T2w --output_space T1w
+hippunfold test_data/bids_T1w test_out participant -np --modality T1w --use-template-seg
diff --git a/.github/workflows/python-testing.yml b/.github/workflows/python-testing.yml
@@ -55,7 +55,7 @@ jobs:
 
       - name: Setup env for hippunfold
         run: |
-          echo "HIPPUNFOLD_CACHE_DIR=`pwd`/test_data/fake_models" >> $GITHUB_ENV
+          echo "HIPPUNFOLD_CACHE_DIR=`pwd`/test_data/hippunfold_cache_dir" >> $GITHUB_ENV
 
       - name: Test single T2w bids
         run: |

diff --git a/hippunfold/config/snakebids.yml b/hippunfold/config/snakebids.yml
@@ -141,7 +141,40 @@ parse_args:
       - segT1w
       - segT2w
       - cropseg
-
+
+
+  --template:
+    choices:
+      - 'CITI168'
+      - 'dHCP'
+      - 'MBMv2'
+      - 'MBMv3'
+      - 'CIVM'
+    default: 'CITI168'
+    help: 'Set the template to use for registration to coronal oblique (and optionally for template-based segmentation if --use-template-seg is enabled). CITI168 is for adult human data, dHCP is for neonatal human data, MBMv2 is for ex vivo marmoset data, MBMv3 is for in vivo marmoset data, and CIVM is for in vivo macaque data. (default: %(default)s)'
+
+
+
+
+  --inject_template:
+    choices:
+      - 'upenn'
+      - 'MBMv2'
+      - 'MBMv3'
+      - 'CIVM'
+    default: 'upenn'
+    help: 'Set the template to use for shape injection. (default: %(default)s)'
+
+  --use_template_seg:
+    help: 'Use template-based segmentation for hippocampal tissue *instead of* nnUnet and shape injection. This is only to be used if nnUnet models are not trained for the data you are using, e.g. for non-human primate data with the MBMv2 (ex vivo marmoset), MBMv3 (in vivo marmoset), or CIVM (ex vivo macaque) template. (default: %(default)s)'
+    default: False
+    action: 'store_true'
+
+  --template_seg_smoothing_factor: 
+    help: 'Scales the default smoothing sigma for gradient and warp in greedy registration for template-based segmentation. Using a value higher than 1 will use result in a smoother warp. (default: %(default)s)'
+    default: 2.0
+
+
   --derivatives:
     help: 'Path to the derivatives folder (e.g. for finding manual segs) (default: %(default)s) '
     default: False
@@ -178,13 +211,7 @@ parse_args:
     default: False
     action: 'store_true'
 
-  --template:
-    choices:
-      - 'CITI168'
-      - 'dHCP'
-    default: 'CITI168'
-    help: 'Set the template to use for registration to coronal oblique. (default: %(default)s)'
-
+
   --t1_reg_template:
     help: 'Use T1w to register to template space, instead of the segmentation modality. Note: this was the default behavior prior to v1.0.0.  (default: %(default)s)'
     default: false
@@ -338,6 +365,26 @@ singularity:
   autotop: 'docker://khanlab/hippunfold_deps:v0.5.0'
 
 xfm_identity: resources/etc/identity_xfm.txt
+xfm_identity_itk: resources/etc/identity_xfm_itk.txt
+
+#templates enabled for template-based segmentation are here
+# TODO: should also perhaps include modalities avaialble, and any custom crop_native_res settings
+template_based_segmentation:
+  CITI168:
+    hemi:
+      - R
+      - L
+  CIVM:
+    hemi:
+      - R
+      - L
+  MBMv2:
+    hemi:
+      - R
+  MBMv3:
+    hemi:
+      - R
+
 
 template_files:
   CITI168:
@@ -347,13 +394,42 @@ template_files:
     crop_ref: T2w_300umCoronalOblique_hemi-{hemi}.nii.gz
     crop_refT1w: T1w_300umCoronalOblique_hemi-{hemi}.nii.gz
     Mask_crop: Mask_300umCoronalOblique_hemi-{hemi}.nii.gz
+    dseg:  sub-CITI168_hemi-{hemi}_space-cropT1w_desc-postproc_dseg.nii.gz 
+    coords: sub-CITI168_dir-{dir}_hemi-{hemi}_space-cropT1w_label-{autotop}_desc-laplace_coords.nii.gz 
   dHCP:
     T1w: tpl-dHCP_cohort-1_res-1_T1w.nii.gz
     T2w: tpl-dHCP_cohort-1_res-1_T2w.nii.gz
     xfm_corobl: tpl-dHCP_cohort-1_to-corobl_affine.txt
     crop_ref: tpl-dHCP_cohort-1_res-1_space-corobl_hemi-{hemi}_T2w.nii.gz
     crop_refT1w: tpl-dHCP_cohort-1_res-1_space-corobl_hemi-{hemi}_T1w.nii.gz
     Mask_crop: tpl-dHCP_cohort-1_res-1_space-corobl_hemi-{hemi}_desc-hipp_mask.nii.gz
+  MBMv2:
+    T1w: Template_sym_MTR_80um.nii.gz
+    T2w: Template_sym_T2_80um.nii.gz
+    xfm_corobl: tpl-MBMv2_from-native_to-corobl_type-itk_affine.txt
+    crop_ref: tpl-MBMv2_hemi-{hemi}_space-corobl_{modality}.nii.gz
+    Mask_crop: tpl-MBMv2_hemi-{hemi}_space-corobl_desc-hipp_mask.nii.gz
+    dseg: tpl-MBMv2_hemi-R_space-corobl_desc-tissuemanual_dseg.nii.gz
+    coords: tpl-MBMv2_dir-{dir}_hemi-R_space-corobl_label-{autotop}_desc-laplace_coords.nii.gz
+  MBMv3:
+    T1w: tpl-MBMv3_T1w.nii.gz
+    T2w: tpl-MBMv3_T2w.nii.gz
+    xfm_corobl: tpl-MBMv3_from-native_to-corobl_type-itk_affine.txt
+    crop_ref: tpl-MBMv3_hemi-{hemi}_space-corobl_T1w.nii.gz
+    Mask_crop: tpl-MBMv3_hemi-{hemi}_space-corobl_desc-hipp_mask.nii.gz
+    dseg: tpl-MBMv3_hemi-{hemi}_space-corobl_desc-tissuemanual_dseg.nii.gz
+    coords: tpl-MBMv3_dir-{dir}_hemi-{hemi}_space-corobl_label-{autotop}_desc-laplace_coords.nii.gz
+  CIVM:
+    T1w: tpl-CIVM_T1w.nii.gz
+    xfm_corobl: tpl-CIVM_from-native_to-corobl_type-itk_affine.txt
+    crop_ref: tpl-CIVM_hemi-{hemi}_space-corobl_{modality}.nii.gz
+    Mask_crop: tpl-CIVM_hemi-{hemi}_space-corobl_desc-hipp_mask.nii.gz
+    dseg: tpl-CIVM_hemi-{hemi}_space-corobl_desc-tissuemanual_dseg.nii.gz
+    coords: tpl-CIVM_dir-{dir}_hemi-{hemi}_space-corobl_label-{autotop}_desc-laplace_coords.nii.gz
+  upenn:
+    T1w: tpl-upenn_desc-hipptissue_dseg.nii.gz
+    dseg: tpl-upenn_desc-hipptissue_dseg.nii.gz
+    coords: tpl-upenn_dir-{dir}_label-{autotop}_coords.nii.gz
 
 atlas_files:
   multihist7:
@@ -404,6 +480,8 @@ no_reg_template: False
 
 modality: T2w
 
+inject_template: upenn
+
 template: CITI168
 
 atlas: 
@@ -433,6 +511,7 @@ resource_urls:
     MBMv2: 'files.ca-1.osf.io/v1/resources/v8acf/providers/osfstorage/65395c0887852d133ca597dd/?zip='
     MBMv3: 'files.ca-1.osf.io/v1/resources/v8acf/providers/osfstorage/65395c0e8a28b11240ffc6e9/?zip='
     upenn: 'files.ca-1.osf.io/v1/resources/v8acf/providers/osfstorage/65395c1613d27b122a94ca09/?zip='
+    CIVM: 'files.ca-1.osf.io/v1/resources/v8acf/providers/osfstorage/65395bf62827451220b86e24/?zip='
 #to get hash, see https://github.com/CenterForOpenScience/osf.io/issues/8256#issuecomment-379833911
 
 
@@ -572,3 +651,5 @@ t1_reg_template: False
 generate_myelin_map: False
 no_unfolded_reg: False
 root: results
+use_template_seg: False
+template_seg_smoothing_factor: 2
diff --git a/hippunfold/dags/proc_subgraph.py b/hippunfold/dags/proc_subgraph.py
@@ -19,7 +19,6 @@
 subgraphs = dict()
 
 for smk in glob("../workflow/rules/*.smk"):
-
     smk_name = re.findall(re.compile(r"/(\w+).smk"), smk)[0]
 
     with open(smk, "r") as f:

diff --git a/hippunfold/resources/etc/identity_xfm_itk.txt b/hippunfold/resources/etc/identity_xfm_itk.txt
@@ -0,0 +1,5 @@
+#Insight Transform File V1.0
+#Transform 0
+Transform: MatrixOffsetTransformBase_double_3_3
+Parameters: 1 0 0 0 1 0 0 0 1 0 0 0
+FixedParameters: 0 0 0
diff --git a/hippunfold/workflow/Snakefile b/hippunfold/workflow/Snakefile
@@ -59,7 +59,7 @@ elif "T1w" in config["modality"]:
 
 # only generate hipp surface if skipping template injection
 #   i.e. drop dentate
-if config["skip_inject_template_labels"]:
+if config["skip_inject_template_labels"]:  # TODO do we need this?
     config["autotop_labels"] = ["hipp"]
 
 
@@ -103,6 +103,7 @@ wildcard_constraints:
     density="[a-zA-Z0-9]+",
     atlas="[a-zA-Z0-9]+",
     autotop="[a-zA-Z0-9]+",
+    template="[a-zA-Z0-9]+",
 
 
 root = os.path.join(config["root"], "hippunfold")
@@ -144,9 +145,6 @@ include: "rules/autotop.smk"
 include: "rules/warps.smk"
 
 
-include: "rules/shape_inject.smk"
-
-
 include: "rules/gifti.smk"
 
 
@@ -162,6 +160,26 @@ include: "rules/qc.smk"
 include: "rules/myelin_map.smk"
 
 
+if config["use_template_seg"]:
+    if (
+        "dseg" in config["template_files"][config["template"]]
+        and "coords" in config["template_files"][config["template"]]
+    ):
+
+        include: "rules/templateseg.smk"
+
+    else:
+        print(
+            "use_template_seg is not possible, template dseg or coords do not exist for the chosen template"
+        )
+
+        include: "rules/shape_inject.smk"
+
+else:
+
+    include: "rules/shape_inject.smk"
+
+
 rule all:
     input:
         get_final_output(),

diff --git a/hippunfold/workflow/rules/common.smk b/hippunfold/workflow/rules/common.smk
@@ -24,7 +24,6 @@ def get_modality_suffix(modality):
 
 
 def get_final_spec():
-
     if len(config["hemi"]) == 2:
         specs = expand(
             bids(
@@ -293,20 +292,21 @@ def get_final_qc():
             )
         )
     if (config["modality"] == "T1w") or (config["modality"] == "T2w"):
-        qc.extend(
-            expand(
-                bids(
-                    root=root,
-                    datatype="qc",
-                    desc="unetf3d",
-                    suffix="dice.tsv",
-                    hemi="{hemi}",
-                    **config["subj_wildcards"],
-                ),
-                hemi=config["hemi"],
-                allow_missing=True,
+        if not config["use_template_seg"]:
+            qc.extend(
+                expand(
+                    bids(
+                        root=root,
+                        datatype="qc",
+                        desc="unetf3d",
+                        suffix="dice.tsv",
+                        hemi="{hemi}",
+                        **config["subj_wildcards"],
+                    ),
+                    hemi=config["hemi"],
+                    allow_missing=True,
+                )
             )
-        )
     return qc
 
 
@@ -322,7 +322,6 @@ def get_final_subj_output():
 
 
 def get_final_output():
-
     if config["keep_work"]:
         subj_output = get_final_subj_output()
     else:

diff --git a/hippunfold/workflow/rules/download.smk b/hippunfold/workflow/rules/download.smk
@@ -3,15 +3,25 @@
 download_dir = get_download_dir()
 
 
-rule download_extract_atlas_or_template:
+rule download_extract_atlas:
     params:
-        url=lambda wildcards: config["resource_urls"][wildcards.resource_type][
-            wildcards.atlas
-        ],
+        url=lambda wildcards: config["resource_urls"]["atlas"][wildcards.atlas],
     output:
-        unzip_dir=directory(
-            Path(download_dir) / "{resource_type,atlas|template}" / "{atlas}"
-        ),
+        unzip_dir=directory(Path(download_dir) / "atlas" / "{atlas}"),
+    container:
+        config["singularity"]["autotop"]
+    shadow:
+        "minimal"
+    shell:
+        "wget https://{params.url} -O temp.zip && "
+        " unzip -d {output.unzip_dir} temp.zip"
+
+
+rule download_extract_template:
+    params:
+        url=lambda wildcards: config["resource_urls"]["template"][wildcards.template],
+    output:
+        unzip_dir=directory(Path(download_dir) / "template" / "{template}"),
     container:
         config["singularity"]["autotop"]
     shadow:

diff --git a/hippunfold/workflow/rules/nnunet.smk b/hippunfold/workflow/rules/nnunet.smk
@@ -42,7 +42,6 @@ def get_nnunet_input(wildcards):
 
 
 def get_model_tar():
-
     if config["force_nnunet_model"]:
         model_name = config["force_nnunet_model"]
     else:
@@ -201,7 +200,6 @@ rule unflip_nnunet_nii:
 
 
 def get_f3d_ref(wildcards, input):
-
     if config["modality"] == "T2w":
         nii = Path(input.template_dir) / config["template_files"][config["template"]][
             "crop_ref"

diff --git a/hippunfold/workflow/rules/preproc_t1.smk b/hippunfold/workflow/rules/preproc_t1.smk
@@ -71,7 +71,6 @@ else:
 
 
 def reg_to_template_cmd(wildcards, input, output):
-
     ref = str(
         Path(input.template_dir)
         / config["template_files"][config["template"]][wildcards.modality].format(
@@ -245,7 +244,7 @@ rule template_xfm_itk2ras:
             datatype="warps",
             **config["subj_wildcards"],
             suffix="xfm.txt",
-            from_="{native_modality,T1w|T2w}",
+            from_="{native_modality}",
             to="corobl",
             desc="affine",
             type_="itk"
@@ -292,7 +291,9 @@ rule warp_t1_to_corobl_crop:
         template_dir=Path(download_dir) / "template" / config["template"],
     params:
         ref=lambda wildcards, input: Path(input.template_dir)
-        / config["template_files"][config["template"]]["crop_ref"].format(**wildcards),
+        / config["template_files"][config["template"]]["crop_ref"].format(
+            **wildcards, modality="T1w"
+        ),
     output:
         t1=bids(
             root=work,