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


img = cv2.imread("../data/zidane.jpg")

model = YOLO("yolov8s-seg.pt")

yolo_detect_result = model.predict([img], conf=0.5)[0]
detections = sv.Detections.from_yolov8(yolo_detect_result)

annotated_img = img.copy()

cat_maps = {0: "person", 27: "tie"}
labels = [f"{cat_maps[class_id]}" for _, _, _, class_id, _ in detections]

tracker_ids = np.array([1, 2, 3], dtype=int)
detections.tracker_id = tracker_ids  # Add tracker id to visualize color by tracker id


color_by_track = False   # Test with True and False
font_path = "font_path" # Test with actual path and None
skip_label = False  # Test with True and False

detection_annotator = sv.DetectionAnnotator(color_by_track=color_by_track, skip_label=skip_label)
segmentation_annotator = sv.SegmentationAnnotator(color_by_track=color_by_track, skip_label=skip_label)

box_annotator = sv.BoxAnnotator(color_by_track=color_by_track)
label_annotator = sv.LabelAnnotator(color_by_track=color_by_track)
mask_annotator = sv.MaskAnnotator(color_by_track=color_by_track)

boxmask_annotator = sv.BoxMaskAnnotator(color_by_track=color_by_track)
cornered_box_annotator = sv.CorneredBoxAnotator(color_by_track=color_by_track)
pil_label_annotator = sv.PillowLabelAnnotator(color_by_track=color_by_track, font=font_path )


annotators = {"label": label_annotator, "pill-label": pil_label_annotator, "box": box_annotator, "cornered_box": cornered_box_annotator, "box-mask": boxmask_annotator, "mask": mask_annotator, "detection":
              detection_annotator, "segmentation": segmentation_annotator}


for key, value in annotators.items():
    if key in ["label", "detection", "segmentation", "pill-label"]:
        annotated_img = annotators[key].annotate(scene=img.copy(), detections=detections, labels=labels)
    elif key == "pill-label":
        annotated_img = pil_label_annotator.annotate(scene=img.copy(), detections=detections, labels=labels)
    else:
        annotated_img = annotators[key].annotate(scene=img.copy(), detections=detections)
    cv2.imshow(key, annotated_img)
    cv2.waitKey(0)