Remove k wildcard from atropos_seg, add k-retry logic with gmm_init_k/gmm_min_k config

Agent-Logs-Url: https://github.com/khanlab/SPIMquant/sessions/e025f01c-8e6d-495f-b403-90c8c2aa6710

Co-authored-by: akhanf <11492701+akhanf@users.noreply.github.com>

Author: Copilot

spimquant/config/.ipynb_checkpoints/snakebids-checkpoint.yml

@@ -236,7 +236,8 @@ templatereg:
     - 40
 
 masking:
-  gmm_k: 9
+  gmm_init_k: 9
+  gmm_min_k: 2
   gmm_bg_class: 1
   pre_atropos_downsampling: '50%'
 

spimquant/config/snakebids.yml

@@ -352,7 +352,8 @@ templatereg:
     - 40
 
 masking:
-  gmm_k: 9
+  gmm_init_k: 9
+  gmm_min_k: 2
   gmm_bg_class: 1
   pre_atropos_downsampling: '50%'
 

spimquant/workflow/rules/masking.smk

@@ -74,17 +74,22 @@ rule pre_atropos:
 
 rule atropos_seg:
     """Perform tissue classification using Atropos (k-class GMM).
-    
+
     Uses ANTs Atropos to classify tissue into k intensity classes via Gaussian
     mixture modeling with Markov random field (MRF) spatial smoothing. Outputs
     a discrete segmentation and posterior probability maps for each class.
+
+    Automatically decrements k from init_k down to min_k if Atropos fails,
+    handling cases where the image lacks enough distinct intensity classes.
     """
     input:
         downsampled=rules.pre_atropos.output.downsampled,
         mask=rules.pre_atropos.output.mask,
     params:
         mrf_smoothing=0.3,
         mrf_radius="2x2x2",
+        init_k=config["masking"]["gmm_init_k"],
+        min_k=config["masking"]["gmm_min_k"],
     output:
         dseg=temp(
             bids(
@@ -93,7 +98,6 @@ rule atropos_seg:
                 stain="{stain}",
                 level="{level}",
                 desc="dsAtropos",
-                k="{k}",
                 suffix="dseg.nii",
                 **inputs["spim"].wildcards,
             ),
@@ -106,27 +110,19 @@ rule atropos_seg:
                     stain="{stain}",
                     level="{level}",
                     desc="Atropos",
-                    k="{k}",
                     suffix="posteriors",
                     **inputs["spim"].wildcards,
                 )
             ),
         ),
     conda:
         "../envs/ants.yaml"
-    shadow:
-        "minimal"
     threads: 16
     resources:
         mem_mb=32000,
         runtime=45,
-    shell:
-        "mkdir -p {output.posteriors_dir} && "
-        "ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS={threads} "
-        "Atropos -v -d 3 --initialization KMeans[{wildcards.k}] "
-        " --intensity-image {input.downsampled} "
-        " --output [{output.dseg},{output.posteriors_dir}/class-%02d.nii] "
-        " --mask-image {input.mask} --mrf [{params.mrf_smoothing},{params.mrf_radius}]"
+    script:
+        "../scripts/atropos_seg.py"
 
 
 rule post_atropos:
@@ -148,7 +144,6 @@ rule post_atropos:
                 stain="{stain}",
                 level="{level}",
                 desc="Atropos",
-                k="{k}",
                 suffix="dseg.nii",
                 **inputs["spim"].wildcards,
             ),
@@ -273,7 +268,6 @@ rule create_mask_from_gmm_and_prior:
             stain="{stain}",
             level="{level}",
             desc="Atropos",
-            k=config["masking"]["gmm_k"],
             suffix="dseg.nii",
             **inputs["spim"].wildcards,
         ),
@@ -285,8 +279,6 @@ rule create_mask_from_gmm_and_prior:
             suffix="mask.nii.gz",
             **inputs["spim"].wildcards,
         ),
-    params:
-        k=config["masking"]["gmm_k"],
     output:
         mask=bids(
             root=root,
@@ -313,7 +305,6 @@ rule create_mask_from_gmm:
             stain="{stain}",
             level="{level}",
             desc="Atropos",
-            k=config["masking"]["gmm_k"],
             suffix="dseg.nii",
             **inputs["spim"].wildcards,
         ),

spimquant/workflow/scripts/atropos_seg.py

@@ -0,0 +1,44 @@
+"""
+Run Atropos segmentation with automatic k-retry logic.
+
+Tries GMM segmentation with init_k classes, and decrements k if Atropos fails,
+down to min_k. This handles cases where the image doesn't have enough distinct
+intensity classes to support the initial k value.
+"""
+
+import os
+import shutil
+import subprocess
+
+init_k = snakemake.params.init_k
+min_k = snakemake.params.min_k
+threads = snakemake.threads
+mrf_smoothing = snakemake.params.mrf_smoothing
+mrf_radius = snakemake.params.mrf_radius
+
+downsampled = snakemake.input.downsampled
+mask = snakemake.input.mask
+dseg = snakemake.output.dseg
+posteriors_dir = snakemake.output.posteriors_dir
+
+os.makedirs(posteriors_dir, exist_ok=True)
+
+for k in range(init_k, min_k - 1, -1):
+    cmd = (
+        f"ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS={threads} "
+        f"Atropos -v -d 3 --initialization KMeans[{k}] "
+        f" --intensity-image {downsampled} "
+        f" --output [{dseg},{posteriors_dir}/class-%02d.nii] "
+        f" --mask-image {mask} --mrf [{mrf_smoothing},{mrf_radius}]"
+    )
+    result = subprocess.run(cmd, shell=True)
+    if result.returncode == 0:
+        break
+    elif k == min_k:
+        raise subprocess.CalledProcessError(result.returncode, cmd)
+    else:
+        # Clean up and retry with lower k
+        if os.path.exists(dseg):
+            os.remove(dseg)
+        shutil.rmtree(posteriors_dir, ignore_errors=True)
+        os.makedirs(posteriors_dir, exist_ok=True)

spimquant/workflow/scripts/create_mask_from_gmm_and_prior.py

@@ -13,7 +13,10 @@
 
 out_mask = np.zeros(template_mask_vol.shape)
 
-for i in range(1, snakemake.params.k + 1):
+# Detect the number of tissue classes from the segmentation labels
+tissue_labels = sorted(np.unique(tissue_vol[tissue_vol > 0]).astype(int))
+
+for i in tissue_labels:
 
     # if more voxels in foreground than in background, we assign it to the mask
     nvox_fg = (fg_tissue == i).sum()