# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Vehicle Detection: 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

VEHICLE_CLASS_IDS = {2: "car", 3: "motorcycle", 5: "bus", 7: "truck"}
CONF_THRESHOLD = 0.4
# Select the model: "yolo26n" (default) or "detr-resnet-50".
MODEL = "yolo26n"
INPUT_SIZE = 800 if MODEL.startswith("detr") else 640

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

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

image = cv2.imread("test.jpg")
h0, w0 = image.shape[:2]

blob = cv2.resize(image, (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])[compiled.output(0)][0]
mask = (output[:, 4] >= CONF_THRESHOLD) & np.isin(output[:, 5].astype(int), list(VEHICLE_CLASS_IDS.keys()))
dets = output[mask]

sx, sy = w0 / INPUT_SIZE, h0 / INPUT_SIZE
vehicle_count = len(dets)
print(f"Detected vehicles: {vehicle_count}")

colors = {"car": (0, 255, 0), "motorcycle": (255, 0, 0),
          "bus": (0, 165, 255), "truck": (0, 0, 255)}
for det in dets:
    x1 = int(det[0] * sx)
    y1 = int(det[1] * sy)
    x2 = int(det[2] * sx)
    y2 = int(det[3] * sy)
    cid = int(det[5])
    conf = float(det[4])
    name = VEHICLE_CLASS_IDS[cid]
    label = f"{name} {conf:.2f}"
    color = colors[name]
    cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
    cv2.putText(image, label, (x1, y1 - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
    print(f"  {label} at ({x1},{y1})-({x2},{y2})")

cv2.imwrite("output_openvino.jpg", image)