feat: improve zoom montage QC PNGs - better mask color, no-mask output, n4 bg option

Agent-Logs-Url: https://github.com/khanlab/SPIMquant/sessions/d973790d-eb1b-4730-b51a-704815237200

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

Author: Copilot

spimquant/config/snakebids.yml

@@ -184,6 +184,11 @@ parse_args:
       - distributed
     default: distributed
 
+  --qc_roi_zoom_bg_n4:
+    help: "Use n4 bias-field corrected OME-Zarr as background image in the ROI zoom montage QC figures instead of the raw SPIM (only effective when --correction_method n4) (default: %(default)s)"
+    action: store_true
+    default: False
+
   --sloppy:
     help: "Use low-quality parameters for speed (USE FOR TESTING ONLY)"
     action: store_true

spimquant/workflow/Snakefile

@@ -682,6 +682,25 @@ rule all_qc:
         )
         if do_seg
         else [],
+        # Segmentation ROI zoom montage - no mask overlay
+        inputs["spim"].expand(
+            bids(
+                root=root,
+                datatype="qc",
+                seg="{seg}",
+                from_="{template}",
+                stain="{stain}",
+                desc="{desc}nomask",
+                suffix="roimontage.png",
+                **inputs["spim"].wildcards,
+            ),
+            seg=atlas_segs,
+            template=config["template"],
+            stain=stains_for_seg,
+            desc=config["seg_method"],
+        )
+        if do_seg
+        else [],
         # Vessel overview figures
         inputs["spim"].expand(
             bids(
@@ -716,6 +735,25 @@ rule all_qc:
         )
         if do_vessels
         else [],
+        # Vessel ROI zoom montage - no mask overlay
+        inputs["spim"].expand(
+            bids(
+                root=root,
+                datatype="qc",
+                seg="{seg}",
+                from_="{template}",
+                stain="{stain}",
+                desc="{desc}nomask",
+                suffix="vesselroimontage.png",
+                **inputs["spim"].wildcards,
+            ),
+            seg=atlas_segs,
+            template=config["template"],
+            stain=stain_for_vessels,
+            desc=config["vessel_seg_method"],
+        )
+        if do_vessels
+        else [],
         # Z-profile QC (per stain, per seg method)
         inputs["spim"].expand(
             bids(

spimquant/workflow/rules/qc.smk

@@ -27,6 +27,13 @@ Rules 2 and the ROI zoom are also generated for vessel segmentations.
 All outputs are written to the ``qc`` datatype directory for each subject.
 """
 
+# Whether to use the n4-corrected OME-Zarr as the background in zoom montage QC.
+# Enabled via the --qc_roi_zoom_bg_n4 CLI option (only meaningful when
+# correction_method=="n4").
+_use_n4_bg = config.get("qc_roi_zoom_bg_n4", False) and (
+    config.get("correction_method") == "n4"
+)
+
 
 rule qc_intensity_histogram:
     """Per-channel intensity histogram QC.
@@ -147,6 +154,9 @@ region's bounding box (in subject space) and displays the best axial slice
 with the field-fraction overlay.  Aspect ratio is corrected from NIfTI
 voxel dimensions.  Provides detail-level visualisation of segmentation
 quality within individual brain regions.
+
+Two PNGs are produced: one with the mask overlay (``roimontage.png``) and
+one without (``desc-{desc}nomask_roimontage.png``).
 """
     input:
         spim=inputs["spim"].path,
@@ -174,6 +184,19 @@ quality within individual brain regions.
             seg="{seg}",
             suffix="dseg.tsv",
         ),
+        **{
+            "spim_n4": bids(
+                root=work,
+                datatype="seg",
+                stain="{stain}",
+                level=str(config["segmentation_level"]),
+                desc="correctedn4",
+                suffix="SPIM.ome.zarr",
+                **inputs["spim"].wildcards,
+            )
+        }
+        if _use_n4_bg
+        else {},
     output:
         png=bids(
             root=root,
@@ -185,6 +208,16 @@ quality within individual brain regions.
             suffix="roimontage.png",
             **inputs["spim"].wildcards,
         ),
+        png_nomask=bids(
+            root=root,
+            datatype="qc",
+            seg="{seg}",
+            from_="{template}",
+            stain="{stain}",
+            desc="{desc}nomask",
+            suffix="roimontage.png",
+            **inputs["spim"].wildcards,
+        ),
     threads: 4
     resources:
         mem_mb=32000,
@@ -194,6 +227,7 @@ quality within individual brain regions.
         n_cols=lambda wildcards: 5 if wildcards.seg == "coarse" else 10,
         patch_size=lambda wildcards: 2000 if wildcards.seg == "coarse" else 500,
         level=config["segmentation_level"],
+        use_n4_bg=_use_n4_bg,
     script:
         "../scripts/qc_segmentation_roi_zoom.py"
 
@@ -204,6 +238,9 @@ rule qc_vessels_roi_zoom:
 Identical to ``qc_segmentation_roi_zoom`` but applied to the vessel
 binary mask.  Uses ZarrNii to load full-resolution data and
 ZarrNiiAtlas for atlas-based ROI cropping.
+
+Two PNGs are produced: one with the mask overlay (``vesselroimontage.png``)
+and one without (``desc-{desc}nomask_vesselroimontage.png``).
 """
     input:
         spim=inputs["spim"].path,
@@ -231,6 +268,19 @@ ZarrNiiAtlas for atlas-based ROI cropping.
             seg="{seg}",
             suffix="dseg.tsv",
         ),
+        **{
+            "spim_n4": bids(
+                root=work,
+                datatype="seg",
+                stain="{stain}",
+                level=str(config["segmentation_level"]),
+                desc="correctedn4",
+                suffix="SPIM.ome.zarr",
+                **inputs["spim"].wildcards,
+            )
+        }
+        if _use_n4_bg
+        else {},
     output:
         png=bids(
             root=root,
@@ -242,6 +292,16 @@ ZarrNiiAtlas for atlas-based ROI cropping.
             suffix="vesselroimontage.png",
             **inputs["spim"].wildcards,
         ),
+        png_nomask=bids(
+            root=root,
+            datatype="qc",
+            seg="{seg}",
+            from_="{template}",
+            stain="{stain}",
+            desc="{desc}nomask",
+            suffix="vesselroimontage.png",
+            **inputs["spim"].wildcards,
+        ),
     threads: 4
     resources:
         mem_mb=32000,
@@ -251,6 +311,7 @@ ZarrNiiAtlas for atlas-based ROI cropping.
         n_cols=lambda wildcards: 5 if wildcards.seg == "coarse" else 10,
         patch_size=lambda wildcards: 2000 if wildcards.seg == "coarse" else 500,
         level=config["segmentation_level"],
+        use_n4_bg=_use_n4_bg,
     script:
         "../scripts/qc_segmentation_roi_zoom.py"
 

spimquant/workflow/scripts/qc_segmentation_roi_zoom.py

@@ -48,9 +48,11 @@ def main():
     subject = snakemake.wildcards.subject
     max_rois = snakemake.params.max_rois
     n_cols = snakemake.params.n_cols
+    use_n4_bg = snakemake.params.get("use_n4_bg", False)
 
+    spim_bg_path = snakemake.input.spim_n4 if use_n4_bg else snakemake.input.spim
     spim_img = ZarrNii.from_ome_zarr(
-        snakemake.input.spim,
+        spim_bg_path,
         level=snakemake.params.level,
         downsample_near_isotropic=True,
         channel_labels=[snakemake.wildcards.stain],
@@ -66,7 +68,7 @@ def main():
     aspect_axial = 1
 
     spim_img_ds = ZarrNii.from_ome_zarr(
-        snakemake.input.spim,
+        spim_bg_path,
         level=(int(snakemake.params.level) + 5),
         downsample_near_isotropic=True,
         channel_labels=[snakemake.wildcards.stain],
@@ -103,7 +105,22 @@ def main():
             color="gray",
         )
         ax.axis("off")
-        plt.savefig(snakemake.output.png, dpi=120, bbox_inches="tight")
+        plt.savefig(snakemake.output.png, dpi=150, bbox_inches="tight")
+        plt.close()
+        # Also save the no-mask version (same empty figure)
+        fig, ax = plt.subplots(figsize=(18, 12))
+        ax.text(
+            0.5,
+            0.5,
+            "No atlas ROIs found in subject",
+            ha="center",
+            va="center",
+            transform=ax.transAxes,
+            fontsize=12,
+            color="gray",
+        )
+        ax.axis("off")
+        plt.savefig(snakemake.output.png_nomask, dpi=150, bbox_inches="tight")
         plt.close()
         return
 
@@ -126,6 +143,8 @@ def main():
     elif n_cols == 1:
         axes = axes[:, np.newaxis]
 
+    # Cache crops for both masked and no-mask figures
+    cached_slices = []
     for i, row in enumerate(roi_rows):
         ax_row = i // n_cols
         ax_col = i % n_cols
@@ -149,11 +168,12 @@ def main():
         spim_sl = spim_crop.data[0, :, :].squeeze().compute()
         spim_sl = _apply_fixed_percentile_norm(spim_sl, glob_lo, glob_hi)
         mask_sl = mask_crop.data[0, :, :].squeeze().compute()
+        cached_slices.append((label_name, spim_sl, mask_sl))
 
         ax.imshow(spim_sl, cmap="gray")
         mask_masked = np.ma.masked_where(mask_sl < 100, mask_sl)
         ax.imshow(
-            mask_masked, cmap="spring", alpha=0.6, vmin=0, vmax=100, aspect=aspect_axial
+            mask_masked, cmap="Reds", alpha=0.6, vmin=0, vmax=100, aspect=aspect_axial
         )
         ax.set_title(label_name, fontsize=7, pad=2)
         ax.set_xticks([])
@@ -167,6 +187,41 @@ def main():
     plt.close()
     print(f"Saved ROI zoom montage to {snakemake.output.png}")
 
+    # --- No-mask version (reuses cached slices) ---
+    n_rows_nm = int(np.ceil(n_rois / n_cols))
+    fig_nm, axes_nm = plt.subplots(
+        n_rows_nm,
+        n_cols,
+        figsize=(n_cols * 3, n_rows_nm * 3),
+        constrained_layout=True,
+    )
+    fig_nm.suptitle(
+        f"ROI Zoom Montage QC (no mask overlay)\n"
+        f"Subject: {subject}  |  Stain: {stain}  |  Method: {desc}",
+        fontsize=11,
+        fontweight="bold",
+    )
+    if n_rows_nm == 1 and n_cols == 1:
+        axes_nm = np.array([[axes_nm]])
+    elif n_rows_nm == 1:
+        axes_nm = axes_nm[np.newaxis, :]
+    elif n_cols == 1:
+        axes_nm = axes_nm[:, np.newaxis]
+
+    for i, (label_name, spim_sl, _) in enumerate(cached_slices):
+        ax_nm = axes_nm[i // n_cols, i % n_cols]
+        ax_nm.imshow(spim_sl, cmap="gray")
+        ax_nm.set_title(label_name, fontsize=7, pad=2)
+        ax_nm.set_xticks([])
+        ax_nm.set_yticks([])
+
+    for i in range(n_rois, n_rows_nm * n_cols):
+        axes_nm[i // n_cols, i % n_cols].axis("off")
+
+    plt.savefig(snakemake.output.png_nomask, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"Saved no-mask ROI zoom montage to {snakemake.output.png_nomask}")
+
 
 if __name__ == "__main__":
     main()