feat(primitives): (B,Z,Y,X) batch-shape helper (to_bzyx)

src/cp_measure/primitives/shapes.py

@@ -0,0 +1,83 @@
+"""Canonical ``(B, Z, Y, X)`` input normalisation (numpy, backend-agnostic).
+
+Every optimised feature works on a batch of volumes. Rather than a single dense
+``(B, Z, Y, X)`` array (which cannot hold ragged, differently-sized images), the
+batch is represented as a **list of B ``(Z, Y, X)`` arrays** — equal-size and
+ragged batches share one code path, and a single image is just ``B == 1``.
+
+Normalisation rules (the only public entry, :func:`to_bzyx`):
+
+================================ ============================ ========
+input                            yields                       batch?
+================================ ============================ ========
+2D ``(H, W)`` ndarray            ``[(1, H, W)]``              no
+3D ``(Z, Y, X)`` ndarray         ``[(Z, Y, X)]`` (one volume) no
+4D ``(B, Z, Y, X)`` ndarray      ``[(Z, Y, X)] * B``          yes
+list/tuple of 2D/3D arrays       one ``(Z, Y, X)`` per item   yes
+================================ ============================ ========
+
+A 3D ndarray is therefore ALWAYS one volume, never a batch — this preserves the
+existing single-volume semantics. To pass a batch of 2D images as an array, use
+``(B, 1, H, W)``; for ragged sizes, pass a list. ``unwrap`` then re-shapes the
+per-image results back to a single dict (single input) or the list (batch).
+
+This is a pure batch normaliser: ``masks`` and ``pixels`` must share the same
+batch/ndim structure. It does NOT broadcast a lower-dimensional mask over a
+volume (e.g. a 2D mask applied to a 3D stack) — per-element ndim handling is the
+caller's job (each backend dispatches its own 2D vs 3D path).
+"""
+
+import numpy
+from numpy.typing import NDArray
+
+
+def _to_zyx(arr: NDArray) -> NDArray:
+    """Promote a single 2D/3D image to ``(Z, Y, X)`` (2D gets a unit Z axis)."""
+    a = numpy.asarray(arr)
+    if a.ndim == 2:
+        return a[numpy.newaxis]
+    if a.ndim == 3:
+        return a
+    raise ValueError(f"expected a 2D or 3D image, got ndim={a.ndim}")
+
+
+def to_bzyx(masks, pixels):
+    """Normalise ``(masks, pixels)`` to the canonical batch-of-volumes form.
+
+    Returns ``(masks_zyx, pixels_zyx, unwrap)`` where ``masks_zyx`` and
+    ``pixels_zyx`` are length-``B`` lists of ``(Z, Y, X)`` arrays (one per image),
+    and ``unwrap(results)`` maps a length-``B`` list of per-image results back to
+    a single result (non-batch input) or the list itself (batch input).
+    """
+    masks_is_seq = isinstance(masks, (list, tuple))
+    pixels_is_seq = isinstance(pixels, (list, tuple))
+    if masks_is_seq != pixels_is_seq:
+        raise ValueError("masks and pixels must both be sequences, or both arrays")
+
+    if masks_is_seq:
+        masks_zyx = [_to_zyx(m) for m in masks]
+        pixels_zyx = [_to_zyx(p) for p in pixels]
+        is_batch = True
+    else:
+        m = numpy.asarray(masks)
+        p = numpy.asarray(pixels)
+        if (m.ndim == 4) != (p.ndim == 4):
+            raise ValueError("masks and pixels must both be 4D for a stacked batch")
+        if m.ndim == 4:
+            masks_zyx = list(m)
+            pixels_zyx = list(p)
+            is_batch = True
+        else:
+            masks_zyx = [_to_zyx(m)]
+            pixels_zyx = [_to_zyx(p)]
+            is_batch = False
+
+    if len(masks_zyx) != len(pixels_zyx):
+        raise ValueError(
+            f"batch size mismatch: {len(masks_zyx)} masks vs {len(pixels_zyx)} images"
+        )
+
+    def unwrap(results):
+        return results if is_batch else results[0]
+
+    return masks_zyx, pixels_zyx, unwrap

