你好，我按照您发的链接下载了.bin文件也在RDK上跑通了，不过我是调用rtsp的视频流作为输入，目前存在几个问题：1.画框的偏移，我调用的函数是官方的，2.还是我原帖的问题，按照论坛里超哥的方法把onnx转为bin这一步里面图片验证集的参数规格能否具体给出，我的报错是2025-08-14 14:32:32,098 ERROR cannot reshape array of size 1228832 into shape (1,3,640,640)，但是我的图片已经是float32的640X640我还给他转成。npy格式了，下面是我的代码。
#!/usr/bin/env python3
# Copyright (c) 2024，WuChao D-Robotics.
# 注意：此程序在 RDK 板端运行
# 改动：支持 RTSP 视频流实时推理
import os
import cv2
import numpy as np
from time import time
import argparse
import logging
# 日志配置
logging.basicConfig(
level=logging.INFO,
format='[%(name)s] [%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s',
datefmt='%H:%M:%S')
logger = logging.getLogger("RDK_YOLO_RTSP")
# ---------- 依赖检查 ----------
try:
from scipy.special import softmax
except ImportError:
logger.warning("scipy 未安装，正在自动安装...")
os.system("pip3 install scipy -i https://pypi.tuna.tsinghua.edu.cn/simple")
from scipy.special import softmax
try:
from hobot_dnn import pyeasy_dnn as dnn
except ImportError:
logger.error("未找到 hobot_dnn，请使用板端系统 Python3 环境")
exit(1)
# ---------- 主程序 ----------
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model-path', type=str,
default='/home/sunrise/rdk_model_zoo/demos/Vision/ultralytics_YOLO/py/yolo13_detect.bin',
help='量化 *.bin 模型路径')
parser.add_argument('--rtsp', type=str,
default='rtsp://192.168.1.120:8554/test',
help='RTSP 流地址')
parser.add_argument('--classes-num', type=int, default=80)
parser.add_argument('--nms-thres', type=float, default=0.7)
parser.add_argument('--score-thres', type=float, default=0.25)
parser.add_argument('--reg', type=int, default=16)
parser.add_argument('--strides', type=lambda s: list(map(int, s.split(','))),
default=[8, 16, 32])
opt = parser.parse_args()
logger.info(opt)
# 初始化模型
model = Ultralytics_YOLO_Detect_Bayese_YUV420SP(
model_path=opt.model_path,
classes_num=opt.classes_num,
nms_thres=opt.nms_thres,
score_thres=opt.score_thres,
reg=opt.reg,
strides=opt.strides
)
# 打开 RTSP 流
cap = cv2.VideoCapture(opt.rtsp, cv2.CAP_FFMPEG)
if not cap.isOpened():
logger.error("无法打开 RTSP 流: {}".format(opt.rtsp))
exit(1)
logger.info("已连接 RTSP: {}".format(opt.rtsp))
# 降低缓存，减少延迟
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
# 主循环
while True:
ret, frame = cap.read()
if not ret:
logger.warning("读取帧失败，重试中...")
continue
# 推理
input_tensor = model.preprocess_yuv420sp(frame)
outputs = model.c2numpy(model.forward(input_tensor))
results = model.postProcess(outputs)
# 画框
for class_id, score, x1, y1, x2, y2 in results:
draw_detection(frame, (x1, y1, x2, y2), score, class_id)
# 本地显示（需要桌面环境）
cv2.imshow("RTSP YOLO Detection - RDK", frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
# ---------- 以下与原始文件一致，无改动 ----------
class Ultralytics_YOLO_Detect_Bayese_YUV420SP:
def __init__(self, model_path, classes_num, nms_thres, score_thres, reg, strides):
begin_time = time()
self.quantize_model = dnn.load(model_path)
logger.debug("Load model time = %.2f ms" % (1000 * (time() - begin_time)))
self.REG = reg
self.CLASSES_NUM = classes_num
self.SCORE_THRESHOLD = score_thres
self.NMS_THRESHOLD = nms_thres
self.CONF_THRES_RAW = -np.log(1 / self.SCORE_THRESHOLD - 1)
self.input_H, self.input_W = self.quantize_model[0].inputs[0].properties.shape[2:4]
self.strides = strides
self.weights_static = np.arange(reg, dtype=np.float32)[np.newaxis, np.newaxis, :]
self.grids = []
for stride in self.strides:
grid_h, grid_w = self.input_H // stride, self.input_W // stride
yv, xv = np.meshgrid(np.arange(grid_h), np.arange(grid_w), indexing='ij')
grid = np.stack([yv, xv], axis=-1).reshape(-1, 2) + 0.5
self.grids.append(grid)
def preprocess_yuv420sp(self, img):
self.img_h, self.img_w = img.shape[:2]
scale = min(self.input_H / self.img_h, self.input_W / self.img_w)
new_w, new_h = int(self.img_w * scale), int(self.img_h * scale)
self.x_shift = (self.input_W - new_w) // 2
self.y_shift = (self.input_H - new_h) // 2
resized = cv2.resize(img, (new_w, new_h))
letter = cv2.copyMakeBorder(resized,
self.y_shift, self.input_H - new_h - self.y_shift,
self.x_shift, self.input_W - new_w - self.x_shift,
cv2.BORDER_CONSTANT, value=(127, 127, 127))
return self.bgr2nv12(letter)
def bgr2nv12(self, bgr_img):
height, width = bgr_img.shape[:2]
yuv420p = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2YUV_I420)
y = yuv420p[:height, :]
uv = yuv420p[height:, :]
uv = uv.reshape((2, height // 2, width // 2)).transpose(1, 2, 0).reshape(height // 2, width)
nv12 = np.empty(height * 3 // 2 * width, dtype=np.uint8)
nv12[:height * width] = y.flatten()
nv12[height * width:] = uv.flatten()
return nv12
def forward(self, input_tensor):
return self.quantize_model[0].forward(input_tensor)
def c2numpy(self, outputs):
return [t.buffer for t in outputs]
def postProcess(self, outputs):
clses = [outputs[0].reshape(-1, self.CLASSES_NUM),
outputs[2].reshape(-1, self.CLASSES_NUM),
outputs[4].reshape(-1, self.CLASSES_NUM)]
bboxes = [outputs[1].reshape(-1, self.REG * 4),
outputs[3].reshape(-1, self.REG * 4),
outputs[5].reshape(-1, self.REG * 4)]
dbboxes, ids, scores = [], [], []
for cls, bbox, stride, grid in zip(clses, bboxes, self.strides, self.grids):
max_scores = np.max(cls, axis=1)
mask = max_scores >= self.CONF_THRES_RAW
cls_sel, bbox_sel = cls[mask], bbox[mask]
if cls_sel.size == 0:
continue
ids.append(np.argmax(cls_sel, axis=1))
scores.append(1 / (1 + np.exp(-max_scores[mask])))
ltrb = softmax(bbox_sel.reshape(-1, 4, self.REG), axis=2)
ltrb = np.sum(ltrb * self.weights_static, axis=2)
grid_sel = grid[np.where(mask)[0]]
xy1 = (grid_sel - ltrb[:, :2]) * stride
xy2 = (grid_sel + ltrb[:, 2:]) * stride
dbboxes.append(np.hstack([xy1, xy2]))
if not dbboxes:
return []
dbboxes = np.vstack(dbboxes)
scores = np.hstack(scores)
ids = np.hstack(ids)
xywh = np.hstack([(dbboxes[:, :2] + dbboxes[:, 2:]) / 2,
dbboxes[:, 2:] - dbboxes[:, :2]])
results = []
for i in range(self.CLASSES_NUM):
idx = ids == i
if not idx.any():
continue
boxes = xywh[idx]
sc = scores[idx]
indices = cv2.dnn.NMSBoxes(boxes, sc, self.SCORE_THRESHOLD, self.NMS_THRESHOLD)
for j in indices:
x1, y1, x2, y2 = dbboxes[idx][j]
x1 = max(0, min(int((x1 - self.x_shift) / (self.input_W / self.img_w)), self.img_w))
y1 = max(0, min(int((y1 - self.y_shift) / (self.input_H / self.img_h)), self.img_h))
x2 = max(0, min(int((x2 - self.x_shift) / (self.input_W / self.img_w)), self.img_w))
y2 = max(0, min(int((y2 - self.y_shift) / (self.input_H / self.img_h)), self.img_h))
results.append((i, sc[j], x1, y1, x2, y2))
return results
# ---------- 画图 ----------
coco_names = [
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
"wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
"broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed",
"dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
"toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
]
#my
# rdk_colors = [
# (56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255),
# (49, 210, 207), (10, 249, 72), (23, 204, 146), (134, 219, 61),
# (52, 147, 26), (187, 212, 0), (168, 153, 44), (255, 194, 0),
# (147, 69, 52), (255, 115, 100), (236, 24, 0), (255, 56, 132),
# (133, 0, 82), (255, 56, 203), (200, 149, 255), (199, 55, 255)
# ]
#offical
rdk_colors = [
(56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255),(49, 210, 207), (10, 249, 72), (23, 204, 146), (134, 219, 61),
(52, 147, 26), (187, 212, 0), (168, 153, 44), (255, 194, 0),(147, 69, 52), (255, 115, 100), (236, 24, 0), (255, 56, 132),
(133, 0, 82), (255, 56, 203), (200, 149, 255), (199, 55, 255)]
#my
# def draw_detection(img, bbox, score, class_id) -> None:
# x1, y1, x2, y2 = map(int, bbox) # 确保坐标是 int
# color = tuple(map(int, rdk_colors[class_id % len(rdk_colors)])) # 确保颜色是 int 元组
# cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
# label = f"{coco_names[class_id]}:{score:.2f}"
# (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
# cv2.rectangle(img, (x1, y1 - h - 4), (x1 + w, y1), color, -1)
# cv2.putText(img, label, (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
#offical
def draw_detection(img, bbox, score, class_id) -> None:
"""
Draws a detection bounding box and label on the image.
Parameters:
img (np.array): The input image.
bbox (tuple[int, int, int, int]): A tuple containing the bounding box coordinates (x1, y1, x2, y2).
score (float): The detection score of the object.
class_id (int): The class ID of the detected object.
"""
x1, y1, x2, y2 = bbox
color = rdk_colors[class_id%20]
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
label = f"{coco_names[class_id]}: {score:.2f}"
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
label_x, label_y = x1, y1 - 10 if y1 - 10 > label_height else y1 + 10
cv2.rectangle(
img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
)
cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
if __name__ == "__main__":
main()
