rpx_benchmark.metrics.tracking_paper

  1"""Paper-protocol metrics for RPX D3 multi-object tracking.
  2
  3RPX publishes temporally consistent instance masks: a positive pixel value is
  4the same physical object ID throughout a ``(scene, phase)`` clip.  D3 models
  5are initialized from the first-frame mask.  For the MOTChallenge metric set we
  6derive tight boxes from those masks and delegate the sequence calculations to
  7the official TrackEval implementation.
  8"""
  9
 10from __future__ import annotations
 11
 12from typing import Any, Sequence
 13
 14import numpy as np
 15
 16from ..exceptions import MetricError
 17
 18PAPER_TRACKING_METRICS = ("hota", "deta", "assa", "idf1", "mota", "idsw")
 19MOT_IOU_THRESHOLD = 0.5
 20
 21
 22def _validate_mask_sequence(masks: Sequence[np.ndarray], name: str) -> list[np.ndarray]:
 23    if not masks:
 24        raise MetricError(f"{name} mask sequence is empty.")
 25    result: list[np.ndarray] = []
 26    expected_shape: tuple[int, int] | None = None
 27    for index, value in enumerate(masks):
 28        mask = np.asarray(value)
 29        if mask.ndim != 2:
 30            raise MetricError(f"{name}[{index}] must be 2-D, got shape {mask.shape}.")
 31        if not np.issubdtype(mask.dtype, np.integer):
 32            raise MetricError(f"{name}[{index}] must have an integer dtype, got {mask.dtype}.")
 33        if np.any(mask < 0):
 34            raise MetricError(f"{name}[{index}] contains negative instance IDs.")
 35        if expected_shape is None:
 36            expected_shape = mask.shape
 37        elif mask.shape != expected_shape:
 38            raise MetricError(
 39                f"{name}[{index}] has shape {mask.shape}; expected {expected_shape}."
 40            )
 41        result.append(mask.astype(np.int32, copy=False))
 42    return result
 43
 44
 45def _boxes_and_ids(mask: np.ndarray, id_map: dict[int, int]) -> tuple[np.ndarray, np.ndarray]:
 46    source_ids = np.unique(mask)
 47    source_ids = source_ids[source_ids > 0]
 48    boxes: list[list[float]] = []
 49    ids: list[int] = []
 50    for source_id in source_ids.tolist():
 51        ys, xs = np.nonzero(mask == source_id)
 52        if xs.size == 0:
 53            continue
 54        # Exclusive maxima follow the common xyxy convention and make a
 55        # one-pixel-wide object have positive area.
 56        boxes.append(
 57            [
 58                float(xs.min()),
 59                float(ys.min()),
 60                float(xs.max() + 1),
 61                float(ys.max() + 1),
 62            ]
 63        )
 64        ids.append(id_map[int(source_id)])
 65    return np.asarray(ids, dtype=np.int64), np.asarray(boxes, dtype=np.float64).reshape(-1, 4)
 66
 67
 68def _box_iou_matrix(gt_boxes: np.ndarray, pred_boxes: np.ndarray) -> np.ndarray:
 69    if gt_boxes.size == 0 or pred_boxes.size == 0:
 70        return np.zeros((len(gt_boxes), len(pred_boxes)), dtype=np.float64)
 71    left_top = np.maximum(gt_boxes[:, None, :2], pred_boxes[None, :, :2])
 72    right_bottom = np.minimum(gt_boxes[:, None, 2:], pred_boxes[None, :, 2:])
 73    intersection_wh = np.maximum(0.0, right_bottom - left_top)
 74    intersection = intersection_wh[..., 0] * intersection_wh[..., 1]
 75    gt_area = np.prod(np.maximum(0.0, gt_boxes[:, 2:] - gt_boxes[:, :2]), axis=1)
 76    pred_area = np.prod(np.maximum(0.0, pred_boxes[:, 2:] - pred_boxes[:, :2]), axis=1)
 77    union = gt_area[:, None] + pred_area[None, :] - intersection
 78    return np.divide(
 79        intersection,
 80        union,
 81        out=np.zeros_like(intersection),
 82        where=union > 0,
 83    )
 84
 85
 86def _trackeval_data(
 87    gt_masks: Sequence[np.ndarray],
 88    pred_masks: Sequence[np.ndarray],
 89) -> dict[str, Any]:
 90    gt = _validate_mask_sequence(gt_masks, "ground_truth")
 91    pred = _validate_mask_sequence(pred_masks, "prediction")
 92    if len(gt) != len(pred):
 93        raise MetricError(
 94            f"prediction has {len(pred)} frames; ground truth has {len(gt)} frames."
 95        )
 96    if pred[0].shape != gt[0].shape:
 97        raise MetricError(
 98            f"prediction shape {pred[0].shape} does not match GT shape {gt[0].shape}."
 99        )
