Merge pull request #302 from Bradley-Karat/AD

AD branch
khanlab · Jun 20, 2024 · e1c8c87 · e1c8c87
2 parents 11e3a73 + 7a76473
commit e1c8c87
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/hippunfold/config/snakebids.yml b/hippunfold/config/snakebids.yml
@@ -326,6 +326,7 @@ parse_args:
       - synthseg_v0.1
       - synthseg_v0.2
       - neonateT1w_v2
+      - AD_T1w
 
 autotop_labels:
   - 'hipp'
@@ -529,6 +530,7 @@ resource_urls:
     T1T2w: 'zenodo.org/record/4508747/files/trained_model.3d_fullres.Task103_hcp1200_T1T2w.nnUNetTrainerV2.model_best.tar'
     synthseg_v0.1: 'zenodo.org/record/8184230/files/trained_model.3d_fullres.Task102_synsegGenDetailed.nnUNetTrainerV2.model_best.tar'
     synthseg_v0.2: 'zenodo.org/record/8184230/files/trained_model.3d_fullres.Task203_synthseg.nnUNetTrainerV2.model_best.tar'
+    AD_T1w: 'trained_model.3d_fullres.Task301_ADNI_T1w_successful.nnUNetTrainer.tar'
   atlas:
     multihist7: 'files.ca-1.osf.io/v1/resources/v8acf/providers/osfstorage/65395b782827451220b86dd8/?zip='
     bigbrain: 'files.ca-1.osf.io/v1/resources/v8acf/providers/osfstorage/65395b8b13d27b123094c96f/?zip='

diff --git a/hippunfold/workflow/rules/nnunet.smk b/hippunfold/workflow/rules/nnunet.smk
@@ -53,10 +53,19 @@ def get_model_tar():
 
     return (Path(download_dir) / "model" / Path(local_tar).name).absolute()
 
+def get_url():
+    if config["force_nnunet_model"]:
+        if config["force_nnunet_model"] == 'AD_T1w':
+            url='osf.io/ug3dz/download'  
+        else:
+            url=config["resource_urls"]["nnunet_model"][config["force_nnunet_model"]] 
+    else:
+        url=config["resource_urls"]["nnunet_model"][config["modality"]] 
+    return(url)
 
 rule download_nnunet_model:
     params:
-        url=config["resource_urls"]["nnunet_model"][config["force_nnunet_model"]]
+        url=get_url()
         if config["force_nnunet_model"]
         else config["resource_urls"]["nnunet_model"][config["modality"]],
         model_dir=Path(download_dir) / "model",