返回

YOLOv8预测结果无box/masks?原因及解决办法

Ai

YOLOv8 预测结果中缺少 'box' 和 'masks' 属性? 别急,这儿有解!

最近用YOLOv8训练了一个模型,用来识别路口的物体(比如汽车、道路之类的)。 模型跑起来没啥问题,能把感兴趣的物体都分割出来,图像结果看着挺好。

但是,我需要的是原始的几何形状(多边形)数据,好把它们保存到txt文件里。 我照着官方文档(https://docs.ultralytics.com/modes/predict/#key-features-of-predict-mode)试了试,结果返回的对象和文档里说的不一样。

实际上,返回的结果是一堆 tensorflow 数字:

[图片]

代码长这样:

import argparse
import cv2
import numpy as np
from pathlib import Path
from ultralytics.yolo.engine.model import YOLO

# Parse command line arguments
parser = argparse.ArgumentParser()
parser.add_argument('--source', type=str, required=True, help='Source image directory or file')
parser.add_argument('--output', type=str, default='output', help='Output directory')
args = parser.parse_args()

# Create output directory if it doesn't exist
Path(args.output).mkdir(parents=True, exist_ok=True)

# Model path
model_path = r'C:\\_Projects\\best_100img.pt'

# Load your model directly
model = YOLO(model_path)
model.fuse()

# Load image(s)
if Path(args.source).is_dir():
    image_paths = list(Path(args.source).rglob('*.tiff'))
else:
    image_paths = [args.source]

# Process each image
for image_path in image_paths:
    img = cv2.imread(str(image_path))
    if img is None:
        continue

    # Perform inference
    predictions = model.predict(image_path, save=True, save_txt=True)

print("Processing complete.")

问题就在这儿:返回的对象 ( predictions 变量) 没有 boxes、masks、keypoints 之类的属性。

我估计你也会遇到一样的问题,我当时主要有两个疑惑:

  • 为啥返回的结果和文档里差这么多?
  • 是不是需要啥转换的步骤?

一、 问题原因:TensorFlow Lite 模型的限制

仔细研究后,我发现问题出在模型格式上。 我的模型 (best_100img.pt) 很有可能是被转换成了 TensorFlow Lite 格式。 虽然YOLOv8能直接加载并使用这种模型, 但是TensorFlow Lite 模型在预测时返回的结果和 PyTorch 版本的模型不一样。

TensorFlow Lite 模型通常用于嵌入式设备或者移动设备,它为了效率和大小做了优化, 输出结果比较“底层”,不像PyTorch版本那样直接给出 boxesmasks 这种方便使用的属性。

二、 解决方案:别慌! 一步步来

1. 使用 PyTorch 模型(强烈推荐)

最直接的办法,就是用 PyTorch 版本的模型(通常是 .pt 结尾的文件)进行预测。 PyTorch 模型的输出结果和文档里的是一致的,可以直接访问 boxesmasks 等属性。

  • 原理: PyTorch 是一个深度学习框架, 提供了完整的模型定义、训练和推理功能。YOLOv8 的官方实现就是基于 PyTorch 的, 所以用 PyTorch 模型能保证结果的完整性和易用性。

  • 代码示例:
    如果你的原始训练模型是 PyTorch 的,那就简单了, 直接加载 .pt 文件。 如果你在训练时候就已经导出了tflite格式, 可以尝试去训练日志的文件夹找一下.

    from ultralytics import YOLO
    
    # 加载 PyTorch 模型
    model = YOLO('path/to/your/pytorch_model.pt')  # 换成你的 PyTorch 模型路径
    
    # 进行预测
    results = model.predict(source='your_image.jpg')
    
    # 获取第一个预测结果的 boxes
    if results and results[0].boxes:
      boxes = results[0].boxes
      print(boxes.xyxy)  # 边界框坐标 (x1, y1, x2, y2)
      print(boxes.conf)  # 置信度
      print(boxes.cls)   # 类别ID
    
      if results[0].masks: #获取masks (只有在使用了segmentation模型时候才有)
          masks = results[0].masks
          print(masks.xy) #多边形的点
    
    

2. 手动解析 TensorFlow Lite 模型的输出 (不推荐,除非没办法)

如果你实在没办法拿到 PyTorch 模型,只能用 TensorFlow Lite 模型,那就要费点劲了,得手动解析输出结果。

  • 原理: TensorFlow Lite 模型的输出通常是一个多维数组, 包含了所有可能的检测结果。你需要根据 YOLOv8 的输出格式规范, 从这个数组里提取出边界框、置信度、类别 ID 等信息, 然后再根据这些信息构建出你需要的 boxesmasks

  • 代码示例: (注意,这只是个大概的思路,具体实现要根据你的模型和输出格式做调整)

    import numpy as np
    from ultralytics import YOLO
    
    # 加载 TensorFlow Lite 模型 (这里假设你已经有了 .tflite 文件)
    model = YOLO('path/to/your/tflite_model.tflite')
    
    # 进行预测
    results = model.predict(source='your_image.jpg')
    
    # 解析输出 (这部分需要根据你的模型输出格式做调整)
    # 假设 output 是一个 [1, num_detections, 6] 的数组
    # 其中 num_detections 是检测到的物体数量, 6 表示 [x_center, y_center, width, height, confidence, class_id]
    output = results[0].numpy()  # 把输出转换为 NumPy 数组, 取决于result结构,可能需要微调
    
    boxes = []
    for detection in output:
        x_center, y_center, width, height, confidence, class_id = detection
        if confidence > 0.5:  # 设置一个置信度阈值
            # 把中心点坐标转换为左上角和右下角坐标
            x1 = int(x_center - width / 2)
            y1 = int(y_center - height / 2)
            x2 = int(x_center + width / 2)
            y2 = int(y_center + height / 2)
            boxes.append([x1, y1, x2, y2, confidence, int(class_id)])
    print (boxes)
    
    

    安全提示: 处理用户上传的模型文件时,一定要小心,防止恶意代码注入。
    这个过程会非常繁琐,需要你对 YOLOv8 的输出格式非常熟悉。强烈建议还是用上面的方法一,省心省力。

3. (进阶技巧)训练时指定输出格式

如果你在训练模型,可以在导出模型的时候指定输出格式。如果明确指定了格式为torchscript onnx或者pt,那么会直接在预测时返回 boxes属性等.

yolo export model=yolov8n.pt format=torchscript

三、 总结一下

遇到YOLOv8预测结果里没有 boxmasks 属性的问题,大概率是因为用了 TensorFlow Lite 模型。 最好的解决办法是用 PyTorch 模型,如果实在不行再考虑手动解析 TensorFlow Lite 模型的输出.
如果能重新训练,导出模型时记得选择torchscript格式或者不指定格式(默认会保留pytorch格式)。