itasc 0.2.0.dev0__py3-none-any.whl

itasc/cellpose/divergence_maps.py

@@ -0,0 +1,380 @@
+"""Divergence-based foreground & contour maps from Cellpose prob/dp outputs.
+
+Replaces the (cellprob x z) mask sweep with a direct computation:
+
+    foreground = reduce_z(sigmoid(prob))
+    contours   = reduce_z(clip(div(filter(dp)), 0, inf))
+
+See ``notes/divergence_maps_spec.md`` for the rationale.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+import tifffile
+from scipy.ndimage import gaussian_filter, median_filter
+
+from itasc.core.cancellation import CancelledError
+from itasc.core.tiff import imwrite_grayscale
+
+
+def _check_cancel(cancel: Callable[[], bool] | None) -> None:
+    if cancel is not None and cancel():
+        raise CancelledError("Operation cancelled.")
+
+
+ZReduction = Literal["mean", "max"]
+
+
+@dataclass(frozen=True, slots=True)
+class DivergenceMapsReport:
+    """Summary returned by :func:`build_divergence_maps`."""
+
+    frames: int
+    foreground_z_reduction: ZReduction
+    contour_z_reduction: ZReduction
+    smoothing_sigma: float
+    median_radius: int
+    contours_path: Path
+    foreground_path: Path
+    foreground_smoothing_sigma: float = 0.0
+    foreground_median_radius: int = 0
+
+
+def sigmoid(x: np.ndarray) -> np.ndarray:
+    """Numerically stable sigmoid on float32 logits."""
+    x = np.clip(np.asarray(x, dtype=np.float32), -88.0, 88.0)
+    return (1.0 / (1.0 + np.exp(-x))).astype(np.float32, copy=False)
+
+
+def _reduce_z(arr_tzyx: np.ndarray, reduction: ZReduction) -> np.ndarray:
+    if reduction == "mean":
+        return arr_tzyx.mean(axis=1, dtype=np.float32).astype(np.float32, copy=False)
+    if reduction == "max":
+        return arr_tzyx.max(axis=1).astype(np.float32, copy=False)
+    raise ValueError(f"reduction must be 'mean' or 'max', got {reduction!r}")
+
+
+def _filter_2d(
+    arr_yx: np.ndarray, *, smoothing_sigma: float, median_radius: int,
+) -> np.ndarray:
+    """Apply median -> gaussian to a single ``(Y, X)`` map. Order matters (spec)."""
+    out = np.asarray(arr_yx, dtype=np.float32)
+    if median_radius > 0:
+        out = median_filter(out, size=2 * int(median_radius) + 1)
+    if smoothing_sigma > 0.0:
+        out = gaussian_filter(out, sigma=float(smoothing_sigma))
+    return out.astype(np.float32, copy=False)
+
+
+def foreground_from_prob(
+    prob_tzyx: np.ndarray,
+    *,
+    reduction: ZReduction,
+    smoothing_sigma: float = 0.0,
+    median_radius: int = 0,
+) -> np.ndarray:
+    """``sigmoid(prob)`` reduced across z, optionally smoothed.
+
+    Returns ``(T, Y, X)`` float32 in [0, 1]. When ``smoothing_sigma`` or
+    ``median_radius`` is set, each reduced frame is median- then gaussian-
+    filtered (defaults of ``0`` leave the foreground untouched).
+    """
+    p = sigmoid(prob_tzyx)
+    fg = _reduce_z(p, reduction)
+    if smoothing_sigma > 0.0 or median_radius > 0:
+        fg = np.stack(
+            [
+                _filter_2d(
+                    fg[t], smoothing_sigma=smoothing_sigma, median_radius=median_radius,
+                )
+                for t in range(fg.shape[0])
+            ]
+        )
+    return fg
+
+
+def divergence_2d(flow_yx: np.ndarray) -> np.ndarray:
+    """Divergence of a ``(2, Y, X)`` flow field with channels ``[dy, dx]``."""
+    flow_yx = np.asarray(flow_yx, dtype=np.float32)
+    if flow_yx.ndim != 3 or flow_yx.shape[0] != 2:
+        raise ValueError(
+            f"flow must be (2, Y, X) with channels [dy, dx]; got {flow_yx.shape}"
+        )
+    # np.gradient needs >= 2 samples along an axis; a singleton Y or X (e.g. a
+    # 1-pixel crop) has zero divergence contribution there, so skip it.
+    d_dy = (
+        np.gradient(flow_yx[0], axis=0)
+        if flow_yx.shape[1] >= 2
+        else np.zeros_like(flow_yx[0])
+    )
+    d_dx = (
+        np.gradient(flow_yx[1], axis=1)
+        if flow_yx.shape[2] >= 2
+        else np.zeros_like(flow_yx[1])
+    )
+    return (d_dy + d_dx).astype(np.float32, copy=False)
+
+
+def _filter_flow(
+    flow_2yx: np.ndarray, *, smoothing_sigma: float, median_radius: int,
+) -> np.ndarray:
+    """Apply median -> gaussian per channel. Order matters (spec)."""
+    out = flow_2yx
+    if median_radius > 0:
+        size = 2 * int(median_radius) + 1
+        out = np.stack(
+            [median_filter(out[0], size=size), median_filter(out[1], size=size)],
+            axis=0,
+        )
+    if smoothing_sigma > 0.0:
+        sigma = float(smoothing_sigma)
+        out = np.stack(
+            [gaussian_filter(out[0], sigma=sigma), gaussian_filter(out[1], sigma=sigma)],
+            axis=0,
+        )
+    return out
+
+
+def contour_from_dp(
+    dp_tzcyx: np.ndarray,
+    *,
+    smoothing_sigma: float,
+    median_radius: int,
+    reduction: ZReduction,
+) -> np.ndarray:
+    """Per (t, z): filter -> divergence -> clip(>=0); then reduce across z.
+
+    ``dp_tzcyx``: shape ``(T, Z, 2, Y, X)`` with channels ``[dy, dx]``.
+    Returns ``(T, Y, X)`` float32.
+    """
+    arr = np.asarray(dp_tzcyx, dtype=np.float32)
+    if arr.ndim != 5 or arr.shape[2] != 2:
+        raise ValueError(
+            f"dp must be (T, Z, 2, Y, X) with channels [dy, dx]; got {arr.shape}"
+        )
+    n_t, n_z, _, n_y, n_x = arr.shape
+    pos = np.empty((n_t, n_z, n_y, n_x), dtype=np.float32)
+    for t in range(n_t):
+        pos[t] = _positive_divergence_z_stack(
+            arr[t],
+            smoothing_sigma=smoothing_sigma,
+            median_radius=median_radius,
+        )
+    return _reduce_z(pos, reduction)
+
+
+def _positive_divergence_z_stack(
+    dp_zcyx: np.ndarray,
+    *,
+    smoothing_sigma: float,
+    median_radius: int,
+    progress_cb: Callable[[int, int], None] | None = None,
+) -> np.ndarray:
+    n_z, _, n_y, n_x = dp_zcyx.shape
+    pos = np.empty((n_z, n_y, n_x), dtype=np.float32)
+    for z in range(n_z):
+        filt = _filter_flow(
+            dp_zcyx[z], smoothing_sigma=smoothing_sigma, median_radius=median_radius,
+        )
+        div = divergence_2d(filt)
+        np.clip(div, 0.0, None, out=div)
+        pos[z] = div
+        if progress_cb is not None:
+            progress_cb(z + 1, n_z)
+    return pos
+
+
+class _LazyTiffStack:
+    """Frame-at-a-time reader for a (T,Z,Y,X) / (T,Z,2,Y,X) TIFF stack.
+
+    Reads only the TIFF header on construction (so callers can validate shapes
+    without paying for a full decompress) and decodes one frame's pages per
+    :meth:`frame` call. Compressed stacks store one page per 2D plane, so a
+    per-frame ``key`` slice touches only that frame's pages — this is what lets
+    :func:`build_divergence_maps` start reporting progress immediately instead
+    of stalling while both whole stacks are read up front.
+    """
+
+    def __init__(self, path: str | Path, *, ndim: int, name: str) -> None:
+        self._path = Path(path)
+        self._name = name
+        with tifffile.TiffFile(str(self._path)) as tf:
+            series = tf.series[0]
+            disk_shape = tuple(int(x) for x in series.shape)
+            axes = str(series.axes)
+        shape = self._canonical_shape(disk_shape, axes, ndim, name)
+        self.shape = shape
+        self._frame_shape = shape[1:]
+        # Pages per frame = product of the non-(Y,X) per-frame axes (Z, or Z*2).
+        self._pages_per_frame = int(np.prod(shape[1:-2], dtype=int))
+
+    @staticmethod
+    def _canonical_shape(
+        disk_shape: tuple[int, ...], axes: str, ndim: int, name: str,
+    ) -> tuple[int, ...]:
+        """Recover the logical ``(T,Z,Y,X)`` / ``(T,Z,2,Y,X)`` shape.
+
+        TIFF drops singleton leading axes on write, so a 2D+t prob ``(T,1,Y,X)``
+        lands on disk as ``(T,Y,X)`` — indistinguishable from a single z-stack by
+        shape alone. When the writer recorded axis labels (our ``write_outputs``),
+        map the surviving letters back to canonical order, inserting size-1 for
+        any dropped ``T``/``Z``/``C``. Otherwise fall back to the length heuristic
+        for legacy/metadata-less files (which are never singleton-squeezed).
+        """
+        canonical = "TZYX" if ndim == 4 else "TZCYX"
+        label = "Z×Y×X or T×Z×Y×X" if ndim == 4 else "Z×2×Y×X or T×Z×2×Y×X"
+        if len(axes) == len(disk_shape) and set("YX") <= set(axes) <= set(canonical):
+            present = dict(zip(axes, disk_shape))
+            shape = tuple(present.get(a, 1) for a in canonical)
+        else:
+            shape = disk_shape
+            # A frameless 3D/4D stack is a single timepoint: prepend T=1.
+            if len(shape) == ndim - 1:
+                shape = (1,) + shape
+        if len(shape) != ndim:
+            raise ValueError(f"{name} must be {label}.")
+        if ndim == 5 and shape[2] != 2:
+            if shape[1] == 3:
+                raise ValueError(
+                    f"{name} has shape {shape}, which is Cellpose 3D flows "
+                    f"(T×3×Z×Y×X from do_3d=True). The divergence builder needs "
+                    f"per-z 2-component flows (T×Z×2×Y×X); re-run nucleus Cellpose "
+                    f"with do_3d=False for the divergence pipeline."
+                )
+            raise ValueError(f"{name} must be Z×2×Y×X or T×Z×2×Y×X.")
+        return shape
+
+    def frame(self, t: int) -> np.ndarray:
+        """Frame ``t`` as float32, shaped ``(Z,Y,X)`` or ``(Z,2,Y,X)``."""
+        start = t * self._pages_per_frame
+        arr = tifffile.imread(
+            str(self._path), key=slice(start, start + self._pages_per_frame),
+        )
+        return np.asarray(arr, dtype=np.float32).reshape(self._frame_shape)
+
+
+def build_divergence_maps(
+    prob_path: str | Path,
+    dp_path: str | Path,
+    contours_out: str | Path,
+    foreground_out: str | Path,
+    *,
+    foreground_z_reduction: ZReduction,
+    contour_z_reduction: ZReduction,
+    smoothing_sigma: float,
+    median_radius: int,
+    foreground_smoothing_sigma: float = 0.0,
+    foreground_median_radius: int = 0,
+    progress_cb: Callable[[int, int, str], None] | None = None,
+    cancel: Callable[[], bool] | None = None,
+) -> DivergenceMapsReport:
+    """Compute and write ``contours`` and ``foreground`` from Cellpose prob/dp.
+
+    Output stacks are ``T x Y x X`` float32. ``progress_cb`` is called for each
+    foreground frame, contour z-slice, and output write. ``smoothing_sigma`` /
+    ``median_radius`` filter the flow before divergence (contours); the
+    ``foreground_*`` knobs filter the reduced foreground map independently.
+    """
+    prob_stack = _LazyTiffStack(prob_path, ndim=4, name="prob")
+    dp_stack = _LazyTiffStack(dp_path, ndim=5, name="dp")
+    if prob_stack.shape[0] != dp_stack.shape[0]:
+        raise ValueError("prob and dp must have the same frame count.")
+    if prob_stack.shape[1] != dp_stack.shape[1]:
+        raise ValueError("prob and dp must have the same z count.")
+    if prob_stack.shape[2:] != dp_stack.shape[3:]:
+        raise ValueError("prob and dp must have the same Y×X shape.")
+
+    n_t = int(prob_stack.shape[0])
+    n_z = int(prob_stack.shape[1])
+    if n_t == 0:
+        raise ValueError("prob/dp stack has no frames (T=0); nothing to compute.")
+    total_steps = n_t + (n_t * n_z) + 2
+    progress_done = 0
+    contour_frames: list[np.ndarray] = []
+    foreground_frames: list[np.ndarray] = []
+    for t in range(n_t):
+        _check_cancel(cancel)
+        fg = foreground_from_prob(
+            prob_stack.frame(t)[np.newaxis],
+            reduction=foreground_z_reduction,
+            smoothing_sigma=foreground_smoothing_sigma,
+            median_radius=foreground_median_radius,
+        )
+        foreground_frames.append(fg[0])
+        progress_done += 1
+        if progress_cb is not None:
+            progress_cb(
+                progress_done,
+                total_steps,
+                f"Divergence maps: foreground frame {t + 1}/{n_t}",
+            )
+
+        def _progress_z(z_done: int, z_total: int) -> None:
+            nonlocal progress_done
+
+            _check_cancel(cancel)
+            progress_done += 1
+            if progress_cb is not None:
+                progress_cb(
+                    progress_done,
+                    total_steps,
+                    (
+                        f"Divergence maps: contours frame {t + 1}/{n_t} "
+                        f"z {z_done}/{z_total}"
+                    ),
+                )
+
+        pos_div = _positive_divergence_z_stack(
+            dp_stack.frame(t),
+            smoothing_sigma=smoothing_sigma,
+            median_radius=median_radius,
+            progress_cb=_progress_z,
+        )
+        contour_frames.append(
+            _reduce_z(pos_div[np.newaxis, ...], contour_z_reduction)[0]
+        )
+
+    _check_cancel(cancel)
+    contours_out = Path(contours_out)
+    foreground_out = Path(foreground_out)
+    contours_out.parent.mkdir(parents=True, exist_ok=True)
+    foreground_out.parent.mkdir(parents=True, exist_ok=True)
+    imwrite_grayscale(
+        contours_out, np.stack(contour_frames).astype(np.float32),
+        compression="zlib", metadata={"axes": "TYX"},
+    )
+    progress_done += 1
+    if progress_cb is not None:
+        progress_cb(progress_done, total_steps, "Divergence maps: writing contours")
+    imwrite_grayscale(
+        foreground_out, np.stack(foreground_frames).astype(np.float32),
+        compression="zlib", metadata={"axes": "TYX"},
+    )
+    progress_done += 1
+    if progress_cb is not None:
+        progress_cb(progress_done, total_steps, "Divergence maps: writing foreground")
+    return DivergenceMapsReport(
+        frames=n_t,
+        foreground_z_reduction=foreground_z_reduction,
+        contour_z_reduction=contour_z_reduction,
+        smoothing_sigma=float(smoothing_sigma),
+        median_radius=int(median_radius),
+        contours_path=contours_out,
+        foreground_path=foreground_out,
+    )
+
+
+__all__ = [
+    "CancelledError",
+    "DivergenceMapsReport",
+    "build_divergence_maps",
+    "contour_from_dp",
+    "divergence_2d",
+    "foreground_from_prob",
+    "sigmoid",
+]

itasc/cellpose/flow_following.py

@@ -0,0 +1,283 @@
+"""Flow-following cell segmentation, Qt-free, for the standalone distro.
+
+When the standalone tool has **both** a nucleus and a cell channel, it can do
+better than independent native masks: it assigns each cell-foreground pixel to a
+*nucleus*, giving one cell per nucleus with the cell sharing the nucleus' id.
+The cell's *extent* is decided by Cellpose's own flow field rather than by naive
+distance.
+
+This is the proven two-phase algorithm (recovered from the pre-publication
+``segmentation/flow_following.py``):
+
+1. **Integrate.** Each foreground pixel is Euler-integrated along a blend of the
+   Cellpose flow vector and an EDT *gravity* vector pointing toward the nearest
+   nucleus (``flow_weight`` mixes the two). A pixel that lands on a nucleus pixel
+   during integration inherits that nucleus' label immediately.
+2. **Grow.** Remaining foreground pixels are assigned by growing the labelled
+   region outward in shells through the *displaced* (post-integration) positions,
+   so labels chain-propagate along the flow topology. Pixels whose displaced
+   position has no labelled neighbour within ``max_assign_radius`` stay
+   background (orphans are dropped, not force-assigned).
+
+The numba kernel is lazily JIT-compiled when numba is importable (it ships with
+the Cellpose stack) and falls back to the identical pure-Python function
+otherwise, so importing this module never requires numba.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+
+__all__ = [
+    "FlowFollowingParams",
+    "flow_follow_frame",
+    "flow_follow_movie",
+]
+
+
+@dataclass(frozen=True)
+class FlowFollowingParams:
+    """Parameters for flow-following cell assignment."""
+
+    fg_threshold: float = 0.5
+    """Cell-foreground cutoff on the sigmoid probability (``sigmoid(prob) >
+    fg_threshold``). Exposed knob — pixels below it are not assigned."""
+
+    flow_weight: float = 0.5
+    """Blend between the Cellpose flow vector (``1``) and the EDT gravity vector
+    toward the nearest nucleus (``0``). ``0.5`` mixes both equally."""
+
+    flow_step_scale: float = 0.2
+    """Euler step size (multiplies the unit blended vector each iteration)."""
+
+    max_iterations: int = 100
+    """Maximum integration steps per pixel."""
+
+    shell_width: float = 5.0
+    """Per-iteration capture distance when growing labels in shells (px)."""
+
+    max_assign_radius: float = 30.0
+    """A displaced foreground pixel with no labelled pixel within this radius
+    stays background — orphan foreground is dropped, not force-assigned."""
+
+
+# ── Phase 1: Euler integration (lazy-JIT kernel + pure-Python twin) ──────────
+
+def _integrate_positions_py(
+    nuclear_labels,   # (H, W) int32
+    flow,             # (2, H, W) float32 — flow[0]=dy, flow[1]=dx
+    grav_y,           # (H, W) float32
+    grav_x,           # (H, W) float32
+    prob_mask,        # (H, W) bool
+    n_steps,          # int
+    flow_step_scale,  # float
+    flow_weight,      # float
+):
+    """Integrate each foreground pixel along ``flow_weight*flow + (1-w)*gravity``.
+
+    Returns ``(result, final_y, final_x)``: a label map where nucleus pixels keep
+    their label and foreground pixels that land directly on a nucleus inherit it,
+    plus the displaced position of every pixel after integration.
+    """
+    H = nuclear_labels.shape[0]
+    W = nuclear_labels.shape[1]
+    result = nuclear_labels.copy()
+    final_y = np.empty((H, W), dtype=np.float32)
+    final_x = np.empty((H, W), dtype=np.float32)
+
+    for i in range(H):
+        for j in range(W):
+            final_y[i, j] = np.float32(i)
+            final_x[i, j] = np.float32(j)
+
+            if nuclear_labels[i, j] > 0 or not prob_mask[i, j]:
+                continue
+
+            py = np.float32(i)
+            px = np.float32(j)
+            for _ in range(n_steps):
+                iy = min(max(int(py), 0), H - 1)
+                ix = min(max(int(px), 0), W - 1)
+                fy = (
+                    flow_weight * flow[0, iy, ix]
+                    + (np.float32(1.0) - flow_weight) * grav_y[iy, ix]
+                )
+                fx = (
+                    flow_weight * flow[1, iy, ix]
+                    + (np.float32(1.0) - flow_weight) * grav_x[iy, ix]
+                )
+                py = py + fy * flow_step_scale
+                px = px + fx * flow_step_scale
+                py = min(max(py, np.float32(0.0)), np.float32(H - 1))
+                px = min(max(px, np.float32(0.0)), np.float32(W - 1))
+                if nuclear_labels[int(py), int(px)] > 0:
+                    result[i, j] = nuclear_labels[int(py), int(px)]
+                    break
+
+            final_y[i, j] = py
+            final_x[i, j] = px
+
+    return result, final_y, final_x
+
+
+def _integrate_positions(*args):
+    """Dispatch to the numba-compiled kernel when available, else pure Python."""
+    global _INTEGRATE
+    if _INTEGRATE is None:
+        try:  # numba ships with the cellpose stack; optional otherwise.
+            import numba
+
+            _INTEGRATE = numba.njit(cache=True)(_integrate_positions_py)
+        except Exception:  # pragma: no cover - exercised only without numba
+            _INTEGRATE = _integrate_positions_py
+    return _INTEGRATE(*args)
+
+
+_INTEGRATE = None
+
+
+# ── Phase 2: progressive shell growth (orphan-bounded) ───────────────────────
+
+def _progressive_shell_assign(
+    labels: np.ndarray,
+    final_y: np.ndarray,
+    final_x: np.ndarray,
+    foreground: np.ndarray,
+    *,
+    shell_width: float,
+    max_assign_radius: float,
+) -> np.ndarray:
+    """Grow labels outward through displaced positions; drop distant orphans.
+
+    Each iteration computes the EDT from currently-labelled pixels and assigns any
+    unassigned foreground pixel whose *displaced* position is within
+    ``shell_width`` of a labelled pixel. Newly labelled pixels seed the next EDT,
+    so labels chain-propagate along the flow topology. The growth stops once no
+    displaced position is within ``max_assign_radius`` of a label — remaining
+    pixels stay background.
+    """
+    H, W = labels.shape
+    result = labels.copy()
+    unassigned = foreground & (result == 0)
+
+    fy = np.clip(np.round(final_y).astype(np.intp), 0, H - 1)
+    fx = np.clip(np.round(final_x).astype(np.intp), 0, W - 1)
+
+    while unassigned.any():
+        unlabelled = result == 0
+        if not unlabelled.any():
+            break
+        dist, indices = distance_transform_edt(unlabelled, return_indices=True)
+        d = dist[fy, fx]
+        nearest = result[indices[0][fy, fx], indices[1][fy, fx]]
+        can_assign = unassigned & (d <= shell_width) & (nearest > 0)
+        if not can_assign.any():
+            break
+        result[can_assign] = nearest[can_assign]
+        unassigned &= ~can_assign
+
+    # Bounded fallback (orphan drop): assign any remaining pixel only if a label
+    # sits within max_assign_radius of its displaced position; else leave it 0.
+    if unassigned.any() and (result > 0).any():
+        unlabelled = result == 0
+        if unlabelled.any():
+            dist, indices = distance_transform_edt(unlabelled, return_indices=True)
+            d = dist[fy, fx]
+            nearest = result[indices[0][fy, fx], indices[1][fy, fx]]
+            close = unassigned & (d <= max_assign_radius) & (nearest > 0)
+            result[close] = nearest[close]
+
+    return result
+
+
+# ── Gravity field ────────────────────────────────────────────────────────────
+
+def _gravity_toward_nuclei(labels_yx: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Unit vectors at every pixel pointing toward the nearest nucleus pixel."""
+    nucleus = labels_yx > 0
+    _dist, indices = distance_transform_edt(~nucleus, return_indices=True)
+    H, W = labels_yx.shape
+    yy, xx = np.mgrid[:H, :W]
+    dy = (indices[0] - yy).astype(np.float32)
+    dx = (indices[1] - xx).astype(np.float32)
+    norm = np.hypot(dy, dx).astype(np.float32)
+    norm[norm == 0.0] = 1.0
+    return (dy / norm).astype(np.float32), (dx / norm).astype(np.float32)
+
+
+# ── Public: single frame + movie ─────────────────────────────────────────────
+
+def flow_follow_frame(
+    foreground_yx: np.ndarray,
+    dp_cyx: np.ndarray,
+    labels_yx: np.ndarray,
+    params: FlowFollowingParams,
+) -> np.ndarray:
+    """Assign one frame's cell foreground to nuclei via flow-following.
+
+    ``foreground_yx`` is a bool mask, ``dp_cyx`` the Cellpose cell flow
+    ``(2, Y, X)`` (``dy, dx``), ``labels_yx`` the tracked nucleus labels. Returns
+    ``(Y, X)`` int32 cell labels carrying the nucleus ids (background ``0``).
+    """
+    foreground_yx = np.ascontiguousarray(foreground_yx, dtype=np.bool_)
+    labels_yx = np.ascontiguousarray(labels_yx, dtype=np.int32)
+    if not foreground_yx.any() or labels_yx.max() == 0:
+        return np.zeros(foreground_yx.shape, dtype=np.int32)
+
+    # Normalise flow by mean foreground magnitude so the step scale is comparable
+    # across frames/datasets.
+    flow = np.ascontiguousarray(dp_cyx, dtype=np.float32).copy()
+    mag = np.sqrt(flow[0] ** 2 + flow[1] ** 2)
+    mean_mag = float(mag[foreground_yx].mean()) if foreground_yx.any() else 1.0
+    if mean_mag > 0:
+        flow /= np.float32(mean_mag)
+
+    grav_y, grav_x = _gravity_toward_nuclei(labels_yx)
+    result, final_y, final_x = _integrate_positions(
+        labels_yx,
+        flow,
+        np.ascontiguousarray(grav_y),
+        np.ascontiguousarray(grav_x),
+        foreground_yx,
+        int(params.max_iterations),
+        np.float32(params.flow_step_scale),
+        np.float32(params.flow_weight),
+    )
+    return _progressive_shell_assign(
+        result, final_y, final_x, foreground_yx,
+        shell_width=float(params.shell_width),
+        max_assign_radius=float(params.max_assign_radius),
+    ).astype(np.int32)
+
+
+def flow_follow_movie(
+    foreground_tyx: np.ndarray,
+    dp_tcyx: np.ndarray,
+    labels_tyx: np.ndarray,
+    params: FlowFollowingParams,
+    *,
+    progress_cb=None,
+) -> np.ndarray:
+    """Run :func:`flow_follow_frame` over a ``(T, Y, X)`` series.
+
+    ``dp_tcyx`` is ``(T, 2, Y, X)``. Cell labels inherit the (already tracked)
+    nucleus ids per frame, so tracking the nuclei tracks the cells. Returns
+    ``(T, Y, X)`` int32.
+    """
+    foreground = np.asarray(foreground_tyx, dtype=bool)
+    labels = np.asarray(labels_tyx, dtype=np.int32)
+    flow = np.asarray(dp_tcyx, dtype=np.float32)
+    if foreground.ndim != 3 or labels.shape != foreground.shape:
+        raise ValueError("foreground and labels must share shape (T, Y, X)")
+    T, Y, X = foreground.shape
+    if flow.shape != (T, 2, Y, X):
+        raise ValueError(f"dp_tcyx must be (T, 2, Y, X); got {flow.shape}")
+
+    out = np.zeros((T, Y, X), dtype=np.int32)
+    for t in range(T):
+        out[t] = flow_follow_frame(foreground[t], flow[t], labels[t], params)
+        if progress_cb is not None:
+            progress_cb(t + 1, T, f"Flow-following: frame {t + 1}/{T}...")
+    return out