Add batch histogram computation, aggregation, Otsu threshold, and QC rules

Agent-Logs-Url: https://github.com/khanlab/SPIMquant/sessions/1911c21b-2dbc-4019-96d4-4c15c24c890a

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

Author: Copilot

spimquant/workflow/Snakefile

@@ -252,6 +252,32 @@ active ``seg_method`` contains at least one ``otsu+k{}i{}`` entry.
         else [],
 
 
+rule all_batch_otsu_hist_qc:
+    """Target rule to generate batch-wide Otsu threshold sweep QC reports.
+
+Produces one HTML report per stain/method combination (across all subjects)
+for every ``otsu+k{}i{}`` segmentation method.  Each report shows the
+aggregated histogram with Multi-Otsu threshold positions and one mid-volume
+crop per subject at each threshold, enabling visual selection of a threshold
+that generalises across the entire acquisition batch.  Only meaningful when
+the active ``seg_method`` contains at least one ``otsu+k{}i{}`` entry.
+"""
+    input:
+        expand(
+            bids(
+                root=root,
+                datatype="group",
+                stain="{stain}",
+                desc="{desc}",
+                suffix="batchotsuthreshqc.html",
+            ),
+            stain=stains_for_seg,
+            desc=otsu_seg_methods,
+        )
+        if (do_seg and len(otsu_seg_methods) > 0)
+        else [],
+
+
 rule all_register:
     input:
         inputs["spim"].expand(

spimquant/workflow/rules/qc.smk

@@ -459,6 +459,62 @@ Only applicable when ``seg_method`` uses the ``otsu+k{}i{}`` pattern.
         "../scripts/qc_otsu_threshold_sweep.py"
 
 
+rule qc_batch_otsu_threshold_sweep:
+    """Batch threshold sweep QC HTML report for multiotsu segmentation.
+
+Uses the batch-wide Multi-Otsu thresholds (computed from the aggregated
+histogram across all subjects in the dataset) to sweep over a range of
+threshold values and display one mid-volume 2D crop per subject at each
+threshold.  Subjects are shown as columns so that the user can immediately
+see whether a given threshold generalises across the entire acquisition batch.
+
+The aggregated histogram PNG is embedded at the top of the report.
+
+Only applicable when ``seg_method`` uses the ``otsu+k{}i{}`` pattern.
+"""
+    input:
+        corrected_zarrs=inputs["spim"].expand(
+            bids(
+                root=work,
+                datatype="seg",
+                stain="{stain}",
+                level=str(config["segmentation_level"]),
+                desc="corrected{method}".format(method=config["correction_method"]),
+                suffix="SPIM.ome.zarr",
+                **inputs["spim"].wildcards,
+            ),
+            allow_missing=True,
+        ),
+        thresholds_png=bids(
+            root=root,
+            datatype="group",
+            stain="{stain}",
+            level=str(config["segmentation_level"]),
+            desc="{desc}",
+            suffix="thresholds.png",
+        ),
+    output:
+        html=bids(
+            root=root,
+            datatype="group",
+            stain="{stain}",
+            desc="{desc}",
+            suffix="batchotsuthreshqc.html",
+        ),
+    threads: 4
+    resources:
+        mem_mb=32000,
+        runtime=60,
+    params:
+        n_thresholds=10,
+        patch_size=300,
+        level=config["segmentation_level"],
+        hist_percentile_range=[float(x) for x in config["seg_hist_percentile_range"]],
+        zarrnii_kwargs={"orientation": config["orientation"]},
+    script:
+        "../scripts/qc_batch_otsu_threshold_sweep.py"
+
+
 rule qc_roi_summary:
     """Per-ROI summary QC: top-region bar plots for a single subject.
 

spimquant/workflow/rules/segmentation.smk

@@ -1,7 +1,7 @@
 """
 Segmentation and quantification workflow for SPIMquant.
 
-This module performs intensity correction, segmentation, and quantitative analysis of 
+This module performs intensity correction, segmentation, and quantitative analysis of
 pathology markers from SPIM microscopy data. It handles multi-channel data with different
 stains (e.g., beta-amyloid, alpha-synuclein, Iba1) and produces per-region statistics.
 
@@ -27,9 +27,6 @@ rule gaussian_biasfield:
     """simple bias field correction with gaussian"""
     input:
         spim=inputs["spim"].path,
-    params:
-        proc_level=5,
-        zarrnii_kwargs={"orientation": config["orientation"]},
     output:
         corrected=temp(
             directory(
@@ -64,6 +61,9 @@ rule gaussian_biasfield:
         mem_mb=256000,
         disk_mb=2097152,
         runtime=15,
+    params:
+        proc_level=5,
+        zarrnii_kwargs={"orientation": config["orientation"]},
     script:
         "../scripts/gaussian_biasfield.py"
 
@@ -72,11 +72,6 @@ rule n4_biasfield:
     """N4 bias field correction with antspyx"""
     input:
         spim=inputs["spim"].path,
-    params:
-        proc_level=5,
-        zarrnii_kwargs={"orientation": config["orientation"]},
-        shrink_factor=16 if config["sloppy"] else 1,
-        target_chunk_size=512,  #this sets the chunk size for this and downstream masks
     output:
         corrected=temp(
             directory(
@@ -96,18 +91,23 @@ rule n4_biasfield:
     resources:
         mem_mb=500000 if config["dask_scheduler"] == "distributed" else 250000,
         runtime=180,
+    params:
+        proc_level=5,
+        zarrnii_kwargs={"orientation": config["orientation"]},
+        shrink_factor=16 if config["sloppy"] else 1,
+        target_chunk_size=512,  #this sets the chunk size for this and downstream masks
     script:
         "../scripts/n4_biasfield.py"
 
 
 rule multiotsu:
     """Perform multi-Otsu thresholding for segmentation.
-    
-    Applies multi-level Otsu thresholding to identify intensity classes in the
-    corrected image. The k parameter determines number of classes, and the i parameter
-    selects which threshold index to use for creating the binary mask. Outputs a
-    histogram visualization of the threshold selection.
-    """
+
+Applies multi-level Otsu thresholding to identify intensity classes in the
+corrected image. The k parameter determines number of classes, and the i parameter
+selects which threshold index to use for creating the binary mask. Outputs a
+histogram visualization of the threshold selection.
+"""
     input:
         corrected=bids(
             root=work,
@@ -118,12 +118,6 @@ rule multiotsu:
             suffix="SPIM.ome.zarr",
             **inputs["spim"].wildcards,
         ),
-    params:
-        hist_bin_width=float(config["seg_hist_bin_width"]),
-        hist_percentile_range=[float(x) for x in config["seg_hist_percentile_range"]],
-        otsu_k=lambda wildcards: int(wildcards.k),
-        otsu_threshold_index=lambda wildcards: int(wildcards.i),
-        zarrnii_kwargs={"orientation": config["orientation"]},
     output:
         mask=bids(
             root=root,
@@ -148,16 +142,23 @@ rule multiotsu:
         mem_mb=500000 if config["dask_scheduler"] == "distributed" else 250000,
         disk_mb=2097152,
         runtime=180,
+    params:
+        hist_bin_width=float(config["seg_hist_bin_width"]),
+        hist_percentile_range=[float(x) for x in config["seg_hist_percentile_range"]],
+        otsu_k=lambda wildcards: int(wildcards.k),
+        otsu_threshold_index=lambda wildcards: int(wildcards.i),
+        zarrnii_kwargs={"orientation": config["orientation"]},
     script:
         "../scripts/multiotsu.py"
 
 
 rule threshold:
     """Apply simple intensity threshold for segmentation.
-    
-    Creates binary mask by thresholding the corrected image at a specified intensity value.
-    Simpler alternative to multi-Otsu for cases where the threshold is known a priori.
-    """
+
+Creates binary mask by thresholding the corrected image at a specified intensity value.
+Simpler alternative to multi-Otsu for cases where the threshold is known a priori.
+
+"""
     input:
         corrected=bids(
             root=work,
@@ -168,9 +169,6 @@ rule threshold:
             suffix="SPIM.ome.zarr",
             **inputs["spim"].wildcards,
         ),
-    params:
-        threshold=lambda wildcards: int(wildcards.threshold),
-        zarrnii_kwargs={"orientation": config["orientation"]},
     output:
         mask=bids(
             root=root,
@@ -185,17 +183,136 @@ rule threshold:
     resources:
         mem_mb=500000 if config["dask_scheduler"] == "distributed" else 250000,
         runtime=180,
+    params:
+        threshold=lambda wildcards: int(wildcards.threshold),
+        zarrnii_kwargs={"orientation": config["orientation"]},
     script:
         "../scripts/threshold.py"
 
 
+rule compute_spim_histogram:
+    """Compute an intensity histogram from a raw SPIM channel at a given level.
+
+Reads the OME-Zarr at the specified pyramid *level* using Dask and accumulates
+a histogram over a fixed intensity range with a fixed number of bins.  The bin
+edges and counts are stored as a compressed NumPy archive (``.npz``) for later
+aggregation and batch-wide threshold optimisation.
+"""
+    input:
+        spim=inputs["spim"].path,
+    output:
+        histogram=bids(
+            root=root,
+            datatype="seg",
+            stain="{stain}",
+            level="{level}",
+            suffix="histogram.npz",
+            **inputs["spim"].wildcards,
+        ),
+    threads: 32
+    resources:
+        mem_mb=32000,
+        runtime=30,
+    params:
+        level="{level}",
+        hist_bins=config.get("spim_hist_bins", 500),
+        hist_range=config.get("spim_hist_range", [0, 65535]),
+        zarrnii_kwargs={"orientation": config["orientation"]},
+    script:
+        "../scripts/compute_spim_histogram.py"
+
+
+rule aggregate_histograms:
+    """Aggregate per-subject histograms into a single batch histogram.
+
+Collects histogram ``.npz`` files from all subjects (for a given stain and
+pyramid level) and produces a single combined histogram by summing counts.
+All input histograms must have been computed with the same ``hist_range`` and
+``hist_bins`` parameters.
+"""
+    input:
+        histograms=inputs["spim"].expand(
+            bids(
+                root=root,
+                datatype="seg",
+                stain="{stain}",
+                level="{level}",
+                suffix="histogram.npz",
+                **inputs["spim"].wildcards,
+            ),
+            allow_missing=True,
+        ),
+    output:
+        histogram=bids(
+            root=root,
+            datatype="group",
+            stain="{stain}",
+            level="{level}",
+            suffix="histogram.npz",
+        ),
+    threads: 1
+    resources:
+        mem_mb=4000,
+        runtime=10,
+    script:
+        "../scripts/aggregate_histograms.py"
+
+
+rule multiotsu_from_batch_histogram:
+    """Compute Multi-Otsu thresholds from the batch-aggregated histogram.
+
+Reads the aggregated histogram, optionally clips to a percentile-based
+intensity range, and applies multi-level Otsu thresholding.  This yields a
+single set of thresholds that is consistent across all subjects in the batch
+(i.e. acquisitions with the same imaging parameters).
+
+Outputs a TSV of threshold index → value and a histogram PNG with the
+threshold positions marked.
+"""
+    input:
+        histogram=bids(
+            root=root,
+            datatype="group",
+            stain="{stain}",
+            level="{level}",
+            suffix="histogram.npz",
+        ),
+    output:
+        thresholds_png=bids(
+            root=root,
+            datatype="group",
+            stain="{stain}",
+            level="{level}",
+            desc="otsu+k{k,[0-9]+}i{i,[0-9]+}",
+            suffix="thresholds.png",
+        ),
+        thresholds_tsv=bids(
+            root=root,
+            datatype="group",
+            stain="{stain}",
+            level="{level}",
+            desc="otsu+k{k,[0-9]+}i{i,[0-9]+}",
+            suffix="thresholds.tsv",
+        ),
+    threads: 1
+    resources:
+        mem_mb=4000,
+        runtime=10,
+    params:
+        hist_percentile_range=[float(x) for x in config["seg_hist_percentile_range"]],
+        otsu_k=lambda wildcards: int(wildcards.k),
+        otsu_threshold_index=lambda wildcards: int(wildcards.i),
+    script:
+        "../scripts/multiotsu_from_batch_histogram.py"
+
+
 rule clean_segmentation:
     """Clean segmentation mask by removing edge artifacts and small objects.
-    
-    Performs connected component analysis to identify and exclude objects that
-    extend too close to the image boundaries (likely artifacts). Creates both
-    a cleaned mask and an exclusion mask showing what was removed.
-    """
+
+Performs connected component analysis to identify and exclude objects that
+extend too close to the image boundaries (likely artifacts). Creates both
+a cleaned mask and an exclusion mask showing what was removed.
+"""
     input:
         mask=bids(
             root=root,
@@ -206,10 +323,6 @@ rule clean_segmentation:
             suffix="mask.ozx",
             **inputs["spim"].wildcards,
         ),
-    params:
-        max_extent=0.15,
-        proc_level=2,  #level at which to calculate conncomp
-        zarrnii_kwargs={"orientation": config["orientation"]},
     output:
         exclude_mask=bids(
             root=root,
@@ -234,5 +347,9 @@ rule clean_segmentation:
         mem_mb=256000,
         disk_mb=2097152,
         runtime=30,
+    params:
+        max_extent=0.15,
+        proc_level=2,  #level at which to calculate conncomp
+        zarrnii_kwargs={"orientation": config["orientation"]},
     script:
         "../scripts/clean_segmentation.py"