import os
import numpy as np
import supervision as sv
from ultralytics import YOLO

# Download dataset
if not os.path.exists("/tmp/rf_animals"):
    !wget https://universe.roboflow.com/ds/1LLwpXz2td?key=8JnJML5YF6 -O /tmp/rf_animals.zip
    !unzip /tmp/dataset.zip -d /tmp/rf_animals

# Load dataset
dataset = sv.DetectionDataset.from_yolo("/tmp/rf_animals/train/images", "/tmp/rf_animals/train/labels", "/tmp/rf_animals/data.yaml")

# Inference
model = YOLO("yolov8s")
targets, detections = [], []
for image_path, image, target in dataset:
    targets.append(target)
    
    prediction = model(image, verbose=False)[0]
    detection = sv.Detections.from_ultralytics(prediction)
    detection = detection[np.isin(detection['class_name'], dataset.classes)]
    detection.class_id = np.array([dataset.classes.index(class_name) for class_name in detection['class_name']])
    detections.append(detection)
    
# Method #1
mAP = sv.metrics.MeanAveragePrecision().update(detections, targets).compute()
print(f"mAP50: {mAP.map50:.4f}")

# Method #2
mAP = sv.MeanAveragePrecision.from_detections(detections, targets)
print(f"mAP50: {mAP.map50:.4f}")