Add evaluation metrics to EgoTracks

EgoTracks/tools/eval_datasets/eval_ego4d_lt_tracking.py

@@ -171,3 +171,131 @@ def gather_ego4d_lt_tracking_result(result):
         gathered[seq_name] = res
 
     return gathered
+
+
+def calculate_ego4d_lt_tracking_metrics(cfg):
+    use_visual_clip = cfg.EVAL.EGO4DLT.USE_VISUAL_CLIP
+    result_dir = os.path.join(
+        cfg.OUTPUT_DIR,
+        "eval",
+        "EGO4DLTTracking",
+        f"{cfg.EVAL.EGO4DLT.TRACK_MODE}",
+        f"{cfg.MODEL_TYPE}",
+    )
+    intermediate_dir = os.path.join(result_dir, "intermediate_result")
+
+    result = {}
+    # Strange, sometimes certain nodes do not save any result
+    # for shard in range(total_machines):
+    #     path = os.path.join(result_dir, "intermediate_result", f"{shard}.pkl")
+    #     logging.info(path)
+    #     shard_result = pkl.load(pathmgr.open(path, "rb"))
+    #     result.extend(shard_result)
+    files = pathmgr.ls(intermediate_dir)
+    for f in files:
+        path = os.path.join(intermediate_dir, f)
+        logging.info(path)
+        shard_result = pkl.load(pathmgr.open(path, "rb"))
+        result.update(shard_result)
+
+    result = gather_ego4d_lt_tracking_result(result)
+    path = os.path.join(result_dir, "result.pkl")
+    pkl.dump(result, pathmgr.open(path, "wb"))
+    logging.info(f"Total number of predicted clips is {len(result)}.")
+
+    annotation_path = cfg.EVAL.EGO4DLT.ANNOTATION_PATH
+    data_dir = cfg.EVAL.EGO4DLT.DATA_DIR
+    eval_ratio = cfg.EVAL.EGO4DLT.EVAL_RATIO
+
+    # Initilize data loader
+    logging.info("building dataset")
+    eval_dataset = EGO4DLTTrackingDataset(data_dir, annotation_path, ratio=eval_ratio)
+
+    iou_per_video = []
+    all_pred_bboxes = []
+    all_gt_bboxes = []
+    all_pred_scores = []
+
+    for seq in eval_dataset:
+        pred_bbox_dict = {r["frame_number"]: r for r in result[seq.name]["pred_bboxes"]}
+        # Only evaluate on frames annotated, all others ignored
+        exclude_frame_numbers = [
+            b["frame_number"]
+            for b in result[seq.name]["pred_bboxes"]
+            if b["type"] == "gt"
+        ]
+        exclude_frame_numbers = set(exclude_frame_numbers)
+
+        if not use_visual_clip:
+            for frame_number, _ in seq.visual_crop.items():
+                exclude_frame_numbers.add(frame_number)
+        else:
+            raise NotImplementedError
+
+        logging.info(exclude_frame_numbers)
+
+        total_frame_numbers = len(seq.frames)
+        ious = []
+        gt_bboxes = [None for _ in range(total_frame_numbers)]
+        pred_bboxes = [None for _ in range(total_frame_numbers)]
+        pred_scores = [0 for _ in range(total_frame_numbers)]
+        for frame_number, gt_bbox in seq.gt_bbox_dict.items():
+            # Ignore excluded frames.
+            if frame_number in exclude_frame_numbers:
+                continue
+            gt_bboxes[frame_number] = gt_bbox
+            if frame_number not in pred_bbox_dict:
+                ious.append(0)
+            else:
+                pred_bbox = pred_bbox_dict[frame_number]["bbox"]
+                pred_score = pred_bbox_dict[frame_number]["score"]
+                iou = compute_overlaps([gt_bbox], [pred_bbox])[0]
+                ious.append(iou)
+                pred_bboxes[frame_number] = pred_bbox
+                pred_scores[frame_number] = pred_score
+
+        for frame_number, pred_bbox in pred_bbox_dict.items():
+            pred_bboxes[frame_number] = pred_bbox["bbox"]
+            pred_scores[frame_number] = pred_bbox["score"]
+
+        if len(ious):
+            iou_per_video.append(np.mean(ious))
+
+        all_gt_bboxes.append(gt_bboxes)
+        all_pred_bboxes.append(pred_bboxes)
+        all_pred_scores.append(pred_scores)
+
+    average_overlap = np.mean(iou_per_video)
+    precision, recall = compute_precision_and_recall(
+        all_pred_scores, all_pred_bboxes, all_gt_bboxes
+    )
+    f1_score, pr_score, re_score = compute_f_score(precision, recall)
+    logging.info(f"IoU per video: {iou_per_video}")
+    logging.info(f"Eval result mIoU {average_overlap}, f1 {f1_score}.")
+
+    # save evaluation result
+    path = os.path.join(result_dir, f"{cfg.MODEL_TYPE}_eval_result.pkl")
+    pkl.dump(
+        {
+            "total_video": len(iou_per_video),
+            "F1": f1_score,
+            "precision_score": pr_score,
+            "re_score": re_score,
+            "AO": average_overlap,
+            "precision": precision,
+            "recall": recall,
+            "result_dir": result_dir,
+            "MODEL_WEIGHTS": cfg.MODEL.WEIGHTS,
+        },
+        pathmgr.open(path, "wb"),
+    )
+
+    return {
+        "total_video": len(iou_per_video),
+        "F1": f1_score,
+        "precision_score": pr_score,
+        "re_score": re_score,
+        "AO": average_overlap,
+        "result_dir": result_dir,
+        "MODEL_WEIGHTS": cfg.MODEL.WEIGHTS,
+    }

