itasc 0.2.0.dev0__py3-none-any.whl

itasc/cellpose/joint.py

@@ -0,0 +1,98 @@
+"""Joint nucleus-anchored cell segmentation + tracking, Qt-free.
+
+When the standalone tool has both channels, this composes the existing pieces
+into one nucleus-anchored path:
+
+1. **Nucleus native masks** (:func:`native_masks.run_nucleus_masks_stack`) →
+   **track** them (:func:`track_laptrack.track_masks`) → tracked nucleus labels.
+2. **Cell flows** (:func:`cellpose_runner.run_cell_stack` → ``prob``, ``dp``);
+   the cell foreground is ``sigmoid(prob) > fg_threshold``.
+3. Per z-plane, **flow-follow** each cell-foreground pixel to a nucleus
+   (:func:`flow_following.flow_follow_movie`). Cell labels inherit the *tracked*
+   nucleus ids, so the cell stack is tracked by construction — one cell per
+   nucleus, sharing its id.
+
+Everything here reuses functions the app also imports (read-only); it adds no
+new behaviour to the app's Cellpose stage.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import numpy as np
+
+from itasc.cellpose import cellpose_runner, native_masks, track_laptrack
+from itasc.cellpose.cellpose_runner import CellParams, NucleusParams
+from itasc.cellpose.flow_following import FlowFollowingParams, flow_follow_movie
+
+__all__ = ["cell_foreground_from_prob", "joint_segment_track"]
+
+
+def cell_foreground_from_prob(prob: np.ndarray, threshold: float) -> np.ndarray:
+    """Boolean cell foreground from a Cellpose probability map via the sigmoid.
+
+    Matches the app's ``sigmoid(prob)`` foreground convention; ``threshold`` is
+    the exposed cutoff on the ``[0, 1]`` sigmoid.
+    """
+    sig = 1.0 / (1.0 + np.exp(-np.asarray(prob, dtype=np.float32)))
+    return sig > float(threshold)
+
+
+def joint_segment_track(
+    nucleus_stack: np.ndarray,
+    cell_stack: np.ndarray,
+    nucleus_params: NucleusParams,
+    cell_params: CellParams,
+    flow_params: FlowFollowingParams,
+    *,
+    max_distance: float = 15.0,
+    max_frame_gap: int = 0,
+    progress_cb: Callable[[int, int, str], None] | None = None,
+    cancel_cb: Callable[[], bool] | None = None,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Run the joint path; return ``(nucleus_tracked, cell_tracked)`` ``(T,Z,Y,X)``.
+
+    Both outputs are ``int32`` and share label ids: a cell carries the id of the
+    nucleus it was assigned to. ``max_distance``/``max_frame_gap`` tune the
+    nucleus tracker; the cell stack inherits its tracks.
+    """
+    nucleus_stack = np.asarray(nucleus_stack)
+    cell_stack = np.asarray(cell_stack)
+    if nucleus_stack.ndim != 4 or cell_stack.ndim != 4:
+        raise ValueError("nucleus_stack and cell_stack must be (T, Z, Y, X)")
+
+    def _report(frac_done: int, total: int, msg: str) -> None:
+        if progress_cb is not None:
+            progress_cb(frac_done, total, msg)
+
+    # 1. Nucleus masks → tracked nucleus labels.
+    _report(0, 4, "Nucleus masks...")
+    nuc_masks = native_masks.run_nucleus_masks_stack(
+        nucleus_stack, nucleus_params, cancel_cb=cancel_cb,
+    )
+    _report(1, 4, "Tracking nuclei...")
+    nuc_tracked = track_laptrack.track_axiswise(
+        nuc_masks, max_distance=max_distance, max_frame_gap=max_frame_gap,
+    ).astype(np.int32)
+
+    # 2. Cell flows + foreground.
+    _report(2, 4, "Cell flows...")
+    prob, dp = cellpose_runner.run_cell_stack(
+        cell_stack, cell_params, cancel_cb=cancel_cb,
+    )
+    foreground = cell_foreground_from_prob(prob, flow_params.fg_threshold)
+
+    # 3. Flow-follow per z-plane; cell labels inherit the tracked nucleus ids.
+    _report(3, 4, "Assigning cells to nuclei...")
+    T, Z = nuc_tracked.shape[:2]
+    cell_tracked = np.zeros_like(nuc_tracked, dtype=np.int32)
+    for z in range(Z):
+        # dp is (T, Z, 2, Y, X) → per-plane (T, 2, Y, X).
+        cell_tracked[:, z] = flow_follow_movie(
+            foreground[:, z],
+            dp[:, z],
+            nuc_tracked[:, z],
+            flow_params,
+        )
+    _report(4, 4, "Joint segmentation complete.")
+    return nuc_tracked, cell_tracked

