diff --git a/EgoTracks/tools/eval_datasets/eval_ego4d_lt_tracking.py b/EgoTracks/tools/eval_datasets/eval_ego4d_lt_tracking.py index 95ce4ab..f05e0f3 100644 --- a/EgoTracks/tools/eval_datasets/eval_ego4d_lt_tracking.py +++ b/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, + } \ No newline at end of file diff --git a/EgoTracks/tracking/metrics/lt_tracking_metrics.py b/EgoTracks/tracking/metrics/lt_tracking_metrics.py new file mode 100644 index 0000000..7f65acc --- /dev/null +++ b/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])