itasc 0.2.0.dev0__py3-none-any.whl

itasc/tracking_ultrack/reseed.py

@@ -0,0 +1,456 @@
+"""Re-solve with validated tracks using annotated Ultrack seed nodes.
+
+The validate-and-resolve loop lets the user lock in confirmed cell tracks
+(validated_tracks), then re-solve the rest of the stack — possibly with
+tweaked parameters — without losing the validated work.
+
+The contract for `validated_tracks` is `dict[int, set[int]]` where keys are
+cell IDs and values are the set of frames at which that cell is validated.
+This is independent of the on-disk JSON shape (which is managed separately by
+`itasc.tracking_ultrack.validation_state`).
+
+`prune_validated_overlaps` and `_build_frame_masks` are kept for existing
+unit tests and ad hoc use.  `merge_validated_into_export` is called at the
+end of `resolve_with_canonical_segment` to paste validated IDs back.
+"""
+from __future__ import annotations
+
+import logging
+import pickle
+from pathlib import Path
+
+import numpy as np
+
+
+LOG = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helper: build per-frame validated-mask index
+# ---------------------------------------------------------------------------
+
+def _build_frame_masks(
+    validated_tracks: dict[int, set[int]],
+    tracked_labels: np.ndarray,
+) -> dict[int, list[tuple[int, np.ndarray, tuple[int, int, int, int]]]]:
+    """Build a per-frame index of validated masks.
+
+    Returns
+    -------
+    dict mapping frame index ``t`` to a list of
+    ``(cell_id, bool_mask_crop, (y0, x0, y1, x1))`` tuples, one per validated
+    cell at that frame.  Both ``bool_mask_crop`` and the bbox refer to the
+    cropped region (exclusive-max convention).
+
+    Only frames where at least one cell is validated are included.
+    The 2D spatial slice is taken from ``tracked_labels[t]`` (shape Y×X).
+    For a 4-D input ``(T, Z, Y, X)`` the z-axis is max-projected to give a
+    2D footprint for intersection tests (conservative: any z-slice counts).
+    """
+    frame_index: dict[int, list[tuple[int, np.ndarray, tuple[int, int, int, int]]]] = {}
+
+    for cell_id, frames in validated_tracks.items():
+        for t in frames:
+            t = int(t)
+            frame_vol = tracked_labels[t]  # (Y, X) or (Z, Y, X)
+
+            # Project to 2D for intersection (max over Z if 3D)
+            if frame_vol.ndim == 3:
+                frame_2d = (frame_vol == cell_id).any(axis=0)
+            else:
+                frame_2d = (frame_vol == cell_id)
+
+            if not frame_2d.any():
+                # Cell absent at this frame — skip silently
+                continue
+
+            # Tight bbox around the mask
+            rows = np.where(frame_2d.any(axis=1))[0]
+            cols = np.where(frame_2d.any(axis=0))[0]
+            y0, y1 = int(rows[0]), int(rows[-1]) + 1
+            x0, x1 = int(cols[0]), int(cols[-1]) + 1
+            mask_crop = frame_2d[y0:y1, x0:x1]
+
+            if t not in frame_index:
+                frame_index[t] = []
+            frame_index[t].append((cell_id, mask_crop, (y0, x0, y1, x1)))
+
+    return frame_index
+
+
+# ---------------------------------------------------------------------------
+# 1. prune_validated_overlaps
+# ---------------------------------------------------------------------------
+
+def prune_validated_overlaps(
+    working_dir: str | Path,
+    validated_tracks: dict[int, set[int]],
+    tracked_labels: np.ndarray,
+) -> int:
+    """Delete NodeDB rows that overlap any validated mask and cascade to OverlapDB.
+
+    Parameters
+    ----------
+    working_dir:
+        Ultrack working directory containing ``data.db``.
+    validated_tracks:
+        ``{cell_id: {frames}}`` — which frames each validated cell spans.
+    tracked_labels:
+        Current corrected labelmap, shape ``(T, Y, X)`` or ``(T, Z, Y, X)``.
+
+    Returns
+    -------
+    int
+        Number of NodeDB rows deleted.
+
+    Notes
+    -----
+    Conflict criterion: **any pixel intersection** between a candidate node's
+    2D footprint and any validated cell mask at the same frame counts as a
+    conflict.  The node (and all OverlapDB rows referencing it) is deleted.
+
+    The 3D node masks stored in NodeDB are stored as ``(1, h, w)`` crops
+    (the ingest convention in ``ingest.py``).  We squeeze the Z=1 dimension
+    for the 2D intersection test.  True 3D (Z > 1) is currently not supported;
+    the bbox test gracefully falls back to a bounding-box-only check in that
+    case.
+
+    Both NodeDB and referencing OverlapDB rows are deleted in a single
+    SQLite transaction to preserve DB consistency.
+    """
+    import sqlalchemy as sqla
+    from sqlalchemy.orm import Session
+    from ultrack.core.database import NodeDB, OverlapDB
+
+    if not validated_tracks:
+        return 0
+
+    working_dir = Path(working_dir)
+    db_url = f"sqlite:///{working_dir / 'data.db'}"
+    engine = sqla.create_engine(db_url)
+
+    frame_masks = _build_frame_masks(validated_tracks, tracked_labels)
+    if not frame_masks:
+        return 0
+
+    total_deleted = 0
+
+    with Session(engine) as session:
+        for t, val_cells in frame_masks.items():
+            # Query all nodes at this frame (id + pickle blob)
+            rows = (
+                session.query(NodeDB.id, NodeDB.pickle)
+                .where(NodeDB.t == t)
+                .all()
+            )
+            if not rows:
+                continue
+
+            conflict_ids: list[int] = []
+
+            for node_id_val, node_blob in rows:
+                # Deserialize; NodeDB.pickle is a MaybePickleType that
+                # returns a Node object on read.
+                if isinstance(node_blob, (bytes, memoryview)):
+                    try:
+                        node = pickle.loads(bytes(node_blob))
+                    except Exception:
+                        LOG.warning(
+                            "Could not unpickle node id=%s at t=%s — skipping",
+                            node_id_val, t,
+                        )
+                        continue
+                else:
+                    # SQLAlchemy MaybePickleType already deserialized it
+                    node = node_blob
+
+                bbox = np.asarray(node.bbox)
+                ndim = len(bbox) // 2
+                if ndim == 3:
+                    # bbox = [z0, y0, x0, z1, y1, x1]
+                    ny0, nx0, ny1, nx1 = int(bbox[1]), int(bbox[2]), int(bbox[4]), int(bbox[5])
+                    node_mask = getattr(node, "mask", None)
+                    if node_mask is not None:
+                        node_mask = np.asarray(node_mask, dtype=bool)
+                        if node_mask.shape[0] == 1:
+                            node_mask_2d = node_mask[0]  # (h, w)
+                        else:
+                            # Multi-Z: project
+                            node_mask_2d = node_mask.any(axis=0)
+                    else:
+                        node_mask_2d = None
+                elif ndim == 2:
+                    # bbox = [y0, x0, y1, x1]
+                    ny0, nx0, ny1, nx1 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
+                    node_mask = getattr(node, "mask", None)
+                    if node_mask is not None:
+                        node_mask_2d = np.asarray(node_mask, dtype=bool)
+                        if node_mask_2d.ndim == 3 and node_mask_2d.shape[0] == 1:
+                            node_mask_2d = node_mask_2d[0]
+                    else:
+                        node_mask_2d = None
+                else:
+                    LOG.warning(
+                        "Unexpected bbox ndim=%s for node id=%s — skipping",
+                        ndim, node_id_val,
+                    )
+                    continue
+
+                # Check against every validated cell at this frame
+                is_conflict = False
+                for _cell_id, val_crop, (vy0, vx0, vy1, vx1) in val_cells:
+                    # Fast bbox intersection test first
+                    if ny1 <= vy0 or vy1 <= ny0 or nx1 <= vx0 or vx1 <= nx0:
+                        continue  # bboxes don't overlap → no conflict
+
+                    # Compute the overlap region in image coordinates
+                    oy0, oy1 = max(ny0, vy0), min(ny1, vy1)
+                    ox0, ox1 = max(nx0, vx0), min(nx1, vx1)
+
+                    if node_mask_2d is None:
+                        # No mask data — treat bbox overlap as conflict
+                        is_conflict = True
+                        break
+
+                    # Crop both masks to the overlap region
+                    nm_crop = node_mask_2d[
+                        oy0 - ny0: oy1 - ny0,
+                        ox0 - nx0: ox1 - nx0,
+                    ]
+                    vm_crop = val_crop[
+                        oy0 - vy0: oy1 - vy0,
+                        ox0 - vx0: ox1 - vx0,
+                    ]
+                    if nm_crop.shape != vm_crop.shape:
+                        # Shape mismatch due to off-by-one — treat as conflict
+                        is_conflict = True
+                        break
+                    if np.any(nm_crop & vm_crop):
+                        is_conflict = True
+                        break
+
+                if is_conflict:
+                    conflict_ids.append(int(node_id_val))
+
+            if not conflict_ids:
+                continue
+
+            # Delete OverlapDB rows referencing any conflict node, then nodes —
+            # all in one transaction (the session is autocommit=False by default).
+            (
+                session.query(OverlapDB)
+                .where(
+                    sqla.or_(
+                        OverlapDB.node_id.in_(conflict_ids),
+                        OverlapDB.ancestor_id.in_(conflict_ids),
+                    )
+                )
+                .delete(synchronize_session=False)
+            )
+            (
+                session.query(NodeDB)
+                .where(NodeDB.id.in_(conflict_ids))
+                .delete(synchronize_session=False)
+            )
+            total_deleted += len(conflict_ids)
+
+        session.commit()
+
+    LOG.info("prune_validated_overlaps: deleted %d NodeDB rows", total_deleted)
+    return total_deleted
+
+
+# ---------------------------------------------------------------------------
+# 2. merge_validated_into_export
+# ---------------------------------------------------------------------------
+
+def merge_validated_into_export(
+    exported_labels: np.ndarray,
+    validated_tracks: dict[int, set[int]],
+    tracked_labels: np.ndarray,
+) -> tuple[np.ndarray, dict[int, int]]:
+    """Paste validated masks back onto the solver's exported labelmap.
+
+    Validated ``cell_id`` values are reserved and pasted back with their
+    original IDs so the validation store remains stable across resolve runs.
+    If the solver independently used a reserved validated ID elsewhere, those
+    solver pixels are moved to a fresh ID before the validated mask is pasted.
+    Validated pixels overwrite any conflicting content in the exported labelmap.
+
+    Parameters
+    ----------
+    exported_labels:
+        Labelmap produced by Ultrack export, shape ``(T, Y, X)`` or
+        ``(T, Z, Y, X)``.  Modified **in place** and returned.
+    validated_tracks:
+        ``{cell_id: {frames}}`` from the validation store.
+    tracked_labels:
+        Current corrected labelmap (same spatial shape as ``exported_labels``).
+        The source of ground-truth masks.
+
+    Returns
+    -------
+    tuple[np.ndarray, dict[int, int]]
+        ``(merged_labelmap, id_map)`` where ``id_map`` maps original validated
+        IDs to changed output IDs. The normal path preserves validated IDs, so
+        ``id_map`` is usually empty. The labelmap is the same array as
+        ``exported_labels``, modified in place.
+    """
+    if not validated_tracks:
+        return exported_labels, {}
+
+    id_map: dict[int, int] = {}
+    reserved_ids = {int(cell_id) for cell_id, frames in validated_tracks.items() if frames}
+    used_fresh_ids = set(int(v) for v in np.unique(exported_labels))
+
+    def _next_fresh_id() -> int:
+        next_id = int(exported_labels.max()) + 1
+        while next_id in reserved_ids or next_id in used_fresh_ids:
+            next_id += 1
+        used_fresh_ids.add(next_id)
+        return next_id
+
+    # Build validated_masks: per cell_id, list of (t, mask) pairs present in tracked_labels
+    validated_masks: dict[int, list[tuple[int, np.ndarray]]] = {}
+    for cell_id, frames in validated_tracks.items():
+        cell_id = int(cell_id)
+        if not frames:
+            continue
+        present: list[tuple[int, np.ndarray]] = []
+        for t in sorted(frames):
+            t = int(t)
+            if t < 0 or t >= tracked_labels.shape[0] or t >= exported_labels.shape[0]:
+                continue
+            frame_src = tracked_labels[t]
+            mask = np.asarray(frame_src == cell_id)
+            if mask.any():
+                present.append((t, mask))
+        if present:
+            validated_masks[cell_id] = present
+
+    # Build solver_track_remap: solver_id -> validated cell_id
+    # Inspect each validated frame to find the dominant solver track covering that mask region,
+    # then propagate that solver track's ID to the validated cell_id everywhere in the export.
+    solver_track_remap: dict[int, int] = {}
+    for cell_id in sorted(validated_masks.keys()):
+        present_masks = validated_masks[cell_id]
+        dominant_solver_id: int | None = None
+        best_count = 0
+        for t, mask in present_masks:
+            if exported_labels.ndim == 4:
+                # (T, Z, Y, X): project mask to 3D for indexing
+                frame_exp = exported_labels[t]
+                if mask.ndim == 2:
+                    solver_pixels = frame_exp[:, mask]
+                elif mask.ndim == 3:
+                    solver_pixels = frame_exp[mask]
+                else:
+                    continue
+            else:
+                frame_exp = exported_labels[t]
+                if mask.ndim == frame_exp.ndim:
+                    solver_pixels = frame_exp[mask]
+                elif mask.ndim == 2 and frame_exp.ndim == 3:
+                    solver_pixels = frame_exp[:, mask]
+                elif mask.ndim == 3 and frame_exp.ndim == 2:
+                    solver_pixels = frame_exp[mask.any(axis=0)]
+                else:
+                    continue
+            nonzero = solver_pixels[solver_pixels != 0]
+            if nonzero.size == 0:
+                continue
+            unique_ids, counts = np.unique(nonzero, return_counts=True)
+            t_best_idx = int(np.argmax(counts))
+            t_best_id = int(unique_ids[t_best_idx])
+            t_best_count = int(counts[t_best_idx])
+            if t_best_count > best_count:
+                best_count = t_best_count
+                dominant_solver_id = t_best_id
+
+        if dominant_solver_id is None or dominant_solver_id == 0:
+            continue
+        if dominant_solver_id == cell_id:
+            # Solver already used the correct ID — no remap needed, but still handle
+            # collision below where the solver used cell_id for unrelated pixels
+            pass
+        elif dominant_solver_id in solver_track_remap:
+            LOG.warning(
+                "Solver track %d claimed by multiple validated cells; "
+                "keeping first claim (cell %d), skipping cell %d",
+                dominant_solver_id,
+                solver_track_remap[dominant_solver_id],
+                cell_id,
+            )
+            continue
+        else:
+            solver_track_remap[dominant_solver_id] = cell_id
+
+    # Resolve collisions: if any existing pixels have value == target cell_id
+    # but DON'T belong to the source solver track, move them to a fresh ID.
+    # This is the same as the original collision handling, now generalized to
+    # cover both the remap targets and direct-match cases.
+    for solver_id, target_cell_id in list(solver_track_remap.items()):
+        if solver_id == target_cell_id:
+            continue
+        # Pixels currently carrying target_cell_id that are NOT from this solver track
+        collision_mask = np.asarray(exported_labels == target_cell_id)
+        # Remove pixels belonging to the source solver track from collision set
+        src_mask = np.asarray(exported_labels == solver_id)
+        # Exclude src_mask locations — those will become target_cell_id after remap anyway
+        collision_mask = collision_mask & ~src_mask
+        if collision_mask.any():
+            exported_labels[collision_mask] = _next_fresh_id()
+
+    # Handle cells that are not remapped (dominant_solver_id == cell_id or no dominant found)
+    # — original collision logic for reserved IDs used by solver for unrelated pixels.
+    remapped_sources = set(solver_track_remap.keys())
+    remapped_targets = set(solver_track_remap.values())
+    for cell_id, present_masks in validated_masks.items():
+        if cell_id in remapped_targets and cell_id not in remapped_sources:
+            # Already handled in collision resolution above
+            continue
+        if cell_id in remapped_sources:
+            # This cell IS the source solver track being remapped — no separate collision needed
+            continue
+        # No remap found: solver either didn't have this cell or already used correct ID.
+        # Run original collision resolution for the reserved cell_id.
+        solver_collision = np.asarray(exported_labels == cell_id)
+        for t, mask in present_masks:
+            frame_collision = solver_collision[t]
+            if mask.ndim == frame_collision.ndim:
+                frame_collision[mask] = False
+            elif mask.ndim == 2 and frame_collision.ndim == 3:
+                frame_collision[:, mask] = False
+            elif mask.ndim == 3 and frame_collision.ndim == 2:
+                frame_collision[mask.any(axis=0)] = False
+            else:
+                raise ValueError(
+                    "Validated mask shape is incompatible with exported labels: "
+                    f"mask={mask.shape}, exported={frame_collision.shape}"
+                )
+        if solver_collision.any():
+            exported_labels[solver_collision] = _next_fresh_id()
+
+    # Apply solver track remap: rename solver track IDs to validated cell IDs
+    for solver_id, target_cell_id in solver_track_remap.items():
+        if solver_id == target_cell_id:
+            continue
+        exported_labels[exported_labels == solver_id] = target_cell_id
+
+    # Paste validated masks (overrides geometry at validated frames)
+    for cell_id, present_masks in validated_masks.items():
+        for t, mask in present_masks:
+            frame_out = exported_labels[t]
+            if mask.ndim == frame_out.ndim:
+                frame_out[mask] = cell_id
+            elif mask.ndim == 2 and frame_out.ndim == 3:
+                frame_out[:, mask] = cell_id
+            elif mask.ndim == 3 and frame_out.ndim == 2:
+                frame_out[mask.any(axis=0)] = cell_id
+            else:
+                raise ValueError(
+                    "Validated mask shape is incompatible with exported labels: "
+                    f"mask={mask.shape}, exported={frame_out.shape}"
+                )
+
+    return exported_labels, id_map

