From 628e51591748e9e6eb626ef18d81fcb3d7a59cdd Mon Sep 17 00:00:00 2001
From: Ali Khan <alik@robarts.ca>
Date: Thu, 15 Feb 2024 13:36:20 -0500
Subject: [PATCH] moved download model rules back to nnunet

avoids breaking the seg* modalities
---
 hippunfold/workflow/rules/download.smk    | 28 -----------------------
 hippunfold/workflow/rules/nnunet.smk      | 28 +++++++++++++++++++++++
 hippunfold/workflow/rules/preproc_seg.smk |  2 +-
 3 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/hippunfold/workflow/rules/download.smk b/hippunfold/workflow/rules/download.smk
index eedbe1a6..07c1bcdd 100644
--- a/hippunfold/workflow/rules/download.smk
+++ b/hippunfold/workflow/rules/download.smk
@@ -3,34 +3,6 @@
 download_dir = get_download_dir()
 
 
-def get_model_tar():
-
-    if config["force_nnunet_model"]:
-        model_name = config["force_nnunet_model"]
-    else:
-        model_name = config["modality"]
-
-    local_tar = config["resource_urls"]["nnunet_model"].get(model_name, None)
-    if local_tar == None:
-        print(f"ERROR: {model_name} does not exist in nnunet_model in the config file")
-
-    return (Path(download_dir) / "model" / Path(local_tar).name).absolute()
-
-
-rule download_nnunet_model:
-    params:
-        url=config["resource_urls"]["nnunet_model"][config["force_nnunet_model"]]
-        if config["force_nnunet_model"]
-        else config["resource_urls"]["nnunet_model"][config["modality"]],
-        model_dir=Path(download_dir) / "model",
-    output:
-        model_tar=get_model_tar(),
-    container:
-        config["singularity"]["autotop"]
-    shell:
-        "mkdir -p {params.model_dir} && wget https://{params.url} -O {output}"
-
-
 rule download_extract_atlas_or_template:
     params:
         url=lambda wildcards: config["resource_urls"][wildcards.resource_type][
diff --git a/hippunfold/workflow/rules/nnunet.smk b/hippunfold/workflow/rules/nnunet.smk
index a255b9f9..84e76103 100644
--- a/hippunfold/workflow/rules/nnunet.smk
+++ b/hippunfold/workflow/rules/nnunet.smk
@@ -41,6 +41,34 @@ def get_nnunet_input(wildcards):
     return nii
 
 
+def get_model_tar():
+
+    if config["force_nnunet_model"]:
+        model_name = config["force_nnunet_model"]
+    else:
+        model_name = config["modality"]
+
+    local_tar = config["resource_urls"]["nnunet_model"].get(model_name, None)
+    if local_tar == None:
+        print(f"ERROR: {model_name} does not exist in nnunet_model in the config file")
+
+    return (Path(download_dir) / "model" / Path(local_tar).name).absolute()
+
+
+rule download_nnunet_model:
+    params:
+        url=config["resource_urls"]["nnunet_model"][config["force_nnunet_model"]]
+        if config["force_nnunet_model"]
+        else config["resource_urls"]["nnunet_model"][config["modality"]],
+        model_dir=Path(download_dir) / "model",
+    output:
+        model_tar=get_model_tar(),
+    container:
+        config["singularity"]["autotop"]
+    shell:
+        "mkdir -p {params.model_dir} && wget https://{params.url} -O {output}"
+
+
 def parse_task_from_tar(wildcards, input):
     match = re.search("Task[0-9]{3}_[\w]+", input.model_tar)
     if match:
diff --git a/hippunfold/workflow/rules/preproc_seg.smk b/hippunfold/workflow/rules/preproc_seg.smk
index 11a62cfa..1f9afab0 100644
--- a/hippunfold/workflow/rules/preproc_seg.smk
+++ b/hippunfold/workflow/rules/preproc_seg.smk
@@ -50,7 +50,7 @@ rule warp_seg_to_corobl_crop:
             **config["subj_wildcards"],
             suffix="dseg.nii.gz",
             space="corobl",
-            hemi="{hemi}",
+            hemi="{hemi,L|R}",
             from_="{space}"
         ),
     container: