Ims workflow fixes

pyproject.toml

@@ -51,9 +51,9 @@ pyarrow = ">=18.0.0,<19"
 
 [tool.pixi.pypi-dependencies]
 spimquant = { path = ".", editable = true }
-zarrnii = ">=0.19.7,<0.20.0"
-#zarrnii = { git = "https://github.com/khanlab/zarrnii", rev = "main",  editable = true }
-#zarrnii = { path = "../zarrnii", editable = true }
+zarrnii = ">=0.21.0,<0.22.0"
+#zarrnii = { git = "https://github.com/khanlab/zarrnii", rev = "resize-local-mean",  editable = true }
+#zarrnii = { path = "/nfs/khan/trainees/akhan488/apps/zarrnii", editable = true }
 vesselfm = { git = "https://github.com/khanlab/vesselfm", rev = "504baab" }
 h5py = ">=3.14.0"
 simpleitk = ">=2.4.0"

spimquant/config/snakebids.yml

@@ -213,6 +213,11 @@ parse_args:
       - distributed
     default: distributed
 
+  --qc_instance_seg:
+    help: "Generate instance-level QC outputs, including animated gifs (default: %(default)s)"
+    action: store_true
+    default: False
+
   --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

spimquant/workflow/Snakefile

@@ -918,7 +918,7 @@ rule all_qc:
             stain=stains_for_seg,
             desc=config["seg_method"],
         )
-        if do_seg
+        if do_seg and config["qc_instance_seg"]
         else [],
         # Colocalization instance GIFs (per atlas seg, per seg method)
         inputs["spim"].expand(
@@ -946,7 +946,7 @@ rule all_qc:
             template=config["template"],
             desc=config["seg_method"],
         )
-        if do_coloc
+        if do_coloc and config["qc_instance_seg"]
         else [],
 
 

spimquant/workflow/scripts/n4_biasfield.py

@@ -4,7 +4,9 @@
     from zarrnii import ZarrNii
     from zarrnii.plugins import N4BiasFieldCorrection
 
-    with get_dask_client(snakemake.config["dask_scheduler"], snakemake.threads):
+    with get_dask_client(
+        snakemake.config["dask_scheduler"], snakemake.threads, threads_per_worker=16
+    ):
 
         hires_level = int(snakemake.wildcards.level)
         proc_level = int(snakemake.params.proc_level)
@@ -17,19 +19,24 @@
             channel_labels=[snakemake.wildcards.stain],
             level=hires_level,
             downsample_near_isotropic=True,
+            chunks=(256, 256, 256),
+            **snakemake.params.zarrnii_kwargs,
+        )
+        znimg_lowres = ZarrNii.from_file(
+            snakemake.input.spim,
+            channel_labels=[snakemake.wildcards.stain],
+            level=proc_level,
+            downsample_near_isotropic=True,
             **snakemake.params.zarrnii_kwargs,
         )
 
         print("compute bias field correction")
 
-        adjusted_chunk = int(
-            snakemake.params.target_chunk_size / (2**adjusted_downsample_factor)
-        )
-
         # Apply bias field correction
         znimg_corrected = znimg.apply_scaled_processing(
             N4BiasFieldCorrection(shrink_factor=snakemake.params.shrink_factor),
-            downsample_factor=adjusted_downsample_factor,
+            lowres_znimg=znimg_lowres,
+            method="map_blocks",
         )
 
         # write to ome_zarr

spimquant/workflow/scripts/ome_zarr_to_nii.py

@@ -3,12 +3,13 @@
 from zarrnii import ZarrNii
 
 if __name__ == "__main__":
-    with get_dask_client(snakemake.config["dask_scheduler"], snakemake.threads):
+    with get_dask_client("threads", snakemake.threads):
         znimg = ZarrNii.from_file(
             snakemake.input.spim,
             level=int(snakemake.wildcards.level),
             channel_labels=[snakemake.wildcards.stain],
             downsample_near_isotropic=True,
+            chunks=None,
             **snakemake.params.zarrnii_kwargs,
         )
 

spimquant/workflow/scripts/qc_instance_gif.py

@@ -70,6 +70,7 @@ def _load_channel_images(spim_path, channels, level):
             spim_path,
             level=level,
             channel_labels=[ch],
+            **snakemake.params.zarrnii_kwargs,
         )
     return imgs
 
@@ -94,6 +95,7 @@ def _global_norms(imgs, channels):
                     snakemake.input.spim,
                     level=(base_level + extra),
                     channel_labels=[ch],
+                    **snakemake.params.zarrnii_kwargs,
                 )
                 arr = coarse.data.compute()
                 break

spimquant/workflow/scripts/qc_segmentation_roi_zoom.py

@@ -159,6 +159,7 @@ def main():
         # get cropped images for this label
         bbox_min, bbox_max = atlas.get_region_bounding_box(region_ids=label_id)
         center_coord = tuple((x + y) / 2 for x, y in zip(bbox_min, bbox_max))
+
         spim_crop = spim_img.crop_centered(
             center_coord,
             patch_size=(snakemake.params.patch_size, snakemake.params.patch_size, 1),
@@ -170,6 +171,7 @@ 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))
 

spimquant/workflow/scripts/vesselfm.py

@@ -10,6 +10,7 @@
         level=int(snakemake.wildcards.level),
         channel_labels=[snakemake.wildcards.stain],
         downsample_near_isotropic=True,
+        chunks=(256, 256, 256),
         **snakemake.params.zarrnii_kwargs,
     )
     znimg_mask = znimg.segment(VesselFMPlugin, **snakemake.params.vesselfm_kwargs)