Dev deeplearning Integration (nnLM and new models)

autoafids-dev.yml

@@ -11,4 +11,4 @@ dependencies:
   - snakefmt >=0.8.4,<0.9.0
   - yamlfix >=1.11.0,<2.0.0
   - pygraphviz ==1.7
-  - jinja2 >=3.0.3,<4.0.0
+  - jinja2 >=3.0.3,<4.0.0

autoafids/config/snakebids.yml

@@ -180,6 +180,73 @@ parse_args:
     default: work
     type: str
 
+  --detect_with_prior:
+    help: |
+      (Default) Use MNI-registered prior FCSV to place 5 patches around the
+      expected AFID location.  Fast and accurate when a good MNI registration
+      is available.
+    dest: detect_mode
+    action: store_const
+    const: prior
+    default: prior
+
+  --detect_without_prior:
+    help: |
+      Use whole-volume sliding-window inference — no prior registration needed.
+      Slower (scans the full image) but does not depend on MNI registration
+      quality.  Overlap is controlled by --inference-overlap (default 0.5).
+    dest: detect_mode
+    action: store_const
+    const: noprior
+
+  --detect_with_nnlm:
+    help: |
+      Use nnLandmark (nnLM) for whole-volume, single-pass AFID detection.
+      All 32 AFIDs are predicted in one nnU-Net forward pass.
+      The model is downloaded automatically on first use.
+    dest: detect_mode
+    action: store_const
+    const: nnlm
+
+  --nnlm_fold:
+    help: "nnLM fold to use for prediction. (default: %(default)s)"
+    dest: nnlm_fold
+    default: "0"
+    type: str
+
+  --nnlm_plans:
+    help: "nnLM plans identifier. (default: %(default)s)"
+    dest: nnlm_plans
+    default: "nnUNetResEncUNetMPlans"
+    type: str
+
+  --nnlm_checkpoint:
+    help: "nnLM checkpoint filename inside the model folder. (default: %(default)s)"
+    dest: nnlm_checkpoint
+    default: "checkpoint_final.pth"
+    type: str
+
+  --nnlm_device:
+    help: "Device for nnLM inference: cuda or cpu. (default: %(default)s)"
+    dest: nnlm_device
+    default: "cpu"
+    type: str
+
+  --inference-overlap:
+    help: |
+      Sliding-window overlap for --detect_without_prior (0.0 = no overlap,
+      0.5 = 50%% overlap, 0.75 = 75%% overlap). Lower = fewer patches = faster.
+      Overrides afids_inference.overlap in snakebids.yml. (default: 0.5)
+    dest: inference_overlap
+    type: float
+
+  --inference-batch-size:
+    help: |
+      Number of patches per model forward pass. Larger = faster but more memory.
+      Overrides afids_inference.batch_size in snakebids.yml. (default: 7)
+    dest: inference_batch_size
+    type: int
+
 #--- workflow specific configuration ---
 
 # Nifti template
@@ -198,12 +265,61 @@ singularity:
 # It will be downloaded to ~/.cache/autoafids
 resource_urls:
   default: 'files.osf.io/v1/resources/9fptg/providers/osfstorage/?zip='
+  # nnLM model checkpoint (whole-volume 32-AFID nnLandmark model)
+  # TODO: Replace with the actual OSF/Zenodo URL when the model is published
+  nnlm: 'https://zenodo.org/records/18991189/files/nnlm_model.zip'
+
+# Sequential inference configuration
+enable_sequential_inference: False
 
 #Stereotaxy models 
 STN: 'resources/stereotaxy/STN.pkl'
 cZI: 'resources/stereotaxy/STN.pkl' #change to czi model but dummy here for testing
 template_fcsv: 'resources/stereotaxy/target_template.fcsv'
 