itasc/tracking_ultrack/retracker.py

@@ -0,0 +1,153 @@
+"""Constrained greedy retracker using the shared linker similarity score.
+
+Each unlocked target cell is matched to a reference cell by the same
+``scoring.similarity_score`` (area ratio + centroid-corrected IoU - distance)
+used by the linker and the Extend tool, gated only by centroid distance.
+
+Matching is greedy best-first rather than a global minimum-cost assignment: a
+target always takes its highest-scoring still-free reference. A global
+``linear_sum_assignment`` minimises the *total* cost and will trade a cell's
+obvious best match away to lower another cell's cost, which in practice handed
+cells a far, worse reference; best-first never does that.
+"""
+from __future__ import annotations
+
+import numpy as np
+from scipy.spatial.distance import cdist
+from skimage.measure import regionprops
+
+from itasc.tracking_ultrack import scoring
+
+
+def _label_props(
+    labels: np.ndarray, keep_ids: list[int]
+) -> dict[int, tuple[np.ndarray, float, np.ndarray]]:
+    """Return {label_id: (centroid_yx, area, coords)} for labels in *keep_ids*.
+
+    Labels outside *keep_ids* (locked / reserved) are masked out before
+    ``regionprops`` so they cost no centroid, area, or coordinate extraction.
+    """
+    if not keep_ids:
+        return {}
+    keep = np.where(np.isin(labels, keep_ids), labels, 0)
+    props: dict[int, tuple[np.ndarray, float, np.ndarray]] = {}
+    for region in regionprops(keep):
+        props[int(region.label)] = (
+            np.asarray(region.centroid, dtype=np.float32),
+            float(region.area),
+            region.coords.astype(np.float32),
+        )
+    return props
+
+
+def retrack_frame_constrained(
+    ref_labels: np.ndarray,
+    target_labels: np.ndarray,
+    locked_target_ids: set[int],
+    max_dist_px: float = 50.0,
+    reserved_ids: set[int] | None = None,
+    *,
+    area_weight: float = 1.0,
+    iou_weight: float = 1.0,
+    distance_weight: float = 0.05,
+) -> np.ndarray:
+    """Remap target IDs by linker similarity without changing locked targets.
+
+    Target cells whose ID is in ``locked_target_ids`` keep their existing IDs.
+    Those IDs, plus any ``reserved_ids``, are protected from assignment to
+    unlocked target cells even if a matching reference cell exists.
+
+    Unlocked target cells are matched to available reference cells greedily in
+    descending ``similarity_score`` order: each target takes its best still-free
+    reference. A pair is only eligible when its centroid distance is
+    ``<= max_dist_px``; area ratio and centroid-corrected IoU contribute as soft
+    score terms (no hard gate), matching the Extend tool's behaviour.
+    """
+    locked_target_ids = set(locked_target_ids)
+    reserved_ids = set(reserved_ids or set())
+    blocked_ids = locked_target_ids | reserved_ids
+
+    result = np.zeros_like(target_labels)
+
+    tgt_ids = {int(i) for i in np.unique(target_labels) if i != 0}
+    for lid in locked_target_ids:
+        if lid in tgt_ids:
+            result[target_labels == lid] = lid
+
+    unlocked_tgt_ids = [tid for tid in tgt_ids if tid not in locked_target_ids]
+    if not unlocked_tgt_ids:
+        return result
+
+    ref_ids = {int(i) for i in np.unique(ref_labels) if i != 0}
+    available_ref_ids = [rid for rid in ref_ids if rid not in blocked_ids]
+
+    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 unlocked_tgt_ids:
+            if tid not in remap:
+                while next_id in blocked_ids:
+                    next_id += 1
+                remap[tid] = next_id
+                next_id += 1
+        return remap
+
+    if not available_ref_ids:
+        remap = _assign_fresh({})
+        for tid, new_id in remap.items():
+            result[target_labels == tid] = new_id
+        return result
+
+    tgt_props = _label_props(target_labels, unlocked_tgt_ids)
+    ref_props = _label_props(ref_labels, available_ref_ids)
+
+    tgt_centroids = np.array([tgt_props[t][0] for t in unlocked_tgt_ids])
+    ref_centroids = np.array([ref_props[r][0] for r in available_ref_ids])
+    dist = cdist(tgt_centroids, ref_centroids)
+    gate = dist <= max_dist_px
+
+    # Score every distance-eligible (target, reference) pair.
+    n_tgt, _ = dist.shape
+    scored: list[tuple[float, float, int, int]] = []
+    for ti in range(n_tgt):
+        t_centroid, t_area, t_coords = tgt_props[unlocked_tgt_ids[ti]]
+        for ri in np.nonzero(gate[ti])[0]:
+            r_centroid, r_area, r_coords = ref_props[available_ref_ids[ri]]
+            area_ratio = min(t_area, r_area) / max(t_area, r_area)
+            iou = scoring.centroid_corrected_iou_from_coords(
+                r_coords, r_centroid, t_coords, t_centroid
+            )
+            score = scoring.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, so a target always lands on its
+    # best still-available reference instead of being traded away by a global
+    # min-cost solver.
+    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 = unlocked_tgt_ids[ti]
+        if tid in remap or ri in used_refs:
+            continue
+        remap[tid] = available_ref_ids[ri]
+        used_refs.add(ri)
+
+    remap = _assign_fresh(remap)
+    for tid, new_id in remap.items():
+        result[target_labels == tid] = new_id
+
+    return result

itasc/tracking_ultrack/scoring.py

@@ -0,0 +1,74 @@
+"""Shared similarity scoring used by both the linker and the greedy retracker."""
+from __future__ import annotations
+
+import numpy as np
+
+
+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 target so its centroid matches 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 centroid_corrected_iou(mask_a: np.ndarray, mask_b: np.ndarray) -> float:
+    """Convenience wrapper for full-frame boolean masks."""
+    coords_a = np.argwhere(mask_a).astype(np.float32)
+    coords_b = np.argwhere(mask_b).astype(np.float32)
+    if len(coords_a) == 0 or len(coords_b) == 0:
+        return 0.0
+    centroid_a = coords_a.mean(axis=0)
+    centroid_b = coords_b.mean(axis=0)
+    return centroid_corrected_iou_from_coords(coords_a, centroid_a, coords_b, centroid_b)
+
+
+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 penalty in
+    pixels. Result can go negative for far candidates, matching the default
+    Ultrack linker's convention so ILP appear/disappear weights stay calibrated.
+    """
+    return (
+        area_weight * area_ratio
+        + iou_weight * centroid_corrected_iou
+        - distance_weight * distance
+    )