test/test_primitives_shapes.py

@@ -0,0 +1,80 @@
+"""Tests for the canonical (B,Z,Y,X) input normalisation helper."""
+
+import numpy
+import pytest
+
+from cp_measure.primitives.shapes import to_bzyx
+
+
+def _img(shape, seed):
+    rng = numpy.random.default_rng(seed)
+    return rng.random(shape)
+
+
+def test_2d_is_single_unit_z():
+    m, p = numpy.ones((4, 5), int), _img((4, 5), 0)
+    masks, pixels, unwrap = to_bzyx(m, p)
+    assert len(masks) == len(pixels) == 1
+    assert masks[0].shape == (1, 4, 5)
+    assert pixels[0].shape == (1, 4, 5)
+    # single input -> unwrap returns the lone result, not a list
+    assert unwrap(["only"]) == "only"
+
+
+def test_3d_is_single_volume_not_batch():
+    m, p = numpy.ones((3, 4, 5), int), _img((3, 4, 5), 1)
+    masks, pixels, unwrap = to_bzyx(m, p)
+    assert len(masks) == 1
+    assert masks[0].shape == (3, 4, 5)  # one volume, Z preserved
+    assert unwrap(["d"]) == "d"
+
+
+def test_4d_is_batch():
+    m, p = numpy.ones((2, 3, 4, 5), int), _img((2, 3, 4, 5), 2)
+    masks, pixels, unwrap = to_bzyx(m, p)
+    assert len(masks) == len(pixels) == 2
+    assert all(a.shape == (3, 4, 5) for a in masks)
+    out = unwrap([{"a": 1}, {"a": 2}])
+    assert isinstance(out, list) and len(out) == 2  # batch -> list
+
+
+def test_list_of_2d_is_batch_unit_z():
+    imgs = [_img((4, 5), i) for i in range(3)]
+    masks = [numpy.ones((4, 5), int) for _ in range(3)]
+    m, p, unwrap = to_bzyx(masks, imgs)
+    assert len(m) == 3
+    assert all(a.shape == (1, 4, 5) for a in m)
+    assert isinstance(unwrap([1, 2, 3]), list)
+
+
+def test_list_ragged_sizes_ok():
+    imgs = [_img((4, 5), 0), _img((7, 3), 1)]
+    masks = [numpy.ones((4, 5), int), numpy.ones((7, 3), int)]
+    m, p, _ = to_bzyx(masks, imgs)
+    assert m[0].shape == (1, 4, 5)
+    assert m[1].shape == (1, 7, 3)
+
+
+def test_list_of_3d_volumes_batch():
+    vols = [_img((2, 4, 5), 0), _img((3, 4, 5), 1)]
+    masks = [numpy.ones((2, 4, 5), int), numpy.ones((3, 4, 5), int)]
+    m, _, _ = to_bzyx(masks, vols)
+    assert m[0].shape == (2, 4, 5) and m[1].shape == (3, 4, 5)
+
+
+@pytest.mark.parametrize(
+    "masks, pixels, match",
+    [
+        (
+            [numpy.ones((4, 5), int)],
+            [_img((4, 5), 0), _img((4, 5), 1)],
+            "batch size mismatch",
+        ),
+        ([numpy.ones((4, 5), int)], _img((4, 5), 0), "both be sequences"),
+        (numpy.ones((2, 3, 4, 5), int), _img((3, 4, 5), 0), "both be 4D"),
+        (numpy.ones((5,), int), numpy.ones((5,), float), "2D or 3D"),
+    ],
+)
+def test_to_bzyx_raises(masks, pixels, match):
+    with pytest.raises(ValueError, match=match):
+        to_bzyx(masks, pixels)