+# -----------------------------------------------------------------------
+# PyTorch nnUNet inference configuration (used by apply_with_prior.py)
+# Set one checkpoint path per fiducial (afid_01 … afid_32).
+# patch_size, batch_size, and device are optional (defaults shown).
+# -----------------------------------------------------------------------
+afids_inference:
+  checkpoints:
+    afid_01: 'resources/afids_cnn_ckpts/afid-01_epoch-820_mae-0.0002.ckpt'
+    afid_02: 'resources/afids_cnn_ckpts/afid-02_epoch-964_mae-0.0007.ckpt'
+    afid_03: 'resources/afids_cnn_ckpts/afid-03_epoch-778_mae-0.0008.ckpt'
+    afid_04: 'resources/afids_cnn_ckpts/afid-04_epoch-529_mae-0.0007.ckpt'  
+    afid_05: 'resources/afids_cnn_ckpts/afid-05_epoch-628_mae-0.0007.ckpt'  
+    afid_06: 'resources/afids_cnn_ckpts/afid-06_epoch-485_mae-0.0009.ckpt'  
+    afid_07: 'resources/afids_cnn_ckpts/afid-07_epoch-599_mae-0.0008.ckpt'  
+    afid_08: 'resources/afids_cnn_ckpts/afid-08_epoch-455_mae-0.0013.ckpt' 
+    afid_09: 'resources/afids_cnn_ckpts/afid-09_epoch-535_mae-0.0010.ckpt'  
+    afid_10: 'resources/afids_cnn_ckpts/afid-10_epoch-431_mae-0.0009.ckpt'  
+    afid_11: 'resources/afids_cnn_ckpts/afid-11_epoch-735_mae-0.0005.ckpt'  
+    afid_12: 'resources/afids_cnn_ckpts/afid-12_epoch-694_mae-0.0007.ckpt'  
+    afid_13: 'resources/afids_cnn_ckpts/afid-13_epoch-763_mae-0.0006.ckpt'  
+    afid_14: 'resources/afids_cnn_ckpts/afid-14_epoch-405_mae-0.0011.ckpt'  
+    afid_15: 'resources/afids_cnn_ckpts/afid-15_epoch-371_mae-0.0013.ckpt'  
+    afid_16: 'resources/afids_cnn_ckpts/afid-16_epoch-454_mae-0.0012.ckpt'  
+    afid_17: 'resources/afids_cnn_ckpts/afid-17_epoch-291_mae-0.0023.ckpt'  
+    afid_18: 'resources/afids_cnn_ckpts/afid-18_epoch-233_mae-0.0016.ckpt'  
+    afid_19: 'resources/afids_cnn_ckpts/afid-19_epoch-626_mae-0.0011.ckpt'  
+    afid_20: 'resources/afids_cnn_ckpts/afid-20_epoch-720_mae-0.0011.ckpt'  
+    afid_21: 'resources/afids_cnn_ckpts/afid-21_epoch-672_mae-0.0009.ckpt'  
+    afid_22: 'resources/afids_cnn_ckpts/afid-22_epoch-647_mae-0.0008.ckpt'   
+    afid_23: 'resources/afids_cnn_ckpts/afid-23_epoch-656_mae-0.0010.ckpt'   
+    afid_24: 'resources/afids_cnn_ckpts/afid-24_epoch-689_mae-0.0007.ckpt'   
+    afid_25: 'resources/afids_cnn_ckpts/afid-25_epoch-513_mae-0.0013.ckpt'   
+    afid_26: 'resources/afids_cnn_ckpts/afid-26_epoch-784_mae-0.0009.ckpt'  
+    afid_27: 'resources/afids_cnn_ckpts/afid-27_epoch-300_mae-0.0014.ckpt'  
+    afid_28: 'resources/afids_cnn_ckpts/afid-28_epoch-351_mae-0.0009.ckpt'  
+    afid_29: 'resources/afids_cnn_ckpts/afid-29_epoch-394_mae-0.0014.ckpt'  
+    afid_30: 'resources/afids_cnn_ckpts/afid-30_epoch-618_mae-0.0012.ckpt'  
+    afid_31: 'resources/afids_cnn_ckpts/afid-31_epoch-573_mae-0.0009.ckpt'  
+    afid_32: 'resources/afids_cnn_ckpts/afid-32_epoch-561_mae-0.0007.ckpt'  
+  patch_size: 64      # cubic patch size in voxels
+  device: cuda:0       # cpu recommended for parallel mode; cuda:0 for sequential
+  overlap: 0.5        # sliding-window overlap for --detect_without_prior (0.0–0.75)
+
 plugins.validator.skip: False
 root: results
 workdir: null

