openvino_sample.py · Intel/object-classification at main

object-classification / openvino_sample.py

Update native RC1 samples and supported device documentation

0577537 verified about 10 hours ago

3.54 kB

# 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")

	# 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")