Use temporary local Zarr for regionprops to prevent file errors

spimquant/workflow/Snakefile

@@ -314,6 +314,7 @@ if config.get("analysis_level") == "group":
 
 
 wildcard_constraints:
+    acq="[a-zA-Z0-9]+",
     stain="[a-zA-Z0-9]+",
     pairwise_contrast="[a-zA-Z0-9+_-]+",
     n4_spline="[0-9]+",

spimquant/workflow/rules/regionprops.smk

@@ -28,8 +28,8 @@ rule compute_filtered_regionprops:
         ),
     threads: 64 if config["dask_scheduler"] == "distributed" else 32
     resources:
-        mem_mb=256000,
-        runtime=180,
+        mem_mb=500000,
+        runtime=360,
     script:
         "../scripts/compute_filtered_regionprops.py"
 

spimquant/workflow/rules/segstats.smk

@@ -143,7 +143,7 @@ rule merge_into_segstats_tsv:
         ),
     threads: 1
     resources:
-        mem_mb=1500,
+        mem_mb=16000,
         runtime=15,
     script:
         "../scripts/merge_into_segstats_tsv.py"

spimquant/workflow/scripts/compute_filtered_regionprops.py

@@ -1,23 +1,37 @@
 """Compute region properties from filtered segmentation masks using ZarrNii.
-
 This script reads a segmentation mask from an OME-Zarr file, performs
-connected components on chunks with overlap, applys filters based on 
+connected components on chunks with overlap, applies filters based on
 region properties, and outputs region properties on these filtered objects
 """
 
+import tempfile
+import zipfile
+from contextlib import contextmanager
+
 from dask_setup import get_dask_client
 from zarrnii import ZarrNii
 
+
+@contextmanager
+def get_zarr_path(mask_path):
+    """Yield the path to the OME-Zarr store, extracting from zip if needed."""
+    if mask_path.endswith(".ozx") or mask_path.endswith(".zip"):
+        with tempfile.TemporaryDirectory(suffix=".ome.zarr") as temp_dir:
+            print(f"Extracting zip archive to temporary directory: {temp_dir}")
+            with zipfile.ZipFile(mask_path, "r") as zip_ref:
+                zip_ref.extractall(temp_dir)
+            yield temp_dir
+    else:
+        yield mask_path
+
+
 if __name__ == "__main__":
     with get_dask_client(snakemake.config["dask_scheduler"], snakemake.threads):
+        with get_zarr_path(snakemake.input.mask) as zarr_path:
+            znimg = ZarrNii.from_file(zarr_path, level=0)
 
-        znimg = ZarrNii.from_file(
-            snakemake.input.mask,
-            level=0,  # input image is already downsampled to the wildcard level
-        )
-
-        znimg.compute_region_properties(
-            output_path=snakemake.output.regionprops_parquet,
-            region_filters=snakemake.params.region_filters,
-            output_properties=snakemake.params.output_properties,
-        )
+            znimg.compute_region_properties(
+                output_path=snakemake.output.regionprops_parquet,
+                region_filters=snakemake.params.region_filters,
+                output_properties=snakemake.params.output_properties,
+            )