itasc 0.2.0.dev0__py3-none-any.whl

itasc/tracking_ultrack/seed_prior.py

@@ -0,0 +1,175 @@
+"""Resolve-time node probability scoring from image quality and validated seeds."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+from pathlib import Path
+
+import numpy as np
+import tifffile
+
+from itasc.tracking_ultrack._node_geometry import node_bbox_and_mask
+from itasc.tracking_ultrack.config import TrackingConfig
+
+
+@dataclass(frozen=True)
+class SeedPriorReport:
+    scored: int
+    seeds: int
+
+
+@dataclass(frozen=True)
+class _NodeScoreRecord:
+    node_id: int
+    t: int
+    bbox: tuple[int, int, int, int]
+    mask: np.ndarray
+
+
+def _load_signal_stack(path: str | Path) -> np.ndarray:
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Quality signal image not found: {path}")
+
+    arr = np.asarray(tifffile.imread(path), dtype=np.float32)
+    if arr.ndim == 2:
+        return arr[np.newaxis]
+    if arr.ndim == 3:
+        return arr
+    if arr.ndim == 4 and arr.shape[1] == 1:
+        return arr[:, 0]
+    raise ValueError(
+        f"Expected quality signal image to be 2D, 3D, or singleton-Z 4D, got {arr.shape}"
+    )
+
+
+def _binary_dilation_2d(mask: np.ndarray) -> np.ndarray:
+    padded = np.pad(np.asarray(mask, dtype=bool), 1, mode="constant")
+    h, w = mask.shape
+    dilated = np.zeros((h, w), dtype=bool)
+    for dy in range(3):
+        for dx in range(3):
+            dilated |= padded[dy:dy + h, dx:dx + w]
+    return dilated
+
+
+def compute_drop_frac(frame: np.ndarray, bbox: tuple[int, int, int, int], mask: np.ndarray) -> float:
+    y0, x0, y1, x1 = bbox
+    frame = np.asarray(frame, dtype=np.float32)
+    mask = np.asarray(mask, dtype=bool)
+    inside = frame[y0:y1, x0:x1][mask]
+    if inside.size == 0:
+        return 0.0
+    inside_median = float(np.median(inside))
+
+    pad_y0 = max(0, y0 - 1)
+    pad_x0 = max(0, x0 - 1)
+    pad_y1 = min(frame.shape[0], y1 + 1)
+    pad_x1 = min(frame.shape[1], x1 + 1)
+
+    expanded = np.zeros((pad_y1 - pad_y0, pad_x1 - pad_x0), dtype=bool)
+    inner_y0 = y0 - pad_y0
+    inner_x0 = x0 - pad_x0
+    expanded[
+        inner_y0:inner_y0 + mask.shape[0],
+        inner_x0:inner_x0 + mask.shape[1],
+    ] = mask
+    ring = _binary_dilation_2d(expanded) & ~expanded
+    if not ring.any():
+        return 0.0
+
+    ring_values = frame[pad_y0:pad_y1, pad_x0:pad_x1][ring]
+    return float(np.mean(ring_values < inside_median))
+
+
+def compute_mask_circularity(mask: np.ndarray) -> float:
+    from skimage.measure import perimeter
+
+    mask = np.asarray(mask, dtype=bool)
+    area = int(mask.sum())
+    if area == 0:
+        return 0.0
+
+    perimeter_px = float(perimeter(mask, neighborhood=4))
+    if perimeter_px <= 0.0:
+        return 0.0
+
+    circularity = 4.0 * math.pi * float(area) / (perimeter_px * perimeter_px)
+    return float(np.clip(circularity, 0.0, 1.0))
+
+
+def _seed_node_prob(base: float, max_base: float, exponent: float) -> float:
+    """Normalise ``base`` to [0, 1] and raise it to ``exponent``.
+
+    Guards the zero-quality case: ``0 ** 0 == 1`` and ``0 ** negative == inf``
+    would both corrupt the seed prior, so a non-positive normalised base scores
+    a flat ``0.0`` regardless of exponent.
+    """
+    normalized = base / max_base if max_base > 0 else base
+    if normalized <= 0.0:
+        return 0.0
+    return float(normalized ** exponent)
+
+
+def write_seed_prior_node_probs(
+    working_dir: str | Path,
+    intensity_image_path: str | Path,
+    cfg: TrackingConfig,
+) -> SeedPriorReport:
+    import sqlalchemy as sqla
+    from sqlalchemy.orm import Session
+    from ultrack.core.database import NodeDB, VarAnnotation
+
+    signal = _load_signal_stack(intensity_image_path)
+    engine = sqla.create_engine(f"sqlite:///{Path(working_dir) / 'data.db'}")
+
+    with Session(engine) as session:
+        rows = session.query(
+            NodeDB.id,
+            NodeDB.t,
+            NodeDB.pickle,
+            NodeDB.node_annot,
+        ).all()
+
+        anchor_ids: list[int] = []
+        score_updates: list[dict[str, object]] = []
+        max_base = cfg.quality_weight + cfg.circularity_weight
+        for node_id, t, node, annot in rows:
+            if annot == VarAnnotation.REAL:
+                anchor_ids.append(int(node_id))
+                continue
+            bbox, mask = node_bbox_and_mask(int(node_id), node)
+            record = _NodeScoreRecord(
+                node_id=int(node_id),
+                t=int(t),
+                bbox=bbox,
+                mask=mask,
+            )
+            if record.t >= signal.shape[0]:
+                raise ValueError(
+                    f"Quality signal image has {signal.shape[0]} frame(s), "
+                    f"cannot score node at t={record.t}"
+                )
+            drop_frac = compute_drop_frac(signal[record.t], record.bbox, record.mask)
+            circularity = compute_mask_circularity(record.mask)
+            base = cfg.quality_weight * drop_frac + cfg.circularity_weight * circularity
+            node_prob = _seed_node_prob(base, max_base, cfg.quality_exponent)
+            # NodeDB has a composite primary key (id, t); bulk_update_mappings
+            # needs every PK column present in each row mapping.
+            score_updates.append(
+                {"id": record.node_id, "t": record.t, "node_prob": node_prob}
+            )
+
+        # One bulk UPDATE instead of a statement per node (was O(n) round-trips).
+        if score_updates:
+            session.bulk_update_mappings(NodeDB, score_updates)
+        if anchor_ids:
+            session.query(NodeDB).where(NodeDB.id.in_(anchor_ids)).update(
+                {NodeDB.node_prob: 1.0},
+                synchronize_session=False,
+            )
+        session.commit()
+        scored = len(score_updates)
+
+    engine.dispose()
+    return SeedPriorReport(scored=scored, seeds=len(anchor_ids))

itasc/tracking_ultrack/solve.py

@@ -0,0 +1,52 @@
+"""Thin wrapper around ultrack.core.solve.processing.solve."""
+from __future__ import annotations
+
+from pathlib import Path
+from collections.abc import Generator
+
+from itasc.tracking_ultrack.config import TrackingConfig
+from itasc.tracking_ultrack.ingest import _build_ultrack_config
+
+
+def database_has_annotations(working_dir: str | Path) -> bool:
+    """Return whether ``data.db`` contains REAL or FAKE node annotations."""
+    import sqlalchemy as sqla
+    from sqlalchemy.orm import Session
+    from ultrack.core.database import NodeDB, VarAnnotation
+
+    db_path = Path(working_dir) / "data.db"
+    if not db_path.exists():
+        return False
+
+    engine = sqla.create_engine(f"sqlite:///{db_path}")
+    try:
+        with Session(engine) as session:
+            return bool(
+                session.query(NodeDB.id)
+                .where(NodeDB.node_annot.in_([VarAnnotation.REAL, VarAnnotation.FAKE]))
+                .limit(1)
+                .first()
+            )
+    finally:
+        engine.dispose()
+
+
+def run_solve(
+    working_dir: str | Path,
+    cfg: TrackingConfig,
+    *,
+    overwrite: bool = True,
+    use_annotations: bool | None = True,
+) -> Generator[tuple[int, int, str], None, None]:
+    """Run the ILP solver, yielding (step, total, label) progress tuples."""
+    from ultrack.core.solve.processing import solve
+
+    total = 2
+    wd = Path(working_dir)
+    ultrack_cfg = _build_ultrack_config(cfg, wd)
+    if use_annotations is None:
+        use_annotations = True
+
+    yield (0, total, "Running ILP solver…")
+    solve(ultrack_cfg, overwrite=overwrite, use_annotations=use_annotations)
+    yield (total, total, "Solve done.")

itasc/tracking_ultrack/swap_candidate.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from itasc.tracking_ultrack._node_geometry import (
+    intersection_area,
+    mask_bbox,
+    match_mask_to_node,
+    node_bbox_and_mask,
+)
+
+
+@dataclass(frozen=True)
+class SwapCandidate:
+    node_id: int
+    mask_2d: np.ndarray
+    bbox: tuple[int, int, int, int]
+    centroid: tuple[float, float]
+    area: int
+
+
+@dataclass
+class _SwapCursor:
+    source_id: int
+    frame: int
+    candidates: tuple[SwapCandidate, ...]
+    index: int
+    baseline_frame: np.ndarray | None = None
+
+
+def list_swap_candidates(
+    *,
+    db_path: Path,
+    frame: int,
+    source_mask: np.ndarray,
+    frame_shape: tuple[int, int],
+    protected_mask: np.ndarray | None = None,
+) -> list[SwapCandidate]:
+    """Return the matched node's nesting lineage around ``source_mask``.
+
+    Matches ``source_mask`` to the best-overlapping Ultrack node at ``frame``,
+    then walks the candidate **containment lattice** recorded in ``OverlapDB``:
+    every node that overlaps the matched node *and* is nested with it — a
+    superset (larger merged segment that contains it) or a subset (smaller
+    fragment it contains) — plus the matched node itself, returned sorted by
+    area. Partially-overlapping neighbours (siblings/cousins) are excluded: they
+    are different cells, so swapping onto them would relocate the cell rather
+    than resize it. ``Z``/``C`` cycle through this area-sorted lineage by index.
+
+    ``OverlapDB`` is populated by both database builders (atom-union enumeration
+    and canonical Ultrack segmentation), whereas ``hier_parent_id`` is only set
+    by the latter — so the overlap lattice is the portable source of structure.
+    """
+    if not Path(db_path).exists():
+        return []
+
+    if mask_bbox(source_mask) is None:
+        return []
+
+    import sqlalchemy as sqla
+    from sqlalchemy import or_
+    from sqlalchemy.orm import Session
+    from ultrack.core.database import NodeDB, OverlapDB
+
+    engine = sqla.create_engine(
+        f"sqlite:///{db_path}", connect_args={"check_same_thread": False}
+    )
+
+    try:
+        with Session(engine) as session:
+            # Match source mask to a node by max IoU (prefiltered by bbox so only
+            # plausibly-overlapping masks are deserialized).
+            matched = match_mask_to_node(session, frame, source_mask)
+            if matched is None:
+                return []
+            matched_id, matched_bbox, matched_mask = matched
+            matched_area = int(matched_mask.sum())
+
+            # Candidate set = the matched node plus everything that overlaps it
+            # (shares atoms) per OverlapDB. The pair direction is not meaningful
+            # here, so match the id on either column.
+            overlap_ids: set[int] = set()
+            for node_id, ancestor_id in session.query(
+                OverlapDB.node_id, OverlapDB.ancestor_id
+            ).filter(
+                or_(OverlapDB.node_id == matched_id, OverlapDB.ancestor_id == matched_id)
+            ):
+                overlap_ids.add(int(node_id))
+                overlap_ids.add(int(ancestor_id))
+            overlap_ids.add(matched_id)
+
+            node_rows = (
+                session.query(NodeDB)
+                .filter(NodeDB.t == frame, NodeDB.id.in_(overlap_ids))
+                .all()
+            )
+
+            results: list[SwapCandidate] = []
+            for node in node_rows:
+                try:
+                    (y0, x0, y1, x1), mask_crop = node_bbox_and_mask(
+                        int(node.id), node.pickle
+                    )
+                except Exception:
+                    continue
+                if mask_crop.shape != (y1 - y0, x1 - x0):
+                    continue
+
+                # Keep only nodes nested with the matched node: one mask must be
+                # a subset of the other. Partially-overlapping neighbours (which
+                # share atoms but each hold unique ones) are different cells.
+                if int(node.id) != matched_id:
+                    inter = intersection_area(
+                        matched_bbox, matched_mask, (y0, x0, y1, x1), mask_crop
+                    )
+                    cand_area = int(mask_crop.sum())
+                    nested = inter == matched_area or inter == cand_area
+                    if not nested:
+                        continue
+
+                full_mask = np.zeros(frame_shape, dtype=bool)
+                full_mask[y0:y1, x0:x1] = mask_crop
+                if not full_mask.any():
+                    continue
+
+                if protected_mask is not None:
+                    paintable = full_mask & ~protected_mask
+                    if not paintable.any():
+                        continue
+
+                area = int(full_mask.sum())
+                results.append(
+                    SwapCandidate(
+                        node_id=int(node.id),
+                        mask_2d=full_mask,
+                        bbox=(y0, x0, y1, x1),
+                        centroid=(float(node.y), float(node.x)),
+                        area=area,
+                    )
+                )
+    finally:
+        engine.dispose()
+
+    results.sort(key=lambda c: c.area)
+    return results
+
+
+def nearest_area_index(
+    candidates: tuple[SwapCandidate, ...] | list[SwapCandidate],
+    area: int,
+) -> int:
+    """Index of the candidate whose area is closest to ``area``.
+
+    Used to seed the cursor on the lineage member that best matches the
+    currently displayed (matched) cell, so the first ``Z``/``C`` step moves
+    relative to it.
+    """
+    return min(range(len(candidates)), key=lambda i: abs(candidates[i].area - area))
+
+
+def cycle_index(count: int, current: int, *, larger: bool) -> int:
+    """Next index when cycling an area-sorted candidate list.
+
+    ``larger`` steps toward bigger areas (higher index), otherwise toward
+    smaller areas (lower index); both directions wrap around.
+    """
+    step = 1 if larger else -1
+    return (current + step) % count

itasc/tracking_ultrack/track_quality.py

@@ -0,0 +1,142 @@
+"""Per-track quality scoring derived from the solved Ultrack database.
+
+Pure module (no Qt, no napari). The quality of an exported track is the sum of
+its nodes' ``node_prob`` plus the sum of its selected edges' ``LinkDB.weight``;
+higher is better. Ordering by this score lets the UI relabel so ID 1 is the
+best track, 2 the next, and so on.
+
+The score is keyed by the painted label value, which equals the exported
+``track_id`` from :func:`ultrack.core.export.to_tracks_layer` — the same
+dataframe consumed in :mod:`itasc.tracking_ultrack.corrections`.
+"""
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+
+from itasc.tracking_ultrack.config import TrackingConfig
+from itasc.tracking_ultrack.db_query import _engine, annotation_name
+
+# NULL node_prob / link weight render as 1.0, matching db_query.py.
+_NULL_VALUE = 1.0
+
+
+def _resolve_prob(value: float | None) -> float:
+    if value is None:
+        return _NULL_VALUE
+    return float(value)
+
+
+def _edge_weight(info: tuple[float | None, str] | None) -> float:
+    """Weight contributed by one selected edge.
+
+    ``None`` means no link row was found for the pair (contributes 0). A FAKE
+    link contributes 0. Otherwise a NULL weight renders as 1.0.
+    """
+    if info is None:
+        return 0.0
+    weight, annotation = info
+    if annotation == "FAKE":
+        return 0.0
+    return _NULL_VALUE if weight is None else float(weight)
+
+
+def compute_track_scores(
+    tracks_df,
+    node_probs: Mapping[int, float | None],
+    links: Mapping[tuple[int, int], tuple[float | None, str]],
+) -> dict[int, float]:
+    """Pure scoring core: ``track_id -> Σ node_prob + Σ edge weight``.
+
+    ``tracks_df`` needs ``id``, ``t`` and ``track_id`` columns (one row per
+    node). ``node_probs`` maps node id -> raw ``node_prob`` (``None`` allowed).
+    ``links`` maps a (source_id, target_id) pair -> (raw weight, annotation
+    name); the pair is ordered earlier-frame -> later-frame, matching LinkDB.
+    """
+    scores: dict[int, float] = {}
+    for track_id, group in tracks_df.groupby("track_id"):
+        ordered = group.sort_values("t")
+        ids = [int(i) for i in ordered["id"].tolist()]
+        score = sum(_resolve_prob(node_probs.get(nid)) for nid in ids)
+        for src, tgt in zip(ids, ids[1:]):
+            score += _edge_weight(links.get((src, tgt)))
+        scores[int(track_id)] = score
+    return scores
+
+
+def quality_order(scores: Mapping[int, float]) -> list[int]:
+    """track_ids sorted by score desc, then by id asc as a stable tiebreak."""
+    return [tid for tid, _ in sorted(scores.items(), key=lambda kv: (-kv[1], kv[0]))]
+
+
+def _query_node_probs(db_path: Path, node_ids: list[int]) -> dict[int, float | None]:
+    from sqlalchemy.orm import Session
+    from ultrack.core.database import NodeDB
+
+    engine = _engine(db_path)
+    try:
+        with Session(engine) as session:
+            rows = (
+                session.query(NodeDB.id, NodeDB.node_prob)
+                .filter(NodeDB.id.in_(node_ids))
+                .all()
+            )
+    finally:
+        engine.dispose()
+    return {int(nid): (None if prob is None else float(prob)) for nid, prob in rows}
+
+
+def _query_links(
+    db_path: Path, node_ids: list[int]
+) -> dict[tuple[int, int], tuple[float | None, str]]:
+    from sqlalchemy.orm import Session
+    from ultrack.core.database import LinkDB
+
+    wanted = set(node_ids)
+    engine = _engine(db_path)
+    try:
+        with Session(engine) as session:
+            rows = (
+                session.query(
+                    LinkDB.source_id,
+                    LinkDB.target_id,
+                    LinkDB.weight,
+                    LinkDB.annotation,
+                )
+                .filter(LinkDB.source_id.in_(node_ids))
+                .all()
+            )
+    finally:
+        engine.dispose()
+    links: dict[tuple[int, int], tuple[float | None, str]] = {}
+    for src, tgt, weight, annotation in rows:
+        if int(tgt) not in wanted:
+            continue
+        links[(int(src), int(tgt))] = (
+            None if weight is None else float(weight),
+            annotation_name(annotation),
+        )
+    return links
+
+
+def track_quality_scores(db_path: str | Path, cfg: TrackingConfig) -> dict[int, float]:
+    """Map exported ``track_id`` -> quality score (Σ node_prob + Σ edge weight).
+
+    Read-only over the solved database. Returns ``{}`` when the export yields no
+    tracks (e.g. an unsolved or empty database).
+    """
+    from itasc.tracking_ultrack.ingest import _build_ultrack_config
+
+    db_path = Path(db_path)
+    ultrack_cfg = _build_ultrack_config(cfg, db_path.parent)
+
+    from ultrack.core.export import to_tracks_layer
+
+    tracks_df, _graph = to_tracks_layer(ultrack_cfg)
+    if tracks_df is None or tracks_df.empty or "id" not in tracks_df:
+        return {}
+
+    node_ids = [int(i) for i in tracks_df["id"].tolist()]
+    node_probs = _query_node_probs(db_path, node_ids)
+    links = _query_links(db_path, node_ids)
+    return compute_track_scores(tracks_df, node_probs, links)