src/test3_withdeepsort.py

import cv2
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort
import os
from datetime import datetime

model = YOLO("C:\\Users\\PC\\Desktop\\GazCounter\\models\\best (10).pt")

tracker = DeepSort(
    max_age=30,         
    n_init=3,            
    max_iou_distance=0.7,
    embedder="mobilenet",
    half=True,
    embedder_gpu=True,
)

video_path = "C:\\Users\\PC\\Desktop\\GazCounter\\data\\WhatsApp Video 2025-10-02 at 11.29.11.mp4"
cap = cv2.VideoCapture(video_path)

# ── OUTPUT SETUP ────────────────────────────────────────
width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps    = int(cap.get(cv2.CAP_PROP_FPS))

output_dir   = "C:\\Users\\PC\\Desktop\\GazCounter\\results"
os.makedirs(output_dir, exist_ok=True)
timestamp    = datetime.now().strftime("%Y%m%d_%H%M%S")
output_video = os.path.join(output_dir, f"result_{timestamp}.mp4")
output_txt   = os.path.join(output_dir, f"result_{timestamp}.txt")

writer = cv2.VideoWriter(
    output_video,
    cv2.VideoWriter_fourcc(*"mp4v"),
    fps, (width, height)
)

# ── STATE ───────────────────────────────────────────────
count_per_brand = {}
detected_ids    = set()
id_to_box       = {}
id_to_class     = {}

IOU_THRESHOLD = 0.5

def iou(boxA, boxB):
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    inter = max(0, xB - xA) * max(0, yB - yA)
    if inter == 0:
        return 0.0
    areaA = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
    areaB = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
    return inter / float(areaA + areaB - inter)

# ── MAIN LOOP ───────────────────────────────────────────
while True:
    ret, frame = cap.read()
    if not ret:
        break

    # ── Step 1: YOLO detect (no tracking) ───────────────
    results = model.predict(frame, conf=0.4, verbose=False)

    detections = []
    cls_map    = {}  # index → class name

    for r in results:
        if r.boxes is None:
            continue
        for i, box in enumerate(r.boxes):
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            conf     = float(box.conf[0])
            cls_id   = int(box.cls[0])
            cls_name = r.names[cls_id]

            # DeepSORT format: ([x1,y1,w,h], conf, class_name)
            w = x2 - x1
            h = y2 - y1
            detections.append(([x1, y1, w, h], conf, cls_name))

    # ── Step 2: DeepSORT update ──────────────────────────
    tracks = tracker.update_tracks(detections, frame=frame)

    for track in tracks:
        if not track.is_confirmed():
            continue

        track_id = track.track_id
        cls_name = track.get_det_class()
        ltrb     = track.to_ltrb()  # (x1, y1, x2, y2)
        x1, y1, x2, y2 = map(int, ltrb)
        curr_box = (x1, y1, x2, y2)

        if cls_name is None:
            continue

        # Draw
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            frame,
            f"{cls_name} #{track_id}",
            (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2
        )

        # ── Unique count with IoU dedup ──────────────────
        if track_id not in detected_ids:
            is_duplicate = False
            for known_id, known_box in id_to_box.items():
                if known_id not in detected_ids:
                    continue
                if iou(curr_box, known_box) > IOU_THRESHOLD:
                    is_duplicate = True
                    break

            if not is_duplicate:
                detected_ids.add(track_id)
                count_per_brand[cls_name] = count_per_brand.get(cls_name, 0) + 1

        id_to_box[track_id]   = curr_box
        id_to_class[track_id] = cls_name

    # ── Overlay panel ────────────────────────────────────
    overlay = frame.copy()
    panel_h = 30 + len(count_per_brand) * 28
    cv2.rectangle(overlay, (5, 5), (300, panel_h), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.4, frame, 0.6, 0, frame)

    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
    )

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

cap.release()
writer.release()
cv2.destroyAllWindows()

# ── Save report ──────────────────────────────────────────
with open(output_txt, "w") as f:
    f.write(f"GazCounter Results - {timestamp}\n")
    f.write("=" * 40 + "\n")
    for brand, count in sorted(count_per_brand.items()):
        f.write(f"  {brand}: {count} bouteille(s)\n")
    f.write("=" * 40 + "\n")
    f.write(f"  TOTAL: {sum(count_per_brand.values())} bouteilles\n")

print("\n" + "=" * 40)
print(f"Video  → {output_video}")
print(f"Report → {output_txt}")
print("=" * 40)
for brand, count in sorted(count_per_brand.items()):
    print(f"  {brand}: {count} bouteille(s)")
print(f"  TOTAL: {sum(count_per_brand.values())} bouteilles")