YOLOv5批量推理仅返回单图结果，求结构化输出解决方案

问题场景 #

使用自定义训练的YOLOv5 P6模型对4张图片批量推理时，调用results.pandas().xyxy[0]仅返回单张图片的结构化输出；调用results.pandas().xyxy[0:]虽能看到4张图片的结果，但格式混乱，包含空DataFrame，无法得到标准结构化DataFrame用于后续处理。

现有推理代码 #

######INFERENCE ON P6 MODELS*****************************************************************************
import torch
import glob
from natsort import natsorted
import cv2
import numpy as np
import matplotlib.pyplot as plt
#%matplotlib inline
import pandas as pd
import os

model = torch.hub.load('/Users/yolov5', 'custom', path='/User/yolov5/runs/train/exp11/weights/best.pt', source='local', force_reload=True)  # custom trained model
model.conf = 0.25  # NMS confidence threshold

Path = 'User/yolov5/data/images/'

imgs = [cv2.imread(file) for file in natsorted(glob.glob(Path+"/*.jpg"))]


# Inference
results = model(imgs,size=640)

# Results:

#results.save() # or .print() .show(), .save(), .crop(), .pandas(), etc.
results.pandas()
#print(results.print())
#print(results.pandas().xyxy[:])
# results.show()
#results.pandas().xyxy[0]


#print(results)
#print(results.pandas().xyxy[0])
# dfm = pd.DataFrame(results.pandas().xyxy[0])#, columns = ['Loss','Accuracy']
# # #dfm['Classes'] = classes.tolist()
# predict_labs = 'pred_yolo_individual.csv'
# with open(predict_labs, mode='w') as fd:
#      dfm.to_csv(fd)

#results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
#results.render()
results.xyxy[0]  # im predictions (tensor)
results.pandas().xyxy[0] 
results.print()
# pred = results.pandas().xyxy[0]
# for index, row in pred.iterrows():
#     print(row['class'], row['confidence'], row['name'])

当前异常输出 #

YOLOv5 🚀 v7.0-72-g064365d Python-3.10.6 torch-1.13.1 CPU

Fusing layers... 
Model summary: 212 layers, 20856975 parameters, 0 gradients, 47.9 GFLOPs
Adding AutoShape... 
[         xmin       ymin        xmax        ymax  confidence  class      name
0  539.859314  119.92907  602.884216  245.533752    0.353711      1  Stabbing, Empty DataFrame
Columns: [xmin, ymin, xmax, ymax, confidence, class, name]
Index: [], Empty DataFrame
Columns: [xmin, ymin, xmax, ymax, confidence, class, name]
Index: [],          xmin       ymin         xmax        ymax  confidence  class      name
0  709.833496  66.843300  1025.770752  800.782593    0.771696      1  Stabbing
1   84.628845   4.153772   461.863617  833.189636    0.632551      1  Stabbing]

期望输出格式 #

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

results.xyxy[0]  # im predictions (tensor)
results.pandas().xyxy[0]  # im predictions (pandas)
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

解决方案 #

批量推理时，results.pandas().xyxy是一个列表，每个元素对应一张图片的预测结果DataFrame（空DataFrame表示该图片无检测结果）。要得到统一的结构化格式，需遍历所有结果并添加图片标识，再合并为一个完整的DataFrame。

修改后处理代码 #

# 批量推理后处理
all_results = []
img_paths = natsorted(glob.glob(Path+"/*.jpg"))  # 获取图片路径用于区分不同图片

for idx, (img_path, pred_df) in enumerate(zip(img_paths, results.pandas().xyxy)):
    # 添加图片路径列，标记检测结果所属图片
    pred_df['image_path'] = img_path
    # 可选：添加图片索引列
    pred_df['image_idx'] = idx
    all_results.append(pred_df)

# 合并所有图片的结果为一个标准DataFrame
final_df = pd.concat(all_results, ignore_index=True)

# 打印查看结构化结果
print(final_df)

# 可选：保存为CSV文件用于后续处理
final_df.to_csv('yolov5_batch_predictions.csv', index=False)

关键说明 #

• 遍历results.pandas().xyxy列表，逐个处理每张图片的预测结果
• 添加image_path或image_idx列，确保能区分不同图片的检测结果
• 用pd.concat合并所有DataFrame，得到统一的结构化格式
• 若需过滤无检测结果的图片，可在遍历时添加判断：if not pred_df.empty

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