From e03aff40e4a230aa524f747acc69401abb39bc69 Mon Sep 17 00:00:00 2001 From: Tim Treis Date: Wed, 3 Jun 2026 20:52:11 +0200 Subject: [PATCH] feat(primitives): canonical (B,Z,Y,X) batch-shape helper (to_bzyx) Shared foundational helper used by the numba intensity/granularity/zernike backends to normalise any input (2D/3D/4D/list) to the canonical batch-of-volumes form: single image = batch of 1, returning a dict for a lone image/volume and a list of dicts for a batch. Pure numpy, no numba. Extracted to its own PR so it can be reviewed first and unblock the feature backends (#56/#57/#58). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/cp_measure/primitives/shapes.py | 83 +++++++++++++++++++++++++++++ test/test_primitives_shapes.py | 80 +++++++++++++++++++++++++++ 2 files changed, 163 insertions(+) create mode 100644 src/cp_measure/primitives/shapes.py create mode 100644 test/test_primitives_shapes.py diff --git a/src/cp_measure/primitives/shapes.py b/src/cp_measure/primitives/shapes.py new file mode 100644 index 0000000..0c94059 --- /dev/null +++ b/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 diff --git a/test/test_primitives_shapes.py b/test/test_primitives_shapes.py new file mode 100644 index 0000000..0584e41 --- /dev/null +++ b/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)