Support datasets mixing .ims and .ome.zarr image formats

spimquant/config/snakebids.yml

@@ -34,6 +34,7 @@ pybids_inputs:
       - sample
       - acquisition
       - staining
+      - extension
   mri:
     filters:
       suffix: 'T2w'

spimquant/workflow/Snakefile

@@ -4,7 +4,14 @@ from itertools import combinations as _combinations, product as _product
 import pandas as _pd
 from zarrnii import ZarrNii
 from snakemake.utils import format
-from snakebids import bids, generate_inputs, get_wildcard_constraints, set_bids_spec
+from snakebids import (
+    BidsComponent,
+    BidsDataset,
+    bids,
+    generate_inputs,
+    get_wildcard_constraints,
+    set_bids_spec,
+)
 import snakebids
 
 # update v15 spec to add template entity
@@ -58,6 +65,38 @@ inputs = generate_inputs(
 )
 
 
+# The spim component is generated with `extension` as a wildcard so that one
+# path template covers datasets mixing image formats (e.g. .ims files alongside
+# .ome.zarr stores). Rules resolve each scan's concrete image path through
+# spim_input(); the component exposed to the rest of the workflow has the
+# extension stripped so it cannot leak into derived output templates.
+spim_raw = inputs["spim"]
+
+_spim_key_wildcards = [w for w in spim_raw.zip_lists if w != "extension"]
+_spim_path_lookup = {
+    tuple(spim_raw.zip_lists[w][i] for w in _spim_key_wildcards): path
+    for i, path in enumerate(spim_raw.expand())
+}
+
+
+def spim_input(wildcards):
+    """Concrete raw image path (whatever its extension) for a job's wildcards."""
+    return _spim_path_lookup[tuple(getattr(wildcards, w) for w in _spim_key_wildcards)]
+
+
+inputs = BidsDataset(
+    {
+        **inputs,
+        "spim": BidsComponent(
+            name=spim_raw.name,
+            path=spim_raw.path.replace("{extension}", ""),
+            zip_lists={w: spim_raw.zip_lists[w] for w in _spim_key_wildcards},
+        ),
+    },
+    layout=inputs.layout,
+)
+
+
 if "mri" in inputs:
     suffixes = inputs["mri"].entities["suffix"]
     if len(suffixes) == 1:
@@ -81,7 +120,7 @@ include: "rules/common.smk"
 # prepopulate json sidecar metadata for every spim input
 spim_json_overrides = {
     str(spim_path): get_spim_json_overrides(spim_path)
-    for spim_path in inputs["spim"].expand()
+    for spim_path in spim_raw.expand()
 }
 zarrnii_cli_kwargs = {}
 

spimquant/workflow/rules/common.smk

@@ -89,7 +89,7 @@ def get_stains_all_subjects(ignore_stains=None):
     ignore_set = set(ignore_stains) if ignore_stains else set()
 
     stain_sets = []
-    for zarr in inputs["spim"].expand():
+    for zarr in spim_raw.expand():
         channels = set(get_spim_channels(zarr))
         # Remove ignored stains
         channels = channels - ignore_set

spimquant/workflow/rules/counts.smk

@@ -1,7 +1,7 @@
 rule counts_per_voxel:
     """Calculate counts per voxel based on points"""
     input:
