如何结合OpenCV使用YOLOv5导出的TFLite模型并可视化检测结果

你拿到的output_data形状为(1, num_anchors, 5 + num_classes)，各维度含义如下：

• 第0维：batch维度，固定为1，对应单次推理1张图像
• 第1维：所有检测层预设的anchor框总数，输入分辨率为256256时，这个值固定为4032（对应3个检测步长下的anchor总数：(3232 + 1616 + 88)*3）
• 第2维：每个anchor对应的预测结果，按顺序分别是：
  • 前4位：归一化格式的预测框坐标(x_center, y_center, width, height)，所有值都在0-1区间，是相对于模型输入分辨率（256*256）的比例值，不是绝对像素坐标
  • 第5位：目标置信度，代表当前anchor位置存在目标的概率，取值0-1
  • 后续num_classes位：类别置信度，每一位对应一个类别，代表当前框属于该类别的概率，取值0-1

注意：你当前的推理代码存在预处理错误：YOLOv5模型要求输入像素值归一化到0-1区间，你直接传入了0-255范围的像素值，会导致推理结果完全不准。

2.1 核心处理流程 #

你之前的代码缺少输出解析、NMS去重步骤，且坐标映射逻辑写反了宽高参数，按以下步骤处理即可：

• 修正输入预处理：将像素值除以255做归一化，OpenCV读取的BGR格式帧要先转RGB再输入模型
• 压缩输出维度：去掉第0位的batch维度，得到形状为(4032, 5+num_classes)的结果数组
• 计算每个框的实际类别得分：得分 = 目标置信度 * 类别置信度，取每个框得分最高的类别作为预测类别
• 阈值过滤：剔除得分低于设定阈值（一般取0.25）的无效框
• 坐标转换：把归一化的xywh格式坐标，转换成绘制用的xyxy（左上角x、左上角y、右下角x、右下角y）绝对像素坐标
• 非极大值抑制（NMS）：对同一目标的重复检测框做去重，保留得分最高的框
• 结果绘制：把最终保留的框、类别、得分画到图像上

2.2 完整参考代码 #

单图推理版本 #

import numpy as np
import tensorflow as tf
import cv2

# 配置参数
MODEL_PATH = "yolov5s-fp16.tflite"
INPUT_SIZE = (256, 256)
CONF_THRESH = 0.25
IOU_THRESH = 0.45
# 替换成你自己的类别名列表，COCO数据集默认是80类
CLASS_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"]

# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path=MODEL_PATH)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

def nms(boxes, scores, iou_threshold):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= iou_threshold)[0]
        order = order[inds + 1]
    return keep

def detect(img):
    # 预处理
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img_rgb, INPUT_SIZE)
    input_data = img_resized.astype(np.float32) / 255.0
    input_data = np.expand_dims(input_data, 0)
    
    # 推理
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])[0]
    
    # 解析输出
    boxes = []
    scores = []
    classes = []
    img_h, img_w = img.shape[:2]
    for pred in output_data:
        xc, yc, w, h, obj_conf = pred[:5]
        cls_conf = pred[5:]
        cls_id = np.argmax(cls_conf)
        score = float(obj_conf * cls_conf[cls_id])
        if score < CONF_THRESH:
            continue
        # xywh转xyxy绝对坐标
        x1 = int((xc - w/2) * img_w)
        y1 = int((yc - h/2) * img_h)
        x2 = int((xc + w/2) * img_w)
        y2 = int((yc + h/2) * img_h)
        # 边界裁剪
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(img_w, x2)
        y2 = min(img_h, y2)
        boxes.append([x1, y1, x2, y2])
        scores.append(score)
        classes.append(cls_id)
    
    # NMS去重
    if len(boxes) == 0:
        return img
    boxes = np.array(boxes)
    scores = np.array(scores)
    keep_idx = nms(boxes, scores, IOU_THRESH)
    
    # 绘制结果
    for i in keep_idx:
        x1, y1, x2, y2 = boxes[i]
        score = scores[i]
        cls_id = classes[i]
        cls_name = CLASS_NAMES[cls_id]
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{cls_name}: {score:.2f}"
        cv2.putText(img, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
    return img

# 单图测试
img = cv2.imread("dataset/images/val/IMG_6099.JPG")
result_img = detect(img)
cv2.imwrite("result.jpg", result_img)

实时摄像头版本 #

cap = cv2.VideoCapture(0)
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    result_frame = detect(frame)
    cv2.imshow("YOLOv5 TFLite Detect", result_frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

如果你是自训练模型，把代码里的CLASS_NAMES替换成自己训练时用的类别列表即可。你之前写的坐标映射逻辑存在宽高颠倒的问题，参考代码已经修正了坐标转换和边界裁剪逻辑，不会出现框位置错位的问题。

内容的提问来源于stack exchange，提问作者narsan