使用Ultralytics YOLO提取车牌检测边界框遇索引错误及标注失效

问题描述 #

我在用Ultralytics YOLO做车牌检测，尝试从Results.boxes对象中提取边界框坐标时遇到问题，触发了IndexError，而且cv2无法高亮显示车牌。我的代码如下：

class ImageProcessing:
    def __init__(self, model_path: Path, input_image: Path, output_image: Path):
        if not isinstance(model_path, Path):
            raise TypeError("model_path must be a pathlib.Path instance")
        if not isinstance(input_image, Path) or not isinstance(output_image, Path):
            raise TypeError("input_image and output_image must be pathlib.Path instances")
        # Load the YOLO model from the provided path
        self.model = YOLO(str(model_path))
        self.input_image = input_image
        self.output_image = output_image

    def ascertain_license_plates_as_image(self, threshold: float = 0.5, fontscale: float = 1.3, color: tuple = (0, 255, 0), thickness: int = 3):
        image = opencv.imread(str(self.input_image))
        results = self.model(image)

        # Check if results is a list and get the first result
        if isinstance(results, list):
            results = results[0]

        # Iterate through each detected object
        for box in results.boxes:
            # Extract coordinates, confidence, and class ID
            x1, y1, x2, y2, conf, class_id = box.data[0][0], box.data[0][1], box.data[0][2], box.data[0][3], box.conf.item(), int(box.cls.item())
            if conf > threshold:
                opencv.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
                label = results.names[class_id].upper() if results.names else f'class {class_id}'
                opencv.putText(image, label, (int(x1), int(y1) - 10), opencv.FONT_HERSHEY_SIMPLEX, fontscale, color, thickness, opencv.LINE_AA)

        opencv.imwrite(str(self.output_image), image)
        return results

错误原因 #

直接通过box.data[0][0]索引边界框坐标的方式不合理，Ultralytics YOLO的Boxes对象中，单个box的data张量维度为(1,6)，错误的索引逻辑会触发IndexError；同时，box.conf和box.cls返回的是张量，直接调用.item()虽然能取单个值，但如果前置索引错误，会导致后续绘图逻辑失效。

修复方案 #

使用Ultralytics YOLO为Boxes对象提供的便捷属性xyxy、conf、cls来提取数据，避免手动索引张量：

修复后的核心代码 #

# 遍历每个检测框
for box in results.boxes:
    # 直接通过属性获取坐标、置信度、类别ID
    x1, y1, x2, y2 = box.xyxy[0].tolist()  # xyxy返回(1,4)张量，转为列表
    conf = box.conf[0].item()
    class_id = int(box.cls[0].item())
    
    if conf > threshold:
        opencv.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
        label = results.names[class_id].upper() if results.names else f'class {class_id}'
        opencv.putText(image, label, (int(x1), int(y1) - 10), opencv.FONT_HERSHEY_SIMPLEX, fontscale, color, thickness, opencv.LINE_AA)

额外优化建议 #

• 增加图像读取校验：在opencv.imread后添加判断if image is None: raise ValueError("Failed to read input image")，避免后续无效处理。
• 若无需自定义绘图，可直接使用results[0].plot()方法快速生成带标注的图像，简化代码。

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