import cv2
from ultralytics import YOLO

model = YOLO("C:\\Users\\PC\\Desktop\\GazCounter\\best (10).pt")
cap = cv2.VideoCapture("C:\\Users\\PC\\Desktop\\GazCounter\\WhatsApp Video 2025-10-02 at 11.29.10 (2).mp4")

count_per_brand = {}
detected_ids = set()  

while True:
    ret, frame = cap.read()
    if not ret:
        break

    results = model.track(frame, persist=True, conf=0.4)

    for r in results:
        if r.boxes is None:
            continue

        for box in r.boxes:
            #  Vérifier que l'ID du tracker existe
            if box.id is None:
                continue

            track_id = int(box.id[0])  #  Convertir tensor → int
            cls_id   = int(box.cls[0])
            cls_name = r.names[cls_id]
            conf     = float(box.conf[0])

            #  Dessiner la bounding box
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(
                frame,
                f"{cls_name} #{track_id} ({conf:.0%})",
                (x1, y1 - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2
            )

            #  Comptage UNIQUE par ID de tracking
            if track_id not in detected_ids:
                detected_ids.add(track_id)
                count_per_brand[cls_name] = count_per_brand.get(cls_name, 0) + 1

    overlay = frame.copy()
    cv2.rectangle(overlay, (5, 5), (300, 30 + len(count_per_brand) * 28), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.4, frame, 0.6, 0, frame)  # fond semi-transparent

    for i, (brand, count) in enumerate(sorted(count_per_brand.items())):
        cv2.putText(
            frame,
            f"{brand}: {count}",
            (10, 28 + i * 28),
            cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 255, 255), 2
        )

    total = sum(count_per_brand.values())
    cv2.putText(
        frame,
        f"TOTAL: {total} bouteilles",
        (10, frame.shape[0] - 10),
        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2
    )

    cv2.imshow("GazCounter", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

for brand, count in sorted(count_per_brand.items()):
    print(f"  {brand}: {count} bouteille(s)")
print(f"  TOTAL: {sum(count_per_brand.values())} bouteilles")