100
101    gt_source_ids = sorted({int(v) for mask in gt for v in np.unique(mask) if v > 0})
102    pred_source_ids = sorted({int(v) for mask in pred for v in np.unique(mask) if v > 0})
103    gt_id_map = {source_id: index for index, source_id in enumerate(gt_source_ids)}
104    pred_id_map = {source_id: index for index, source_id in enumerate(pred_source_ids)}
105
106    gt_ids: list[np.ndarray] = []
107    tracker_ids: list[np.ndarray] = []
108    similarity_scores: list[np.ndarray] = []
109    num_gt_dets = 0
110    num_tracker_dets = 0
111    for gt_mask, pred_mask in zip(gt, pred, strict=True):
112        frame_gt_ids, gt_boxes = _boxes_and_ids(gt_mask, gt_id_map)
113        frame_pred_ids, pred_boxes = _boxes_and_ids(pred_mask, pred_id_map)
114        gt_ids.append(frame_gt_ids)
115        tracker_ids.append(frame_pred_ids)
116        similarity_scores.append(_box_iou_matrix(gt_boxes, pred_boxes))
117        num_gt_dets += len(frame_gt_ids)
118        num_tracker_dets += len(frame_pred_ids)
119
120    return {
121        "num_timesteps": len(gt),
122        "num_gt_ids": len(gt_source_ids),
123        "num_tracker_ids": len(pred_source_ids),
124        "num_gt_dets": num_gt_dets,
125        "num_tracker_dets": num_tracker_dets,
126        "gt_ids": gt_ids,
127        "tracker_ids": tracker_ids,
128        "similarity_scores": similarity_scores,
129    }
130
131
132def paper_tracking_metrics(
133    pred_masks: Sequence[np.ndarray],
134    gt_masks: Sequence[np.ndarray],
135) -> dict[str, float]:
136    """Compute HOTA, DetA, AssA, IDF1, MOTA and ID switches for one clip.
137
138    TrackEval is intentionally a runtime dependency rather than a locally
139    reimplemented approximation.  HOTA is the standard mean over association
140    thresholds 0.05:0.05:0.95; CLEAR and Identity use the MOTChallenge IoU
141    threshold of 0.5.
142    """
143
144    try:
145        # TrackEval still references NumPy's removed aliases at the pinned
146        # upstream revision. Restore those aliases locally before importing it;
147        # this does not change any calculation.
148        if not hasattr(np, "float"):
149            np.float = float  # type: ignore[attr-defined]
150        if not hasattr(np, "int"):
151            np.int = int  # type: ignore[attr-defined]
152        from trackeval.metrics import CLEAR, HOTA, Identity
153    except ImportError as exc:
154        raise MetricError(
155            "Paper D3 metrics require the official TrackEval package.",
156            hint="Use the RPX tracking Docker image; it pins TrackEval.",
157        ) from exc
158
159    data = _trackeval_data(gt_masks=gt_masks, pred_masks=pred_masks)
160    clear_result = CLEAR({"THRESHOLD": MOT_IOU_THRESHOLD, "PRINT_CONFIG": False}).eval_sequence(
161        data
162    )
163    identity_result = Identity(
164        {"THRESHOLD": MOT_IOU_THRESHOLD, "PRINT_CONFIG": False}
165    ).eval_sequence(data)
166    hota_result = HOTA().eval_sequence(data)
167
168    metrics = {
169        "hota": float(np.mean(hota_result["HOTA"])),
170        "deta": float(np.mean(hota_result["DetA"])),
171        "assa": float(np.mean(hota_result["AssA"])),
172        "idf1": float(identity_result["IDF1"]),
173        "mota": float(clear_result["MOTA"]),
174        "idsw": float(clear_result["IDSW"]),
175    }
176    if not all(np.isfinite(value) for value in metrics.values()):
177        raise MetricError(f"TrackEval returned non-finite metrics: {metrics}")
178    return metrics