如何用PyTorch加载自定义YOLOv7权重实现验证码检测并输出标准结果

我有一个自定义训练的YOLOv7权重文件，想用PyTorch对单张验证码图片做预测。目前已经初始化了YOLOv7模型，但运行预测后返回的是难以解读的张量元组，我需要得到和运行命令python detect.py --source ../captcha.png --weights captcha_model.pt --save-txt --no-trace --exist-ok --project .. --name output --nosave一致的结果——要么是控制台输出的分类标签，要么是指定目录下captcha.txt格式的检测数据。

当前代码 #

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

classes = ['M','Y','8','9','F','B','V','I','Q','H','4','P','T',
'C','W','A','K','G','N','L','5','6','2','0','Z','7','1','J','D','E',
'O','X','3','R']

def pre_image(image_path,model):
  img = Image.open(image_path)
  mean = [0.485, 0.456, 0.406] 
  std = [0.229, 0.224, 0.225]
  transform_norm = transforms.Compose([transforms.ToTensor(), 
  transforms.Resize((224,224)),transforms.Normalize(mean, std)])
  img_normalized = transform_norm(img).float()
  img_normalized = img_normalized.unsqueeze_(0)
  img_normalized = img_normalized.to("cpu")
  with torch.no_grad():
    model.eval()  
    output = model(img_normalized)
    
    
  return output


model = torch.hub.load("WongKinYiu/yolov7","custom","captcha_model.pt",trust_repo=True)
output = pre_image("captcha.png", model)
print(output)

当前输出 #

(tensor([[[-1.26430e+00,  2.61231e+00,  3.59347e+01,  ...,  1.01391e-02,  2.46865e-02,  3.76955e-02],
         [ 7.77992e+00,  3.27832e+00,  3.90596e+01,  ...,  7.17988e-03,  1.89734e-02,  2.55351e-02],
         [ 1.55948e+01,  6.01060e+00,  4.03760e+01,  ...,  6.80088e-03,  1.94541e-02,  2.04232e-02],
         ...,
         [ 1.48366e+02,  2.00934e+02,  1.02123e+03,  ...,  2.46739e-02,  8.09314e-03,  1.13809e-02],
         [ 1.84225e+02,  2.08909e+02,  1.01224e+03,  ...,  2.44055e-02,  6.32820e-03,  7.09502e-03],
         [ 2.13092e+02,  2.05269e+02,  6.91005e+02,  ...,  2.27712e-02,  6.58516e-03,  4.77984e-03]]]), [tensor([[[[[-1.57869e+00, -3.50467e-01,  1.85951e+00,  ..., -4.58117e+00, -3.67650e+00, -3.23979e+00],
           [-1.17338e+00, -1.80912e-01,  2.22052e+00,  ..., -4.92927e+00, -3.94556e+00, -3.64183e+00],
           [-1.23863e+00,  5.13654e-01,  2.40425e+00,  ..., -4.98388e+00, -3.92005e+00, -3.87045e+00],
           ...,
           [-1.44293e+00,  1.19797e-01,  2.43755e+00,  ..., -4.76797e+00, -3.84923e+00, -3.28092e+00],
           [-1.21709e+00,  3.48075e-01,  2.28074e+00,  ..., -5.22396e+00, -3.59381e+00, -2.93395e+00],
           [-1.81546e+00, -4.95381e-02,  2.13998e+00,  ..., -4.43276e+00, -3.90898e+00, -3.05755e+00]],

          [[-1.33876e+00,  4.33999e-01,  1.77264e+00,  ..., -5.10550e+00, -4.48759e+00, -2.86169e+00],
           [-5.99240e-01,  6.20912e-02,  2.25098e+00,  ..., -5.15602e+00, -4.78393e+00, -3.22067e+00],
           [-8.69940e-01,  5.02609e-01,  2.12465e+00,  ..., -5.56788e+00, -3.98619e+00, -3.18369e+00],
           ...,
.
.
.

          [[ 5.62926e-01, -1.30771e+00,  6.28318e-02,  ..., -5.44807e+00, -2.17277e+00, -5.25101e+00],
           [ 3.79302e-01, -1.64500e+00,  3.82732e-01,  ..., -6.31032e+00, -2.52180e+00, -5.23643e+00],
           [-1.41712e-01, -1.54378e+00,  9.03876e-01,  ..., -6.42302e+00, -2.80328e+00, -6.22176e+00],
           ...,
           [-1.35747e-01, -1.81381e+00,  1.08393e+00,  ..., -6.70534e+00, -3.46506e+00, -5.70917e+00],
           [ 6.60871e-02, -1.64646e+00,  4.27860e-01,  ..., -6

期望输出 #

控制台标签格式 #

1 8, 1 I, 1 L, 1 X, Done. (824.8ms) Inference, (2.0ms) NMS

TXT文件格式 #

19 0.278333 0.449074 0.15 0.583333
7 0.636667 0.462963 0.146667 0.648148
2 0.135 0.467593 0.156667 0.583333
31 0.443333 0.444444 0.22 0.62963

解决方案 #

核心问题是直接调用model(img_normalized)返回的是YOLOv7模型的原始特征图输出，没有经过NMS（非极大值抑制）等后处理步骤，而detect.py包含完整后处理逻辑生成最终检测结果。以下是修改后的代码：

import torch
from PIL import Image

# 你的分类标签列表
classes = ['M','Y','8','9','F','B','V','I','Q','H','4','P','T',
'C','W','A','K','G','N','L','5','6','2','0','Z','7','1','J','D','E',
'O','X','3','R']

# 加载自定义模型，hub模型已内置预处理和后处理逻辑
model = torch.hub.load("WongKinYiu/yolov7", "custom", "captcha_model.pt", trust_repo=True)
# 设置置信度阈值（和detect.py默认值一致）
model.conf = 0.25

# 直接传入图片路径或PIL图片，模型自动处理
results = model("captcha.png")

# --- 生成控制台格式输出 ---
detected_labels = []
for pred in results.pred[0]:
    class_idx = int(pred[5])
    class_label = classes[class_idx]
    detected_labels.append(f"1 {class_label}")

console_output = ", ".join(detected_labels) + ", Done."
print(console_output)

# --- 生成TXT文件格式内容 ---
txt_lines = []
img_width, img_height = Image.open("captcha.png").size
for pred in results.pred[0]:
    # 转换为YOLO标准格式：class_idx x_center y_center width height（归一化值）
    x1, y1, x2, y2, conf, class_idx = pred.tolist()
    x_center = (x1 + x2) / 2 / img_width
    y_center = (y1 + y2) / 2 / img_height
    width = (x2 - x1) / img_width
    height = (y2 - y1) / img_height
    txt_line = f"{int(class_idx)} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
    txt_lines.append(txt_line)

# 打印结果
for line in txt_lines:
    print(line)
# 写入文件可取消注释：
# with open("captcha.txt", "w") as f:
#     f.write("\n".join(txt_lines))

关键说明 #

1. 移除自定义预处理：torch.hub加载的YOLOv7模型会自动处理图片缩放、归一化，无需手动用transforms操作，避免和内置逻辑冲突。
2. 利用内置后处理：调用model(image)返回的results对象包含NMS处理后的检测结果，results.pred[0]是当前图片的所有检测框，每个元素格式为[x1, y1, x2, y2, confidence, class_idx]。
3. 解析生成目标格式：
   • 控制台输出：通过class_idx映射到自定义标签列表，拼接成指定格式。
   • TXT格式：将检测框坐标转换为YOLO标准的归一化格式（坐标均相对于图片宽高）。

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