使用SSD MobileNetV3目标检测模型遇IndexError索引越界问题

问题重现 #

运行SSD MobileNetV3模型检测手机拍摄的勺子图片时，触发以下错误：

IndexError: list index out of range

错误出现在标签生成代码行：

label = f"{class_names[class_id - 1].upper()} {probability * 100:.2f}%"

排查发现模型返回的class_id为87，而该模型基于COCO数据集训练，仅支持1-80的类别ID，导致访问class_names列表时索引越界。

核心原因 #

模型输出的class_id可能包含背景类（ID=0）或模型无法确定的异常ID，当模型对检测目标置信度不足时，会返回超出标准COCO类别范围的ID，进而引发索引错误。

解决方案 #

1. 增加Class ID合法性校验 #

在生成标签前，先判断class_id是否在有效范围内（1到len(class_names)），超出范围则标记为未知类别：

# 提取检测对象ID并生成标签
class_id = int(detection[1])
# 新增合法性判断
if 1 <= class_id <= len(class_names):
    label = f"{class_names[class_id - 1].upper()} {probability * 100:.2f}%"
    cv2.putText(image, label, (box[0], box[1] + 15),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
else:
    # 处理未知类别
    label = f"UNKNOWN {probability * 100:.2f}%"
    cv2.putText(image, label, (box[0], box[1] + 15),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

2. 确保标签文件与模型匹配 #

确认labels.txt是COCO数据集的80类标准标签，避免标签数量不足或类别ID不对应。

3. 优化检测置信度阈值 #

适当提高置信度阈值（比如从0.5调整到0.6），过滤模型不确定的低置信度检测结果，减少异常ID出现：

if probability < 0.6:  # 提高阈值
    continue

完整修正后的代码 #

import cv2

image = cv2.imread('Images/IMG20221220090901.jpg')
image = cv2.resize(image, (640, 480))
h = image.shape[0]
w = image.shape[1]

# 模型文件路径
weights = "frozen_inference_graph.pb"
model = "ssd_mobilenet_v3_large_coco_2020_01_14.pbtxt"
net = cv2.dnn.readNetFromTensorflow(weights, model)

# 加载类别标签
class_names = []
with open("labels.txt", "r") as f:
    class_names = f.read().strip().split("\n")

# 生成图像Blob
blob = cv2.dnn.blobFromImage(
    image, 1.0/127.5, (320, 320), [127.5, 127.5, 127.5])
net.setInput(blob)
output = net.forward()

# 遍历检测结果
for detection in output[0, 0, :, :]:
    probability = detection[2]
    # 提高置信度阈值
    if probability < 0.6:
        continue

    # 计算 bounding box
    box = [int(a * b) for a, b in zip(detection[3:7], [w, h, w, h])]
    box = tuple(box)
    cv2.rectangle(image, box[:2], box[2:], (0, 255, 0), thickness=2)

    # 处理类别ID并绘制标签
    class_id = int(detection[1])
    if 1 <= class_id <= len(class_names):
        label = f"{class_names[class_id - 1].upper()} {probability * 100:.2f}%"
        cv2.putText(image, label, (box[0], box[1] + 15),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    else:
        label = f"UNKNOWN {probability * 100:.2f}%"
        cv2.putText(image, label, (box[0], box[1] + 15),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)

cv2.imshow('Image', image)
cv2.waitKey()
cv2.destroyAllWindows()  # 新增窗口销毁代码，避免资源残留

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