-        ref_spim=inputs["spim"].path,
+        ref_spim=spim_input,
         regionprops_parquet=bids(
             root=root,
             datatype="tabular",

spimquant/workflow/rules/import.smk

@@ -30,7 +30,7 @@ rule get_downsampled_nii:
     based on configuration.
     """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
     params:
         zarrnii_kwargs=zarrnii_in_kwargs,
     output:

spimquant/workflow/rules/patches.smk

@@ -31,7 +31,7 @@ rule create_spim_patches:
     but level in the dseg input is the (downsampled) registration_level.
     """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         dseg=bids(
             root=root,
             datatype="parc",
@@ -192,7 +192,7 @@ rule create_imaris_crops:
     to_imaris() function. Level defaults to 0 for high-resolution output.
     """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         dseg=bids(
             root=root,
             datatype="parc",

spimquant/workflow/rules/plaques.smk

@@ -28,7 +28,7 @@ rule run_lantern_plaques:
     changed downstream without re-running inference.
     """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         models=expand(
             "resources/models/lantern-ki3-abeta/fold{fold}/seg_model.pt",
             fold=range(config["plaque_n_folds"]),
@@ -108,7 +108,7 @@ rule binarize_lantern_plaques:
             suffix="probseg.{ext}",
             **inputs["spim"].wildcards,
         ),
-        spim=inputs["spim"].path,
+        spim=spim_input,
     params:
         zarrnii_kwargs=zarrnii_in_kwargs,
         n_folds=config["plaque_n_folds"],

spimquant/workflow/rules/qc.smk

@@ -51,7 +51,7 @@ cumulative distribution, and a summary-statistics panel including the
 saturation/clip fraction (percentage of voxels at the maximum bin).
 """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
     output:
         png=bids(
             root=root,
@@ -83,7 +83,7 @@ overlay, and a max-intensity projection column for each orientation.
 Aspect ratio is corrected using voxel spacings from ``ZarrNii.get_zooms()``.
 """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         mask=bids_oz_in(
             root=root,
             datatype="seg",
@@ -122,7 +122,7 @@ vessel binary mask.  Loads data via ZarrNii with ``downsample_near_isotropic``
 for isotropic display and physically correct aspect ratio.
 """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         mask=bids_oz_in(
             root=root,
             datatype="vessels",
@@ -181,7 +181,7 @@ one without (``desc-{desc}nomask_roimontage.png``).
             if _use_n4_bg
             else {}
         ),
-        spim=inputs["spim"].path,
+        spim=spim_input,
         mask=bids_oz_in(
             root=root,
             datatype="seg",
@@ -268,7 +268,7 @@ and one without (``desc-{desc}nomask_vesselroimontage.png``).
             if _use_n4_bg
             else {}
         ),
-        spim=inputs["spim"].path,
+        spim=spim_input,
         mask=bids_oz_in(
             root=root,
             datatype="vessels",
@@ -660,7 +660,7 @@ Inputs are the aggregated (all-stain) regionprops parquet in template space
 parcellation for atlas-label lookup.
 """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         instance_parquet=bids(
             root=root,
             datatype="tabular",
@@ -730,7 +730,7 @@ two colocalized objects (subject-space ``pos_coloc_x/y/z``), with a circle
 marker drawn at the average radius of the pair.
 """
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         instance_parquet=bids(
             root=root,
             datatype="tabular",

spimquant/workflow/rules/segmentation.smk

@@ -26,7 +26,7 @@ Outputs include:
 rule gaussian_biasfield:
     """simple bias field correction with gaussian"""
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
     output:
         corrected=temp(
             bids_oz_out(
@@ -204,7 +204,7 @@ rule n4_pre_quant_tune:
 rule n4_biasfield:
     """N4 bias field correction with antspyx"""
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         biasfield=bids(
             root=root,
             datatype="micr",

spimquant/workflow/rules/templatereg.smk

@@ -386,7 +386,7 @@ rule resample_labels_to_zarr:
         dseg=bids(root=root, template="{template}", desc="LR", suffix="dseg.nii.gz"),
         label_tsv=bids(root=root, template="{template}", desc="LR", suffix="dseg.tsv"),
         xfm_ras=rules.affine_reg.output.xfm_ras,
-        zarr_zip=inputs["spim"].path,
+        zarr_zip=spim_input,
     params:
         level_to_resample_to=0,
         max_downsampling_layers=config["ome_zarr"]["max_downsampling_layers"],
@@ -420,7 +420,7 @@ rule resample_labels_to_zarr:
 
 rule affine_zarr_to_template_nii:
     input:
-        ome_zarr=inputs["spim"].path,
+        ome_zarr=spim_input,
         xfm_ras=rules.affine_reg.output.xfm_ras,
         ref_nii=get_template_for_reg,
     params:
@@ -445,7 +445,7 @@ rule affine_zarr_to_template_nii:
 
 rule affine_zarr_to_template_ome_zarr:
     input:
-        ome_zarr=inputs["spim"].path,
+        ome_zarr=spim_input,
         xfm_ras=rules.affine_reg.output.xfm_ras,
         ref_nii=get_template_for_reg,
     params:
@@ -474,7 +474,7 @@ rule deform_zarr_to_template_nii:
         xfm_composite=rules.compose_subject_to_template_warp.output.xfm_composite,
         ref_nii=get_template_for_reg,
     params:
-        ome_zarr=inputs["spim"].path,
+        ome_zarr=spim_input,
         flo_opts={"level": 2},  #downsampling level to use (TODO: set this automatically based on ref resolution?)
         do_downsample=True,  #whether to perform further downsampling before transforming
         downsample_opts={"along_z": 4},  #could also be determined automatically 
@@ -499,7 +499,7 @@ rule deform_zarr_to_template_nii:
 
 rule deform_to_template_nii_zoomed:
     input:
-        ome_zarr=inputs["spim"].path,
+        ome_zarr=spim_input,
         xfm_composite=rules.compose_subject_to_template_warp.output.xfm_composite,
         ref_nii=get_template_for_reg,
     params:
@@ -612,7 +612,7 @@ rule deform_template_dseg_to_subject_nii:
 """ this rule needs updating - use atlas/seg wildcard and proper script
 rule deform_transform_labels_to_subj:
     input:
-        ref_ome_zarr=inputs["spim"].path,
+        ref_ome_zarr=spim_input,
         xfm_ras=rules.affine_reg.output.xfm_ras,
         invwarp_nii=rules.deform_reg.output.invwarp,
         flo_nii=bids(

spimquant/workflow/rules/vessels.smk

@@ -10,7 +10,7 @@ rule import_vesselfm_model:
 
 rule run_vesselfm:
     input:
-        spim=inputs["spim"].path,
+        spim=spim_input,
         model_path="resources/models/vesselfm.pt",
     params:
         zarrnii_kwargs=zarrnii_in_kwargs,