Python运行ONNX格式YOLOv7模型输入不匹配报错求助

尝试在Python中运行ONNX格式的YOLOv7模型，但每次都会触发运行时错误，提示预期输入与实际给定的输入不匹配。以下是使用的代码：

from PIL import Image
import onnx
import torch
from PIL import Image
import torchvision.transforms as transforms

img = Image.open("test (2).jpg")

resize = transforms.Resize([640, 640])
img = resize(img)

img_ycbcr = img.convert('YCbCr')
img_y, img_cb, img_cr = img_ycbcr.split()

to_tensor = transforms.ToTensor()
img_y = to_tensor(img_y)
img_y.unsqueeze_(0)

onnx_model = onnx.load("best.onnx")
onnx.checker.check_model(onnx_model)
session = rt.InferenceSession(
"best.onnx", providers=[ 'CPUExecutionProvider'])
session.get_modelmeta()
ort_inputs = {session.get_inputs()[0].name: np.array(img_y, dtype=np.float32)}
ort_outs = session.run(None, ort_inputs)
img_out_y = ort_outs[0]

img_out_y = Image.fromarray(np.uint8((img_out_y[0] * 255.0).clip(0, 255)[0]), mode='L')

# get the output image follow post-processing step from PyTorch implementation
final_img = Image.merge(
"YCbCr", [
    img_out_y,
    img_cb.resize(img_out_y.size, Image.BICUBIC),
    img_cr.resize(img_out_y.size, Image.BICUBIC),
]).convert("RGB")

# Save the image, we will compare this with the output image from mobile device
final_img.save("./runs/detect/my_exp/test (2).jpg")

解决思路 #

• 修正输入通道数
  YOLOv7的ONNX模型默认要求输入为3通道RGB图像，维度为[1, 3, 640, 640]。你的代码仅提取了YCbCr的Y单通道，输入维度变为[1, 1, 640, 640]，这是输入不匹配的核心原因，需改用RGB三通道输入。

• 补充缺失的依赖导入
  代码中使用了rt（onnxruntime）和np（numpy）但未导入，需在开头添加：

  import numpy as np
  import onnxruntime as rt
  

• 确认模型输入规格
  运行以下代码查看ONNX模型的输入要求，确保输入的形状、数据类型完全匹配：

  session = rt.InferenceSession("best.onnx", providers=['CPUExecutionProvider'])
  input_info = session.get_inputs()[0]
  print(f"输入名称: {input_info.name}")
  print(f"输入形状: {input_info.shape}")
  print(f"输入数据类型: {input_info.type}")
  

• 调整预处理流程
  改用符合YOLOv7要求的RGB预处理逻辑，示例代码：

  img = Image.open("test (2).jpg").convert('RGB')
  transform = transforms.Compose([
      transforms.Resize((640, 640)),
      transforms.ToTensor(),  # 自动完成/255归一化
  ])
  img_tensor = transform(img).unsqueeze(0)  # 增加batch维度，形状为[1,3,640,640]
  ort_inputs = {input_info.name: img_tensor.numpy().astype(np.float32)}
  

• 替换错误的后处理逻辑
  YOLOv7的ONNX输出是检测框的坐标、置信度和类别，并非图像张量，你当前的后处理逻辑属于超分辨率模型的处理方式，完全不适用YOLOv7。需替换为YOLOv7的检测后处理：解析输出张量，过滤置信度阈值，绘制检测框到原图。

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