Source code for traccuracy.metrics._compute_complete_tracks_by_length

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from traccuracy._tracking_graph import EdgeFlag, NodeFlag

if TYPE_CHECKING:
    from typing import Literal

    from traccuracy.matchers import Matched

# Sentinel values for the DP grid
EMPTY = -1
CORRECT = 1
INCORRECT = 0


[docs] def compute_complete_tracks_by_length( matched: Matched, max_length: int, lineages: bool = True, error_type: Literal["basic", "ctc"] = "basic", relax_skips_gt: bool = False, relax_skips_pred: bool = False, ) -> dict[int, tuple[int, int]]: r"""Compute the fraction of GT track segments correctly reconstructed, by length. For each component (lineage or tracklet), builds a 2D grid of per-frame-step correctness values (see ``_build_grid``), then uses dynamic programming to combine adjacent frame steps into longer segments. Segments are accumulated via a sliding window over the timeline, one window per segment length. Length is measured in frames (time difference), not edge count. Skip edges spanning multiple frames are interpolated to fill intermediate frame steps with the same correctness status. This means that if a skip edge exceeds the length, the part of it inside the window still counts towards the correctness of the segment. Example — A->A' then A' divides into B and C, with an error on the C->C' edge (1 = correct, 0 = incorrect, empty = no track):: GT graph (6 nodes, 5 edges): t=0 t=1 t=2 t=3 / B ---- B' A ---- A'< \ C -x-- C' (error on C->C') base grid (window=1): one entry per edge Each column is frame span ((0-1) covers frames 0->1, etc.) (0-1)(1-2)(2-3) ┌─────────┬────┬────┬────┐ │track. A │ 1 │ │ │ edge A->A' ├─────────┼────┼────┼────┤ │track. B │ │ 1 │ 1 │ edges A'->B, B->B' ├─────────┼────┼────┼────┤ │track. C │ │ 1 │ 0 │ edges A'->C, C->C' (error) └─────────┴────┴────┴────┘ w=2 grid (combine base[t] with base[t+1] via division links): (0-2)(1-3) ┌─────────┬────┬────┐ │track. A │ 1 │ │ AND(A->A'=1, A'->B=1, A'->C=1) = 1 ├─────────┼────┼────┤ │track. B │ │ 1 │ AND(A'->B=1, B->B'=1) = 1 ├─────────┼────┼────┤ │track. C │ │ 0 │ AND(A'->C=1, C->C'=0) = 0 └─────────┴────┴────┘ w=3 grid: (0-3) ┌─────────┬────┐ │track. A │ 0 │ AND(A->A'=1, w2_B=1, w2_C=0) = 0 ├─────────┼────┤ │track. B │ │ ├─────────┼────┤ │track. C │ │ └─────────┴────┘ At each length, non-empty entries are counted as segments. length 1: 5 (4 correct). length 2: 3 (2 correct). length 3: 1 (0 correct). GT tracks shorter than a given length still count once for that length (correct iff the entire track is correct). Args: matched: Matched data object with annotated errors max_length: Maximum track length to evaluate (in frames) lineages: If True, evaluate on full lineages. If False, on tracklets. error_type: "basic" or "ctc" - which error classification was used relax_skips_gt: If True, SKIP_TRUE_POS edges in GT count as correct relax_skips_pred: If True, SKIP_TRUE_POS edges in pred count as correct Returns: Dictionary mapping track length (int) to tuple of (correct_count, total_count) for each length from 1 to max_length """ is_ctc = error_type == "ctc" # Get components (lineages or tracklets) if lineages: components = matched.gt_graph.get_lineages() else: components = matched.gt_graph.get_tracklets(include_division_edges=False) results: dict[int, tuple[int, int]] = {} for gt_track in components: # Build the w=1 grid and division links for this component. # grid[t][row] = EMPTY, CORRECT, or INCORRECT for step t -> t+1. # divisions: row -> list of daughter row indices. grid, divisions, num_rows = _build_grid( gt_track, matched, is_ctc, relax_skips_gt, relax_skips_pred, ) T = len(grid) # DP to build grids for increasing segment lengths. # For length w, combine base_grid[t] (w=1) with prev_grid[t+1] (w-1). base_grid = grid prev_grid = grid for w in range(1, max_length + 1): if w == 1: cur_grid = base_grid else: cur_len = T - w + 1 if cur_len <= 0: break cur_grid = [[] for _ in range(cur_len)] for t in range(cur_len): cur_grid[t] = [EMPTY] * num_rows for row in range(num_rows): val_base = base_grid[t][row] # Only follow division links when the parent row # is occupied at this time step (val_base != EMPTY). divs = divisions if val_base != EMPTY else None val_prev = _get_continuation_value(row, t, prev_grid, divs) cur_grid[t][row] = _combine(val_base, val_prev) prev_grid = cur_grid correct, total = _count_grid(cur_grid) # aggregate results across different lineages by updating results dict sum_correct, sum_total = results.get(w, (0, 0)) results[w] = (sum_correct + correct, sum_total + total) return results
def _combine(a: int, b: int) -> int: """AND two correctness values: a segment is correct only if all parts are correct. EMPTY means "no data" and is ignored (identity element). """ if a == EMPTY: return b if b == EMPTY: return a if a == INCORRECT or b == INCORRECT: return INCORRECT return CORRECT def _get_continuation_value( row: int, t: int, prev_grid: list[list[int]], divisions: dict[int, list[int]] | None, ) -> int: """Look up the correctness of the continuation track(s) from the edge at (t, row). For a non-dividing tracklet, the continuation is the same row at t+1. At a division, the continuation includes all daughter rows at t+1, AND'd together (a segment spanning a division is only correct if all branches are correct). Returns EMPTY if there is no continuation (t+1 past the grid, or divisions=None and the row has no data — EMPTY acts as identity in _combine so the current cell's value is unchanged). """ next_t = t + 1 if next_t >= len(prev_grid): return EMPTY if divisions is not None and row in divisions: result = EMPTY for daughter_row in divisions[row]: daughter_val = prev_grid[next_t][daughter_row] result = _combine(result, daughter_val) return result return prev_grid[next_t][row] def _count_grid(cur_grid: list[list[int]]) -> tuple[int, int]: """Count how many segments exist in the grid and how many are correct.""" total = 0 correct = 0 for t_col in cur_grid: for val in t_col: if val != EMPTY: total += 1 if val == CORRECT: correct += 1 return correct, total def _build_grid( gt_track: Any, matched: Matched, is_ctc: bool, relax_skips_gt: bool, relax_skips_pred: bool, ) -> tuple[list[list[int]], dict[int, list[int]], int]: """Build the base (window=1) correctness grid for a single GT component. Maps a lineage or tracklet tree onto a 2D grid where each column is a frame step and each row is one branch of the tree. Each entry records whether the corresponding edge was correctly reconstructed. At divisions, the parent row ends and daughters get new rows. Also returns division links so that larger windows can AND across all daughter branches when a segment spans a division. Skip edges spanning multiple frames are interpolated: each intermediate frame step gets the same correctness value. Returns: grid: grid[t][row] = CORRECT, INCORRECT, or EMPTY divisions: parent row -> list of daughter row indices num_rows: total number of rows """ gt_graph = matched.gt_graph frame_key = gt_graph.frame_key # Determine time range from global frames - this ensures that short gt tracks # contribute to larger window sizes, but might be inefficient min_frame = gt_graph.start_frame max_frame = gt_graph.end_frame if min_frame is None or max_frame is None: return [], {}, 0 T = max_frame - min_frame - 1 # Number of frame steps if T <= 0: return [], {}, 0 grid: list[list[int]] = [[] for _ in range(T)] divisions: dict[int, list[int]] = {} # Each component has exactly one root node root = next(n for n in gt_track.nodes() if gt_track.in_degree(n) == 0) num_rows = 1 # DFS through the track tree: each entry is (node, grid row for that node) nodes_to_process = [(root, 0)] while nodes_to_process: node, cur_row = nodes_to_process.pop() out_edges = list(gt_track.out_edges(node)) if not out_edges: continue node_correct = _is_node_correct( node, matched, is_ctc, relax_skips_pred, ) is_division = len(out_edges) > 1 source_frame = gt_graph.nodes[node][frame_key] daughter_rows = [] for edge in out_edges: # Divisions get new rows; single edges continue on cur_row if is_division: edge_row = num_rows num_rows += 1 daughter_rows.append(edge_row) else: edge_row = cur_row target = edge[1] target_frame = gt_graph.nodes[target][frame_key] edge_span = target_frame - source_frame target_correct = _is_node_correct( target, matched, is_ctc, relax_skips_pred, ) edge_correct = ( node_correct and not _has_fp_division(node, matched, is_ctc) and _is_edge_correct( edge, matched, is_ctc, relax_skips_gt, relax_skips_pred, ) and target_correct ) val = CORRECT if edge_correct else INCORRECT for dt in range(edge_span): t_idx = source_frame - min_frame + dt if 0 <= t_idx < T: _grid_set(grid, t_idx, edge_row, val) nodes_to_process.append((target, edge_row)) # Record division link for this row. # If the division is at the root (no pre-division steps on # this row), no link is needed — daughters are independent. if is_division and source_frame > min_frame: divisions[cur_row] = daughter_rows # Pad all columns to the same number of rows for t in range(T): while len(grid[t]) < num_rows: grid[t].append(EMPTY) return grid, divisions, num_rows def _grid_set(grid: list[list[int]], t: int, row: int, val: int) -> None: """Set grid[t][row] = val, extending the row list if needed.""" while len(grid[t]) <= row: grid[t].append(EMPTY) grid[t][row] = val def _is_node_correct( gt_node: Any, matched: Matched, is_ctc: bool, relax_skips_pred: bool, ) -> bool: """Is this GT node correctly detected in the prediction? A node is correct if it is a true positive. Division errors are checked separately (FP_DIV via ``_has_fp_division``, FN_DIV via edge error flags). With skip relaxation, nodes between matched skip edges also count as correct. """ gt_graph = matched.gt_graph node_tp = NodeFlag.CTC_TRUE_POS if is_ctc else NodeFlag.TRUE_POS gt_node_data = gt_graph.nodes[gt_node] # All mapped nodes are marked TP by classify_basic_errors/evaluate_ctc_events, # so checking the GT node flag is sufficient. if node_tp in gt_node_data: return True # If not a TP, check if it's between skip edges (when relaxing) if relax_skips_pred: for prev_edge in gt_graph.graph.in_edges(gt_node): if EdgeFlag.SKIP_TRUE_POS in gt_graph.edges[prev_edge]: return True return False def _has_fp_division( gt_node: Any, matched: Matched, is_ctc: bool, ) -> bool: """Does this GT node have a false positive division in the prediction? A spurious division means the prediction splits the track where the GT doesn't. This invalidates windows starting from this node, since the outgoing pred edges include a division that shouldn't exist. Only applies to basic errors — CTC handles division errors via WRONG_SEMANTIC edge flags instead. """ if is_ctc: return False pred_graph = matched.pred_graph for pred_node in matched.get_gt_pred_matches(gt_node): if NodeFlag.FP_DIV in pred_graph.nodes[pred_node]: return True return False def _is_edge_correct( gt_edge: tuple[Any, Any], matched: Matched, is_ctc: bool, relax_skips_gt: bool, relax_skips_pred: bool, ) -> bool: """Is this GT edge correctly reconstructed in the prediction? An edge is correct if the prediction has a matching true positive link with no semantic errors. With skip relaxation, matched skip edges also count as correct. """ gt_graph = matched.gt_graph pred_graph = matched.pred_graph edge_data = gt_graph.edges[gt_edge] if is_ctc: # CTC errors don't annotate edge TPs, so check for absence of error flags if EdgeFlag.CTC_FALSE_NEG in edge_data: return False # Also check matched pred edges for CTC error flags # Note: INTERTRACK_EDGE just marks division edges, not errors matched_sources = matched.get_gt_pred_matches(gt_edge[0]) matched_targets = matched.get_gt_pred_matches(gt_edge[1]) for src in matched_sources: for tgt in matched_targets: if pred_graph.graph.has_edge(src, tgt): pred_edge_data = pred_graph.edges[(src, tgt)] if EdgeFlag.WRONG_SEMANTIC in pred_edge_data: return False return True else: # Check for regular TP if EdgeFlag.TRUE_POS in edge_data: return True # Check for skip TP if the matching relaxation flag is enabled: # - GT skip edge (spans multiple frames in GT) needs relax_skips_gt # - Non-skip edge with SKIP_TRUE_POS (pred has the skip) needs relax_skips_pred if EdgeFlag.SKIP_TRUE_POS in edge_data: is_skip_edge = gt_graph.is_skip_edge(gt_edge) if is_skip_edge and relax_skips_gt: return True if not is_skip_edge and relax_skips_pred: return True return False