autoafids/workflow/Snakefile

@@ -1,8 +1,16 @@
 import snakebids
-from snakebids import bids, generate_inputs, get_wildcard_constraints
+from snakebids import bids, generate_inputs, get_wildcard_constraints, set_bids_spec
 from appdirs import AppDirs
 import warnings
 
+# Suppress snakebids "unrecognized entity" warnings for the {afid} wildcard
+warnings.filterwarnings(
+    "ignore",
+    message="Path generated with unrecognized entities",
+    category=UserWarning,
+)
+set_bids_spec("v0_0_0")
+
 try:
     from autoafids.workflow.lib import (
         utils as utils,  # Works when run as a package
@@ -345,7 +353,19 @@ rule mni2subfids:
         "scripts/tform_script.py"
 
 
-include: "rules/cnn.smk"
+# ── Detection mode: choose rules to include ──────────────────────────────
+if config.get("detect_mode", "prior") == "nnlm":
+
+    # nnLandmark mode: single-pass whole-volume detection
+    include: "rules/nnlm.smk"
+
+else:
+    # Legacy CNN modes (prior / noprior)
+    # If running on GPU, automatically enable sequential inference to avoid spawning 32 parallel jobs
+    if config.get("afids_inference", {}).get("device") == "cuda:0":
+        config["enable_sequential_inference"] = True
+
+    include: "rules/cnn.smk"
 
 
 if config["fidqc"]:
@@ -360,13 +380,24 @@ if config["LEAD_DBS_DIR"] or config["FMRIPREP_DIR"]:
     include: "rules/regqc.smk"
 
 
+# Select the correct FCSV output descriptor based on inference mode.
+# prior   → applyfidmodel_gather          → desc="afidscnn"
+# noprior → applyfidmodel_noprior_gather  → desc="afidscnn-noprior"
+# nnlm    → nnlm_to_fcsv                 → desc="afidscnn-nnlm"
+_fcsv_desc = {
+    "prior": "afidscnn",
+    "noprior": "afidscnn-noprior",
+    "nnlm": "afidscnn-nnlm",
+}.get(config.get("detect_mode", "prior"), "afidscnn")
+
+
 rule all:
     input:
         models=inputs[config["modality"]].expand(
             bids(
                 root=root,
                 datatype="afids-cnn",
-                desc="afidscnn",
+                desc=_fcsv_desc,
                 suffix="afids.fcsv",
                 **inputs[config["modality"]].wildcards,
             ),

autoafids/workflow/envs/nnlm.yaml

@@ -0,0 +1,15 @@
+---
+name: nnlm
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.12
+  - pytorch=2.5.1
+  - torchvision=0.20.1
+  - pytorch-cuda=12.4
+  - pip
+  - pip:
+      - nnlandmark @ git+https://github.com/MIC-DKFZ/nnLandmark.git

autoafids/workflow/envs/pytorch.yaml

@@ -0,0 +1,13 @@
+---
+name: pytorch
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+dependencies:
+  # Match the known-fast training/inference stack (e.g. torch==2.7.1+cu126)
+  - python=3.10
+  - pytorch=2.7.1
+  - nibabel              # nib.nifti1.load
+  - numpy                # np
+  - pandas               # pd (fcsv loading)