File size: 2,266 Bytes
0d4136a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # 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)
|