import time
|
|
import cv2
|
import numpy as np
|
from openvino.runtime import Core
|
|
|
|
def preprocess_image(image, input_size=(640, 640)):
|
|
img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转换为 RGB
|
img_resized = cv2.resize(img, input_size) # 调整大小
|
img_normalized = img_resized / 255.0 # 归一化
|
img_transposed = np.transpose(img_normalized, (2, 0, 1)) # HWC -> CHW
|
img_batch = np.expand_dims(img_transposed, axis=0).astype(np.float32) # 增加 batch 维度
|
return img_batch, img_resized
|
|
def postprocess(output, conf_threshold=0.5, iou_threshold=0.4):
|
# 解析输出
|
boxes = []
|
scores = []
|
class_ids = []
|
|
for detection in output[0]:
|
confidence = detection[4]
|
if confidence > conf_threshold:
|
class_id = np.argmax(detection[5:])
|
cx, cy, w, h = detection[:4] # 中心点和宽高
|
x_min = int((cx - w / 2) * original_image.shape[1]) # 转换回原始图像尺寸
|
y_min = int((cy - h / 2) * original_image.shape[0])
|
width = int(w * original_image.shape[1])
|
height = int(h * original_image.shape[0])
|
|
boxes.append([x_min, y_min, width, height])
|
scores.append(confidence)
|
class_ids.append(class_id)
|
|
# 非极大值抑制 (NMS)
|
indices = cv2.dnn.NMSBoxes(boxes, scores, conf_threshold, iou_threshold)
|
|
final_boxes = []
|
for idx in indices:
|
box = boxes[idx]
|
final_boxes.append({
|
"box": box,
|
"score": scores[idx],
|
"class_id": class_ids[idx]
|
})
|
|
return final_boxes
|
|
|
def draw_detections(image, detections):
|
for det in detections:
|
x_min, y_min, width, height = det["box"]
|
class_id = det["class_id"]
|
score = det["score"]
|
|
# 确保坐标和尺寸是整数
|
x_min = int(x_min)
|
y_min = int(y_min)
|
width = int(width)
|
height = int(height)
|
|
# 绘制边界框
|
cv2.rectangle(image, (x_min, y_min), (x_min + width, y_min + height), (0, 255, 0), 2)
|
|
# 显示类别和置信度
|
label = f"Class {class_id}: {score:.2f}"
|
cv2.putText(image, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
|
|
return image
|
|
|
if __name__ == "__main__":
|
# 初始化 OpenVINO Runtime
|
core = Core()
|
|
# 加载模型
|
model_path = "model/safe_det/best.xml" # 替换为你的模型路径
|
model = core.read_model(model=model_path)
|
compiled_model = core.compile_model(model=model, device_name="CPU") # 设备可以是 "CPU", "GPU", "MYRIAD" 等
|
|
# 获取输入和输出层
|
input_layer = compiled_model.input(0)
|
output_layer = compiled_model.output(0)
|
|
# 打印输入输出信息
|
print(f"Input shape: {input_layer.shape}")
|
print(f"Output shape: {output_layer.shape}")
|
|
# 预处理图像，改为从摄像头获取
|
cap = cv2.VideoCapture(0)
|
|
# 初始化变量
|
start_time = time.time() # 记录开始时间
|
frame_count = 0 # 帧计数器
|
|
while True:
|
ret, frame = cap.read()
|
if not ret:
|
break
|
|
input_data, original_image = preprocess_image(frame)
|
|
# 推理
|
results = compiled_model([input_data])[output_layer]
|
# 输出结果
|
print(results)
|
detections = postprocess(results)
|
# 打印检测结果
|
for det in detections:
|
print(det)
|
draw_img = draw_detections(frame, detections)
|
|
# 计算FPS
|
frame_count += 1
|
elapsed_time = time.time() - start_time
|
|
fps = frame_count / elapsed_time
|
# start_time = time.time() # 重置开始时间
|
# frame_count = 0 # 重置帧计数器
|
|
# 显示FPS
|
cv2.putText(draw_img, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
|
|
cv2.imshow("Detections", draw_img)
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
break
|