EgoTracks/tracking/metrics/lt_tracking_metrics.py

@@ -0,0 +1,94 @@
+import math
+from typing import Iterable, List
+
+import numpy as np
+from tqdm import tqdm
+from tracking.metrics.miou import compute_overlaps
+
+# Copied from https://github.com/votchallenge/toolkit/blob/master/vot/analysis/tpr.py
+def determine_thresholds(scores: Iterable[float], resolution: int) -> List[float]:
+    scores = [
+        score for score in scores if not math.isnan(score)
+    ]  # and not score is None]
+    scores = sorted(scores, reverse=True)
+
+    if len(scores) > resolution - 2:
+        delta = math.floor(len(scores) / (resolution - 2))
+        idxs = np.round(
+            np.linspace(delta, len(scores) - delta, num=resolution - 2)
+        ).astype(np.int)
+        thresholds = [scores[idx] for idx in idxs]
+    else:
+        thresholds = scores
+
+    thresholds.insert(0, math.inf)
+    thresholds.insert(len(thresholds), -math.inf)
+
+    return thresholds
+
+
+def compute_tpr_curves(confidences, overlaps, thresholds, n_visible):
+    precision = len(thresholds) * [float(0)]
+    recall = len(thresholds) * [float(0)]
+    for i, threshold in enumerate(thresholds):
+
+        subset = confidences >= threshold
+
+        if np.sum(subset) == 0:
+            precision[i] = 1
+            recall[i] = 0
+        else:
+            precision[i] = np.mean(overlaps[subset])
+            recall[i] = np.sum(overlaps[subset]) / n_visible
+
+    return precision, recall
+
+
+def compute_precision_and_recall(
+    all_pred_scores: List[List], all_pred_bboxes: List[List], all_gt_bboxes: List[List]
+):
+    """
+    Compute the precision and recall for the entire dataset. Per score/bbox for each frame of each video.
+    If there is no gt bbox for that frame, then it is set to None.
+
+    all_pred_scores: confidence score for all tracking trajectories in the dataset.
+    all_pred_bboxes: predicted bbox for all tracking trajectories in the dataset.
+    all_gt_bboxes: gt bbox for all tracking trajectories in the dataset.
+    """
+    all_overlaps = [
+        compute_overlaps(pred_bboxes, gt_bbboxes)
+        for pred_bboxes, gt_bbboxes in zip(all_pred_bboxes, all_gt_bboxes)
+    ]
+    resolution = 100
+
+    thresholds = determine_thresholds(
+        [s for conf in all_pred_scores for s in conf], resolution
+    )
+
+    precision = len(thresholds) * [float(0)]
+    recall = len(thresholds) * [float(0)]
+
+    for i in tqdm(range(len(all_gt_bboxes)), total=len(all_gt_bboxes)):
+        confidences = np.array(all_pred_scores[i])
+        overlaps = np.array(all_overlaps[i])
+        gt_bboxes = all_gt_bboxes[i]
+        n_visible = len([b for b in gt_bboxes if b is not None])
+
+        pr, re = compute_tpr_curves(confidences, overlaps, thresholds, n_visible)
+        for j in range(len(thresholds)):
+            precision[j] += pr[j]
+            recall[j] += re[j]
+
+    precision = [pr / len(all_gt_bboxes) for pr in precision]
+    recall = [re / len(all_gt_bboxes) for re in recall]
+
+    return precision, recall
+
+
+def compute_f_score(precision, recall):
+    f_score = [
+        2 * precision[i] * recall[i] / (precision[i] + recall[i])
+        for i in range(len(precision))
+    ]
+    n = np.argmax(f_score)
+    return (f_score[n], precision[n], recall[n])