itasc/cellpose/napari.yaml

@@ -0,0 +1,10 @@
+name: itasc-cellpose
+display_name: ITASC Cellpose
+contributions:
+  commands:
+    - id: itasc-cellpose.segment_track
+      title: "Cellpose Segment + Track"
+      python_name: itasc.napari.cellpose_segment_track_widget:make_cellpose_segment_track_widget
+  widgets:
+    - command: itasc-cellpose.segment_track
+      display_name: "Cellpose Segment + Track"

itasc/cellpose/native_masks.py

@@ -0,0 +1,313 @@
+"""Capture Cellpose-SAM *native* masks, Qt-free, for the standalone distro.
+
+The app's :mod:`itasc.cellpose.cellpose_runner` keeps only the probability /
+flow maps (``_, flows, _ = model.eval(...)``) because the integrated pipeline
+derives its labels from divergence maps + Ultrack. The independently-shipped
+``itasc-cellpose`` tool instead wants the labelled masks Cellpose computes and
+the runner throws away — that is index ``0`` of ``model.eval``.
+
+This module re-runs the same eval loop but keeps the masks, producing a canonical
+``(T, Z, Y, X)`` label stack per channel. It reuses the runner's model loading,
+gamma correction and cancellation so the two paths stay consistent; it does not
+touch the runner itself, so the app's behaviour is unchanged.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal
+from collections.abc import Callable
+
+import numpy as np
+
+from itasc.core.tiff import imwrite_grayscale
+from itasc.cellpose import cellpose_runner
+from itasc.cellpose.cellpose_runner import (
+    CellParams,
+    NucleusParams,
+    _apply_gamma,
+    _check_cancel,
+    _diameter_kwarg,
+    _NORMALIZE,
+)
+
+__all__ = [
+    "run_nucleus_masks_frame",
+    "run_cell_masks_frame",
+    "run_nucleus_masks_stack",
+    "run_cell_masks_stack",
+    "run_nucleus_maps_frame",
+    "iter_nucleus_maps_stack",
+    "write_masks",
+    "offset_slice_labels",
+]
+
+
+def _niter_kwarg(niter: int) -> int | None:
+    """``0`` means auto — pass ``None`` so Cellpose derives niter from the diameter."""
+    return None if int(niter) == 0 else int(niter)
+
+
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    """Map Cellpose's cellprob logits to a ``[0, 1]`` probability map."""
+    return 1.0 / (1.0 + np.exp(-np.asarray(x, dtype=np.float32)))
+
+
+def _flow_to_rgb(dp: np.ndarray) -> np.ndarray:
+    """HSV→RGB visualization of a ``(2, Y, X)`` flow field → ``(Y, X, 3)`` uint8.
+
+    Hue encodes flow direction (the colour wheel Cellpose users recognise) and
+    value encodes magnitude (per-frame normalised), so touching objects pulled
+    apart by their flows are visible. Saturation is fixed at 1.
+    """
+    dp = np.asarray(dp, dtype=np.float32)
+    dy, dx = dp[0], dp[1]
+    hue = (np.arctan2(dy, dx) + np.pi) / (2.0 * np.pi)  # [0, 1)
+    mag = np.hypot(dy, dx)
+    mmax = float(mag.max()) if mag.size else 0.0
+    val = mag / mmax if mmax > 0.0 else np.zeros_like(mag)
+    h6 = hue * 6.0
+    i = np.floor(h6).astype(np.int32) % 6
+    f = h6 - np.floor(h6)
+    p = np.zeros_like(val)         # saturation 1 → p = v*(1-s) = 0
+    q = val * (1.0 - f)
+    tt = val * f
+    conds = [i == 0, i == 1, i == 2, i == 3, i == 4, i == 5]
+    r = np.select(conds, [val, q, p, p, tt, val])
+    g = np.select(conds, [tt, val, val, q, p, p])
+    b = np.select(conds, [p, p, tt, val, val, q])
+    rgb = np.stack([r, g, b], axis=-1)
+    return (np.clip(rgb, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+
+def _eval_maps(img: np.ndarray, **eval_kwargs):
+    """Run ``model.eval`` once and keep masks **and** the maps it would discard.
+
+    Returns ``(masks int32, prob (sigmoid of cellprob) float32, dp float32)`` —
+    the labelled masks, the probability map and the raw flow field — from a single
+    forward pass, so no channel is segmented twice to surface its intermediates.
+    """
+    model = cellpose_runner.get_model()
+    masks, flows, _styles = model.eval(img, normalize=_NORMALIZE, **eval_kwargs)
+    dp = np.asarray(flows[1], dtype=np.float32)
+    cellprob = np.asarray(flows[2], dtype=np.float32)
+    return np.asarray(masks, dtype=np.int32), _sigmoid(cellprob), dp
+
+
+def offset_slice_labels(slices: list[np.ndarray]) -> np.ndarray:
+    """Stack per-z 2D label images into ``(Z, Y, X)`` with frame-unique labels.
+
+    Each 2D Cellpose run labels its objects from ``1`` independently, so naively
+    stacking would collide labels across z. We offset every slice's non-zero
+    labels by the running maximum so a frame's labels are globally unique; the
+    background ``0`` is preserved. For the common single-z case this is a no-op.
+    """
+    out: list[np.ndarray] = []
+    running_max = 0
+    for sl in slices:
+        sl = np.asarray(sl).astype(np.int32, copy=True)
+        nonzero = sl > 0
+        if nonzero.any():
+            sl[nonzero] += running_max
+            running_max = int(sl.max())
+        out.append(sl)
+    return np.stack(out, axis=0)
+
+
+def _eval_masks(img: np.ndarray, **eval_kwargs) -> np.ndarray:
+    """Run ``model.eval`` and return its native masks (index 0) as ``int32``."""
+    model = cellpose_runner.get_model()
+    masks, _flows, _styles = model.eval(img, normalize=_NORMALIZE, **eval_kwargs)
+    return np.asarray(masks, dtype=np.int32)
+
+
+def run_nucleus_masks_frame(
+    frame: np.ndarray,
+    z: int | None,
+    params: NucleusParams,
+) -> np.ndarray:
+    """Native nucleus masks for one frame.
+
+    ``z is None`` runs full 3D over ``(Z, Y, X)`` and returns labels ``(Z, Y, X)``;
+    an integer ``z`` runs 2D on ``frame[z]`` and returns labels ``(Y, X)``.
+    """
+    diameter = _diameter_kwarg(params.diameter)
+    niter = _niter_kwarg(params.niter)
+    if z is None:
+        volume = _apply_gamma(frame, params.gamma)
+        return _eval_masks(
+            volume,
+            do_3D=True,
+            z_axis=0,
+            diameter=diameter,
+            anisotropy=params.anisotropy,
+            min_size=params.min_size,
+            cellprob_threshold=params.cellprob_threshold,
+            flow_threshold=params.flow_threshold,
+            niter=niter,
+        )
+    slice_2d = _apply_gamma(frame[z], params.gamma)
+    return _eval_masks(
+        slice_2d,
+        diameter=diameter,
+        min_size=params.min_size,
+        cellprob_threshold=params.cellprob_threshold,
+        flow_threshold=params.flow_threshold,
+        niter=niter,
+    )
+
+
+def run_nucleus_maps_frame(
+    frame: np.ndarray,
+    z: int,
+    params: NucleusParams,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Native masks **+ sigmoid prob map + RGB flow** for one 2D z-slice.
+
+    A single ``model.eval`` pass yields ``(masks (Y, X) int32, prob (Y, X)
+    float32, flow_rgb (Y, X, 3) uint8)`` — the labels plus the two intermediates
+    (probability and flow) the standalone tool surfaces so Channel 1 can be tuned.
+    The standalone anchor is always per-plane 2D, so ``z`` is an integer slice.
+    """
+    diameter = _diameter_kwarg(params.diameter)
+    slice_2d = _apply_gamma(frame[z], params.gamma)
+    masks, prob, dp = _eval_maps(
+        slice_2d,
+        diameter=diameter,
+        min_size=params.min_size,
+        cellprob_threshold=params.cellprob_threshold,
+        flow_threshold=params.flow_threshold,
+        niter=_niter_kwarg(params.niter),
+    )
+    return masks, prob, _flow_to_rgb(dp)
+
+
+def iter_nucleus_maps_stack(
+    stack: np.ndarray,
+    params: NucleusParams,
+    *,
+    progress_cb: Callable[[int, int, str], None] | None = None,
+    cancel_cb: Callable[[], bool] | None = None,
+):
+    """Yield per-frame ``(t, masks, prob, flow_rgb)`` so results stream into a viewer.
+
+    For a ``(T, Z, Y, X)`` input each yield carries one time-frame: ``masks`` and
+    ``prob`` are ``(Z, Y, X)`` and ``flow_rgb`` is ``(Z, Y, X, 3)`` uint8, with
+    per-frame-unique mask labels (:func:`offset_slice_labels`). Always per-plane 2D
+    (the standalone anchor never runs do_3d). A caller can drop each frame into a
+    napari layer as it arrives instead of waiting for the whole stack.
+    """
+    stack = np.asarray(stack)
+    if stack.ndim != 4:
+        raise ValueError(f"expected (T, Z, Y, X), got shape {stack.shape}")
+    T, Z = stack.shape[:2]
+    for t in range(T):
+        _check_cancel(cancel_cb)
+        mask_slices: list[np.ndarray] = []
+        prob_slices: list[np.ndarray] = []
+        flow_slices: list[np.ndarray] = []
+        for z in range(Z):
+            if progress_cb is not None:
+                progress_cb(t, T, f"Channel 1 maps: frame {t + 1}/{T}, z {z + 1}/{Z}...")
+            m, p, f = run_nucleus_maps_frame(stack[t], z=z, params=params)
+            mask_slices.append(m)
+            prob_slices.append(p)
+            flow_slices.append(f)
+        masks = offset_slice_labels(mask_slices)
+        prob = np.stack(prob_slices, axis=0).astype(np.float32)
+        flow = np.stack(flow_slices, axis=0)
+        if progress_cb is not None:
+            progress_cb(t + 1, T, f"Channel 1 maps: frame {t + 1}/{T}...")
+        yield t, masks, prob, flow
+
+
+def run_cell_masks_frame(
+    frame: np.ndarray,
+    z: int,
+    params: CellParams,
+) -> np.ndarray:
+    """Native cell masks for a single 2D z-slice. Returns labels ``(Y, X)``."""
+    diameter = _diameter_kwarg(params.diameter)
+    slice_2d = _apply_gamma(frame[z], params.gamma)
+    return _eval_masks(slice_2d, diameter=diameter, min_size=params.min_size)
+
+
+def run_nucleus_masks_stack(
+    stack: np.ndarray,
+    params: NucleusParams,
+    *,
+    progress_cb: Callable[[int, int, str], None] | None = None,
+    cancel_cb: Callable[[], bool] | None = None,
+) -> np.ndarray:
+    """Process a ``(T, Z, Y, X)`` stack into a ``(T, Z, Y, X)`` label stack.
+
+    ``params.do_3d`` runs a coherent 3D segmentation per frame; otherwise each
+    z-slice is segmented in 2D and the slices are stacked with frame-unique
+    labels (see :func:`offset_slice_labels`).
+    """
+    if stack.ndim != 4:
+        raise ValueError(f"expected (T, Z, Y, X), got shape {stack.shape}")
+    T, Z = stack.shape[:2]
+    frames: list[np.ndarray] = []
+    for t in range(T):
+        _check_cancel(cancel_cb)
+        if progress_cb is not None:
+            progress_cb(t, T, f"Nucleus masks: frame {t + 1}/{T}...")
+        if params.do_3d:
+            masks = run_nucleus_masks_frame(stack[t], z=None, params=params)
+        else:
+            slices = []
+            for z in range(Z):
+                if progress_cb is not None:
+                    progress_cb(
+                        t, T,
+                        f"Nucleus masks: frame {t + 1}/{T}, z {z + 1}/{Z}...",
+                    )
+                slices.append(run_nucleus_masks_frame(stack[t], z=z, params=params))
+            masks = offset_slice_labels(slices)
+        frames.append(masks.astype(np.int32))
+        if progress_cb is not None:
+            progress_cb(t + 1, T, f"Nucleus masks: frame {t + 1}/{T}...")
+    return np.stack(frames, axis=0)
+
+
+def run_cell_masks_stack(
+    stack: np.ndarray,
+    params: CellParams,
+    *,
+    progress_cb: Callable[[int, int, str], None] | None = None,
+    cancel_cb: Callable[[], bool] | None = None,
+) -> np.ndarray:
+    """Process a ``(T, Z, Y, X)`` stack slice-by-slice into ``(T, Z, Y, X)`` labels."""
+    if stack.ndim != 4:
+        raise ValueError(f"expected (T, Z, Y, X), got shape {stack.shape}")
+    T, Z = stack.shape[:2]
+    frames: list[np.ndarray] = []
+    for t in range(T):
+        _check_cancel(cancel_cb)
+        if progress_cb is not None:
+            progress_cb(t, T, f"Cell masks: frame {t + 1}/{T}...")
+        slices = [run_cell_masks_frame(stack[t], z=z, params=params) for z in range(Z)]
+        frames.append(offset_slice_labels(slices))
+        if progress_cb is not None:
+            progress_cb(t + 1, T, f"Cell masks: frame {t + 1}/{T}...")
+    return np.stack(frames, axis=0)
+
+
+def write_masks(
+    masks_tzyx: np.ndarray,
+    output_dir: Path,
+    channel: Literal["nucleus", "cell"],
+) -> Path:
+    """Write ``{channel}_masks.tif`` (``int32``, axes ``TZYX``); return its path."""
+    if channel not in ("nucleus", "cell"):
+        raise ValueError(f"channel must be 'nucleus' or 'cell', got {channel!r}")
+    if masks_tzyx.ndim != 4:
+        raise ValueError(f"masks must be (T, Z, Y, X), got {masks_tzyx.shape}")
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    path = output_dir / f"{channel}_masks.tif"
+    imwrite_grayscale(
+        path, masks_tzyx.astype(np.int32),
+        compression="zlib", metadata={"axes": "TZYX"},
+    )
+    return path

itasc/cellpose/retrack.py

@@ -0,0 +1,232 @@
+"""DB-free greedy retracker for the standalone segment + track corrector.
+
+This mirrors the constrained retracker in :mod:`itasc.tracking_ultrack`
+(which ships in a separate ``itasc-tracking`` distribution the standalone
+``itasc-cellpose`` tool does *not* depend on), reduced to the validation-free
+case the standalone needs: starting from the current frame, every later frame's
+labels are re-linked to the already-retracked neighbour toward the start frame
+by the same additive similarity score (area ratio + centroid-corrected IoU -
+distance), matched greedily best-first. Unmatched target cells receive fresh
+ids. There is no locked/validated/reserved concept here — the standalone has no
+validation store — so the algorithm is purely the geometric re-linking.
+
+Everything here is Qt-free and depends only on numpy/scipy/skimage, so it ships
+inside the cellpose distro tree and is unit-testable without a viewer.
+"""
+from __future__ import annotations
+
+from typing import Literal
+
+import numpy as np
+from scipy.spatial.distance import cdist
+from skimage.measure import regionprops
+
+
+# ---------------------------------------------------------------------------
+# Similarity scoring (ported from itasc.tracking_ultrack.scoring so the
+# standalone stays self-contained — same formula, so the UX matches the app).
+# ---------------------------------------------------------------------------
+def _rasterize(coords: np.ndarray, mins: np.ndarray, shape: tuple) -> np.ndarray:
+    canvas = np.zeros(shape, dtype=bool)
+    idx = np.rint(coords - mins).astype(int)
+    valid = np.ones(len(idx), dtype=bool)
+    for axis, size in enumerate(shape):
+        valid &= (idx[:, axis] >= 0) & (idx[:, axis] < size)
+    idx = idx[valid]
+    if idx.size:
+        canvas[tuple(idx.T)] = True
+    return canvas
+
+
+def centroid_corrected_iou_from_coords(
+    src_coords: np.ndarray,
+    src_centroid: np.ndarray,
+    target_coords: np.ndarray,
+    target_centroid: np.ndarray,
+) -> float:
+    """IoU after shifting the target so its centroid matches the source's."""
+    if len(src_coords) == 0 or len(target_coords) == 0:
+        return 0.0
+    shifted = target_coords + (src_centroid - target_centroid)
+    all_coords = np.vstack([src_coords, shifted])
+    mins = np.floor(all_coords.min(axis=0)).astype(int) - 1
+    maxs = np.ceil(all_coords.max(axis=0)).astype(int) + 1
+    shape = tuple((maxs - mins + 1).tolist())
+    if any(dim <= 0 for dim in shape):
+        return 0.0
+    mins_f = mins.astype(np.float32)
+    src_canvas = _rasterize(src_coords, mins_f, shape)
+    tgt_canvas = _rasterize(shifted, mins_f, shape)
+    union = np.logical_or(src_canvas, tgt_canvas).sum()
+    if union == 0:
+        return 0.0
+    return float(np.logical_and(src_canvas, tgt_canvas).sum() / union)
+
+
+def similarity_score(
+    *,
+    area_ratio: float,
+    centroid_corrected_iou: float,
+    distance: float,
+    area_weight: float,
+    iou_weight: float,
+    distance_weight: float,
+) -> float:
+    """Additive similarity score (higher = more preferred).
+
+    Shape terms are positive rewards in ``[0, 1]``; distance is a raw pixel
+    penalty, so the score can go negative for far candidates.
+    """
+    return (
+        area_weight * area_ratio
+        + iou_weight * centroid_corrected_iou
+        - distance_weight * distance
+    )
+
+
+# ---------------------------------------------------------------------------
+# Frame / stack retracking
+# ---------------------------------------------------------------------------
+def _label_props(
+    labels: np.ndarray,
+) -> dict[int, tuple[np.ndarray, float, np.ndarray]]:
+    """Return ``{label_id: (centroid_yx, area, coords)}`` for every non-zero id."""
+    props: dict[int, tuple[np.ndarray, float, np.ndarray]] = {}
+    for region in regionprops(labels):
+        props[int(region.label)] = (
+            np.asarray(region.centroid, dtype=np.float32),
+            float(region.area),
+            region.coords.astype(np.float32),
+        )
+    return props
+
+
+def retrack_frame(
+    ref_labels: np.ndarray,
+    target_labels: np.ndarray,
+    *,
+    max_dist_px: float = 50.0,
+    area_weight: float = 1.0,
+    iou_weight: float = 1.0,
+    distance_weight: float = 0.05,
+) -> np.ndarray:
+    """Remap every target id to its best-matching reference id by similarity.
+
+    Each target cell takes its highest-scoring still-free reference cell
+    (greedy best-first), among reference cells whose centroid is within
+    ``max_dist_px``. Targets with no eligible reference get a fresh id above all
+    existing ids, so no two cells ever collide. Returns a relabelled copy of
+    ``target_labels`` (input arrays are not mutated).
+    """
+    result = np.zeros_like(target_labels)
+
+    tgt_ids = [int(i) for i in np.unique(target_labels) if i != 0]
+    if not tgt_ids:
+        return result
+
+    ref_ids = [int(i) for i in np.unique(ref_labels) if i != 0]
+    max_existing = max(
+        int(ref_labels.max()) if ref_labels.size and ref_labels.max() > 0 else 0,
+        int(target_labels.max()) if target_labels.size and target_labels.max() > 0 else 0,
+    )
+
+    def _assign_fresh(remap: dict[int, int]) -> dict[int, int]:
+        next_id = max_existing + 1
+        for tid in tgt_ids:
+            if tid not in remap:
+                remap[tid] = next_id
+                next_id += 1
+        return remap
+
+    if not ref_ids:
+        for tid, new_id in _assign_fresh({}).items():
+            result[target_labels == tid] = new_id
+        return result
+
+    tgt_props = _label_props(np.where(np.isin(target_labels, tgt_ids), target_labels, 0))
+    ref_props = _label_props(np.where(np.isin(ref_labels, ref_ids), ref_labels, 0))
+
+    tgt_centroids = np.array([tgt_props[t][0] for t in tgt_ids])
+    ref_centroids = np.array([ref_props[r][0] for r in ref_ids])
+    dist = cdist(tgt_centroids, ref_centroids)
+    gate = dist <= max_dist_px
+
+    # Score every distance-eligible (target, reference) pair.
+    scored: list[tuple[float, float, int, int]] = []
+    for ti in range(dist.shape[0]):
+        _t_centroid, t_area, t_coords = tgt_props[tgt_ids[ti]]
+        for ri in np.nonzero(gate[ti])[0]:
+            r_centroid, r_area, r_coords = ref_props[ref_ids[ri]]
+            area_ratio = min(t_area, r_area) / max(t_area, r_area)
+            iou = centroid_corrected_iou_from_coords(
+                r_coords, r_centroid, t_coords, tgt_props[tgt_ids[ti]][0]
+            )
+            score = similarity_score(
+                area_ratio=area_ratio,
+                centroid_corrected_iou=iou,
+                distance=float(dist[ti, ri]),
+                area_weight=area_weight,
+                iou_weight=iou_weight,
+                distance_weight=distance_weight,
+            )
+            scored.append((score, float(dist[ti, ri]), ti, int(ri)))
+
+    # Greedy best-first: highest score wins, nearer centroid breaks ties. Each
+    # target and reference is consumed once.
+    scored.sort(key=lambda item: (-item[0], item[1]))
+
+    remap: dict[int, int] = {}
+    used_refs: set[int] = set()
+    for _score, _d, ti, ri in scored:
+        tid = tgt_ids[ti]
+        if tid in remap or ri in used_refs:
+            continue
+        remap[tid] = ref_ids[ri]
+        used_refs.add(ri)
+
+    for tid, new_id in _assign_fresh(remap).items():
+        result[target_labels == tid] = new_id
+    return result
+
+
+def retrack_stack(
+    stack: np.ndarray,
+    *,
+    start_frame: int,
+    direction: Literal["forward", "backward"],
+    max_dist_px: float = 50.0,
+    area_weight: float = 1.0,
+    iou_weight: float = 1.0,
+    distance_weight: float = 0.05,
+) -> np.ndarray:
+    """Retrack a time-first ``(T, Y, X)`` stack outward from ``start_frame``.
+
+    The start frame is kept as the anchor; each later frame in ``direction`` is
+    re-linked (:func:`retrack_frame`) to the already-retracked neighbour toward
+    the start frame, so corrected ids propagate. Returns a new stack; the input
+    is not mutated.
+    """
+    stack = np.asarray(stack)
+    if stack.ndim != 3 or stack.shape[0] < 2:
+        raise ValueError("retrack_stack needs a 3D time-first stack with >= 2 frames.")
+
+    out = stack.copy()
+    if direction == "forward":
+        frame_range = range(start_frame + 1, out.shape[0])
+        neighbour = lambda t: out[t - 1]
+    elif direction == "backward":
+        frame_range = range(start_frame - 1, -1, -1)
+        neighbour = lambda t: out[t + 1]
+    else:
+        raise ValueError(f"unknown retrack direction: {direction!r}")
+
+    for t in frame_range:
+        out[t] = retrack_frame(
+            neighbour(t),
+            out[t],
+            max_dist_px=max_dist_px,
+            area_weight=area_weight,
+            iou_weight=iou_weight,
+            distance_weight=distance_weight,
+        )
+    return out