import torch
import numpy as np
import fiftyone as fo
from box_metrics import BoxMetrics
from seametrics.fo_utils.utils import fo_to_payload 
from tqdm import tqdm

tags = ["WHALES"]
dataset_name = "SENTRY_VIDEOS_DATASET_QA"
model = "cerulean-level-17_11_2023_RL_SPLIT_ep147_CNN"
det_gt_field = "ground_truth_det"

dataset = fo.load_dataset(dataset_name)
dataset_view = fo.load_dataset(dataset_name).match_tags(tags) if tags else fo.load_dataset(dataset_name)
sequences = dataset_view.distinct("sequence")

bbox_metric = BoxMetrics(max_iou=0.01)
payload = fo_to_payload(dataset = dataset_name,
                        gt_field = det_gt_field,
                        models = [model],
                        tracking_mode = True,
                        sequence_list = sequences)
print(payload)
bbox_metric.add_payload(payload)
result = bbox_metric.compute()
print(result)