YOLOv8预测结果无box/masks?原因及解决办法
2025-03-14 00:32:58
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版本那样直接给出 boxes
、masks
这种方便使用的属性。
二、 解决方案:别慌! 一步步来
1. 使用 PyTorch 模型(强烈推荐)
最直接的办法,就是用 PyTorch 版本的模型(通常是 .pt
结尾的文件)进行预测。 PyTorch 模型的输出结果和文档里的是一致的,可以直接访问 boxes
、masks
等属性。
-
原理: 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 等信息, 然后再根据这些信息构建出你需要的
boxes
和masks
。 -
代码示例: (注意,这只是个大概的思路,具体实现要根据你的模型和输出格式做调整)
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预测结果里没有 box
、masks
属性的问题,大概率是因为用了 TensorFlow Lite 模型。 最好的解决办法是用 PyTorch 模型,如果实在不行再考虑手动解析 TensorFlow Lite 模型的输出.
如果能重新训练,导出模型时记得选择torchscript
格式或者不指定格式(默认会保留pytorch格式)。