使用YOLO-NAS模型处理视频时检测标签及边界框未显示问题

可能的问题与修复方案 #

1. 置信度阈值过滤导致有效检测被忽略

YOLO-NAS的predict方法默认会过滤低置信度结果（通常阈值为0.25），如果你的模型预测置信度低于该值，就不会返回任何边界框。需要显式调整阈值，同时验证是否有检测结果返回。

修复：

• 在predict时指定更低的置信度阈值（根据模型性能调整）
• 打印检测数量确认是否有结果

# 调整置信度阈值
result = list(best_model.predict(frame, confidence_threshold=0.1))[0]
# 打印检测数量
print(f"Frame {count}: 检测到 {len(bbox_xyxys)} 个目标")

2. 重复移动模型到设备导致预测异常

循环内每次执行best_model.to(device).predict(frame)会重复将模型迁移到GPU/CPU，可能引发预测错误。应提前完成设备迁移。

修复：
在循环外执行一次设备迁移：

best_model = best_model.to(device)

循环内直接调用预测：

result = list(best_model.predict(frame, confidence_threshold=0.1))[0]

3. 图像格式不匹配（BGR vs RGB）

OpenCV读取的帧是BGR格式，而YOLO-NAS训练时通常使用RGB格式，格式不匹配会导致预测失效。

修复：
预测前将BGR转为RGB：

frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
result = list(best_model.predict(frame_rgb, confidence_threshold=0.1))[0]

4. 视频写入帧率过高导致播放异常

你设置的视频帧率为200，远高于普通视频的24/30帧，输出视频播放时会极快，看起来像是没有检测框。需匹配原视频帧率。

修复：
获取原视频帧率并设置给VideoWriter：

fps = int(cap.get(cv2.CAP_PROP_FPS))
out = cv2.VideoWriter('l.mp4', fourcc, fps, (frame_width, frame_height))

5. 未添加标签绘制逻辑（可选）

你的代码仅绘制了边界框，若需要显示标签，需添加文本绘制代码：

class_name = class_names[cls]
label = f"{class_name}: {confidence:.2f}"
cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,144,30), 2)

修改后的完整代码示例 #

import cv2
import torch
import numpy as np

cap = cv2.VideoCapture('/Users/anika/PycharmProjects/ml_projects/test.mp4')
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
fps = int(cap.get(cv2.CAP_PROP_FPS))

device = 0 if torch.cuda.is_available() else "cpu"
best_model = best_model.to(device)

fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
out = cv2.VideoWriter('l.mp4', fourcc, fps, (frame_width, frame_height))

count = 0
while True:
    ret, frame = cap.read()
    count += 1
    if ret:
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        result = list(best_model.predict(frame_rgb, confidence_threshold=0.1))[0]
        bbox_xyxys = result.prediction.bboxes_xyxy.tolist()
        confidences = result.prediction.confidence
        labels = result.prediction.labels.tolist()
        
        print(f"Frame {count}: 检测到 {len(bbox_xyxys)} 个目标")
        
        class_names = dataset_params['classes']
        for (bbox_xyxy, confidence, cls) in zip(bbox_xyxys, confidences, labels):
            x1, y1, x2, y2 = map(int, bbox_xyxy)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255,144,30), 3)
            class_name = class_names[cls]
            label = f"{class_name}: {confidence:.2f}"
            cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,144,30), 2)
        
        out.write(frame)
        if cv2.waitKey(1) & 0xFF == ord('1'):
            break
    else:
        break

out.release()
cap.release()
cv2.destroyAllWindows()

额外排查步骤 #

• 确认模型加载正确：加载时需指定匹配的类别数量，示例：

  from super_gradients.training import models
  best_model = models.get("yolo_nas_s", num_classes=len(dataset_params['classes']), checkpoint_path="best.pth")
  

• 先用单张测试图片验证模型检测能力，排除视频处理环节的问题。

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