# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Object Classification: OpenVINO Runtime sample.
# See README.md for prerequisites and inputs. Run: python3 openvino_sample.py

import cv2
import numpy as np
import openvino as ov

CONF_THRESHOLD = 0.4
# Select the model: "yolo26n" (default) or "detr-resnet-50".
MODEL = "yolo26n"
INPUT_SIZE = 800 if MODEL.startswith("detr") else 640

# Map the traffic-relevant COCO class ids into higher-level city categories.
# People and Vehicles are the two categories tracked for situational awareness.
CATEGORY_BY_CLASS_ID = {
    0: "People",      # person
    1: "Vehicles",    # bicycle
    2: "Vehicles",    # car
    3: "Vehicles",    # motorcycle
    5: "Vehicles",    # bus
    6: "Vehicles",    # train
    7: "Vehicles",    # truck
}
# BGR overlay colors for each category.
CATEGORY_COLORS = {
    "People": (0, 200, 0),
    "Vehicles": (255, 128, 0),
}

core = ov.Core()
model = core.read_model(f"{MODEL}_openvino_model/{MODEL}.xml")

# YOLO26 embeds the 80 COCO class names in rt_info. Ultralytics separates
# multi-word names with underscores (e.g. "traffic_light"), so restore spaces.
COCO_NAMES = [
    name.replace("_", " ")
    for name in model.get_rt_info()["model_info"]["labels"].value.split()
]

# Change device to "GPU" or "NPU" to run on integrated GPU or NPU.
compiled = core.compile_model(model, "CPU")
output_port = compiled.output(0)

cap = cv2.VideoCapture("test_video.mp4")
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
writer = cv2.VideoWriter(
    "output_openvino.mp4", cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
)

totals = {"People": 0, "Vehicles": 0}
frame_idx = 0
while True:
    ok, frame = cap.read()
    if not ok:
        break
    frame_idx += 1

    blob = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE))
    blob = cv2.cvtColor(blob, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    blob = blob.transpose(2, 0, 1)[np.newaxis, ...]  # NCHW

    # End-to-end output: [1, N, 6] = [x1, y1, x2, y2, confidence, class_id].
    # (N = 300 for YOLO26, 100 for DETR)
    output = compiled([blob])[output_port][0]

    sx, sy = width / INPUT_SIZE, height / INPUT_SIZE
    counts = {"People": 0, "Vehicles": 0}
    for x1, y1, x2, y2, conf, class_id in output:
        if conf < CONF_THRESHOLD:
            continue
        category = CATEGORY_BY_CLASS_ID.get(int(class_id))
        if category is None:
            continue  # not a traffic-relevant object
        counts[category] += 1
        totals[category] += 1
        color = CATEGORY_COLORS[category]
        px1, py1 = int(x1 * sx), int(y1 * sy)
        px2, py2 = int(x2 * sx), int(y2 * sy)
        label = f"{category}: {COCO_NAMES[int(class_id)]} {conf:.2f}"
        cv2.rectangle(frame, (px1, py1), (px2, py2), color, 2)
        cv2.putText(frame, label, (px1, py1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 2.0, color, 2)

    # Overlay the per-category counts for this frame.
    banner = f"People: {counts['People']}  Vehicles: {counts['Vehicles']}"
    cv2.rectangle(frame, (0, 0), (width, 60), (0, 0, 0), -1)
    cv2.putText(frame, banner, (15, 45),
                cv2.FONT_HERSHEY_SIMPLEX, 2.0, (255, 255, 255), 2)

    if frame_idx % 30 == 0:
        print(f"frame {frame_idx}: {banner}", flush=True)

    writer.write(frame)

cap.release()
writer.release()
print(f"Summary: People={totals['People']} Vehicles={totals['Vehicles']}")
print("Saved: output_openvino.mp4")