File size: 3,539 Bytes
0577537 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | # 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")
|