ngio 0.3.5__py3-none-any.whl → 0.4.0a2__py3-none-any.whl

ngio/common/_masking_roi.py

@@ -2,11 +2,10 @@
 
 import itertools
 
-import dask
 import dask.array as da
-import dask.delayed
 import numpy as np
 import scipy.ndimage as ndi
+from dask.delayed import delayed
 
 from ngio.common._roi import Roi, RoiPixels
 from ngio.ome_zarr_meta import PixelSize
@@ -39,7 +38,7 @@ def _adjust_slices(slices, offset):
     return adjusted_slices
 
 
-@dask.delayed
+@delayed
 def _process_chunk(chunk, offset):
     """Process a single chunk.
 
@@ -63,7 +62,7 @@ def _merge_slices(
     return tuple(merged)
 
 
-@dask.delayed
+@delayed
 def _collect_slices(
     local_slices: list[dict[int, tuple[slice, ...]]],
 ) -> dict[int, tuple[slice]]:
@@ -101,7 +100,7 @@ def compute_slices(segmentation: np.ndarray) -> dict[int, tuple[slice, ...]]:
 def lazy_compute_slices(segmentation: da.Array) -> dict[int, tuple[slice, ...]]:
     """Compute slices for each label in a segmentation."""
     global_offsets = _compute_offsets(segmentation.chunks)
-    delayed_chunks = segmentation.to_delayed()
+    delayed_chunks = segmentation.to_delayed()  # type: ignore
 
     grid_shape = tuple(len(c) for c in segmentation.chunks)
 
@@ -125,8 +124,8 @@ def compute_masking_roi(
     Other axes orders are not supported.
 
     """
-    if segmentation.ndim not in [2, 3]:
-        raise NgioValueError("Only 2D and 3D segmentations are supported.")
+    if segmentation.ndim not in [2, 3, 4]:
+        raise NgioValueError("Only 2D, 3D, and 4D segmentations are supported.")
 
     if isinstance(segmentation, da.Array):
         slices = lazy_compute_slices(segmentation)
@@ -136,21 +135,50 @@ def compute_masking_roi(
     rois = []
     for label, slice_ in slices.items():
         if len(slice_) == 2:
-            min_z, min_y, min_x = 0, slice_[0].start, slice_[1].start
-            max_z, max_y, max_x = 1, slice_[0].stop, slice_[1].stop
+            min_t, max_t = None, None
+            min_z, max_z = None, None
+            min_y, min_x = slice_[0].start, slice_[1].start
+            max_y, max_x = slice_[0].stop, slice_[1].stop
         elif len(slice_) == 3:
+            min_t, max_t = None, None
             min_z, min_y, min_x = slice_[0].start, slice_[1].start, slice_[2].start
             max_z, max_y, max_x = slice_[0].stop, slice_[1].stop, slice_[2].stop
+        elif len(slice_) == 4:
+            min_t, min_z, min_y, min_x = (
+                slice_[0].start,
+                slice_[1].start,
+                slice_[2].start,
+                slice_[3].start,
+            )
+            max_t, max_z, max_y, max_x = (
+                slice_[0].stop,
+                slice_[1].stop,
+                slice_[2].stop,
+                slice_[3].stop,
+            )
         else:
             raise ValueError("Invalid slice length.")
+
+        if max_t is None:
+            t_length = None
+        else:
+            t_length = max_t - min_t
+
+        if max_z is None:
+            z_length = None
+        else:
+            z_length = max_z - min_z
+
         roi = RoiPixels(
             name=str(label),
             x_length=max_x - min_x,
             y_length=max_y - min_y,
-            z_length=max_z - min_z,
+            z_length=z_length,
+            t_length=t_length,
             x=min_x,
             y=min_y,
             z=min_z,
+            label=label,
         )
 
         roi = roi.to_roi(pixel_size)

ngio/common/_pyramid.py

@@ -1,5 +1,5 @@
 import math
-from collections.abc import Collection
+from collections.abc import Callable, Sequence
 from typing import Literal
 
 import dask.array as da
@@ -40,7 +40,7 @@ def _on_disk_coarsen(
     source: zarr.Array,
     target: zarr.Array,
     _order: Literal[0, 1] = 1,
-    aggregation_function: np.ufunc | None = None,
+    aggregation_function: Callable | None = None,
 ) -> None:
     """Apply a coarsening operation from a source zarr array to a target zarr array.
 
@@ -132,7 +132,7 @@ def on_disk_zoom(
 
 def _find_closest_arrays(
     processed: list[zarr.Array], to_be_processed: list[zarr.Array]
-) -> tuple[int, int]:
+) -> tuple[np.intp, np.intp]:
     dist_matrix = np.zeros((len(processed), len(to_be_processed)))
     for i, arr_to_proc in enumerate(to_be_processed):
         for j, proc_arr in enumerate(processed):
@@ -147,7 +147,9 @@ def _find_closest_arrays(
                 )
             )
 
-    return np.unravel_index(dist_matrix.argmin(), dist_matrix.shape)
+    indices = np.unravel_index(dist_matrix.argmin(), dist_matrix.shape)
+    assert len(indices) == 2, "Indices must be of length 2"
+    return indices
 
 
 def consolidate_pyramid(
@@ -178,9 +180,9 @@ def consolidate_pyramid(
 def init_empty_pyramid(
     store: StoreOrGroup,
     paths: list[str],
-    ref_shape: Collection[int],
-    scaling_factors: Collection[float],
-    chunks: Collection[int] | None = None,
+    ref_shape: Sequence[int],
+    scaling_factors: Sequence[float],
+    chunks: Sequence[int] | None = None,
     dtype: str = "uint16",
     mode: AccessModeLiteral = "a",
 ) -> None: