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

from collections import Counter

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

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

# YOLO26 embeds the 80 COCO class names in rt_info -- read them instead of
# hardcoding the list. Ultralytics separates multi-word names with
# underscores (e.g. "traffic_light"), so restore spaces for display.
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")

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]
dets = output[output[:, 4] >= CONF_THRESHOLD]

sx, sy = w0 / INPUT_SIZE, h0 / INPUT_SIZE
counts = Counter(COCO_NAMES[int(d[5])] for d in dets)

print(f"Total objects: {len(dets)}")
print("Object counts:")
for name, n in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0])):
    print(f"  {name}: {n}")

colors = np.random.RandomState(42).randint(0, 255, (80, 3)).tolist()
for det in dets:
    x1, y1, x2, y2 = (int(det[0] * sx), int(det[1] * sy),
                      int(det[2] * sx), int(det[3] * sy))
    cid = int(det[5])
    cv2.rectangle(image, (x1, y1), (x2, y2), colors[cid], 2)
    cv2.putText(image, COCO_NAMES[cid], (x1, y1 - 5),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, colors[cid], 2)

summary = ", ".join(f"{n} {name}" for name, n in counts.items())
cv2.putText(image, summary[:60], (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
cv2.imwrite("output_openvino.jpg", image)