如何在Python中使用训练好的YOLOv7 best.pt实现车牌检测应用？

环境准备 #

首先确保安装必要的依赖包：

pip install torch torchvision opencv-python numpy

如果倾向于使用YOLOv7官方原生代码，先克隆仓库并进入目录：

git clone https://github.com/WongKinYiu/yolov7.git
cd yolov7

如果追求更简洁的封装，可安装Ultralytics YOLO库（支持YOLOv7）：

pip install ultralytics

方法一：使用YOLOv7官方原生代码 #

这种方式完全贴合官方训练的模型结构，适合需要自定义预处理/后处理逻辑的场景：

import cv2
import numpy as np
import torch
import os
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device
from utils.datasets import letterbox

# 配置参数
weights_path = "best.pt"  # 你的训练模型路径
test_img_path = "test_car.jpg"  # 待检测的车辆图像路径
save_dir = "cropped_plates"  # 裁剪后车牌的保存目录
conf_threshold = 0.5  # 检测置信度阈值，过滤低置信度结果
iou_threshold = 0.45  # IOU阈值，去除重叠检测框

# 创建保存目录
os.makedirs(save_dir, exist_ok=True)

# 加载计算设备（自动选择GPU/CPU）
device = select_device("")
# 加载训练好的YOLOv7模型
model = attempt_load(weights_path, map_location=device)
stride = int(model.stride.max())  # 获取模型步长，用于图像预处理

# 读取原始图像
img_original = cv2.imread(test_img_path)
# 图像预处理：等比例缩放，保持宽高比
img_processed = letterbox(img_original, new_shape=640, stride=stride)[0]
# 格式转换：BGR→RGB，HWC→CHW
img_processed = img_processed[:, :, ::-1].transpose(2, 0, 1)
img_processed = np.ascontiguousarray(img_processed)

# 模型推理
img_tensor = torch.from_numpy(img_processed).to(device)
img_tensor = img_tensor.float() / 255.0  # 归一化到[0,1]
if img_tensor.ndimension() == 3:
    img_tensor = img_tensor.unsqueeze(0)  # 增加batch维度

# 获取推理结果并进行NMS（非极大值抑制）
predictions = model(img_tensor)[0]
predictions = non_max_suppression(predictions, conf_threshold, iou_threshold)

# 遍历检测结果，裁剪并保存车牌
for det in predictions:
    if len(det) == 0:
        continue
    # 将检测框坐标从预处理图像映射回原始图像尺寸
    det[:, :4] = scale_coords(img_processed.shape[1:], det[:, :4], img_original.shape).round()
    # 处理每个检测到的车牌
    for idx, (*bbox, conf, cls) in enumerate(det):
        x1, y1, x2, y2 = map(int, bbox)
        # 裁剪车牌区域
        plate_img = img_original[y1:y2, x1:x2]
        # 保存为JPEG格式
        save_path = os.path.join(save_dir, f"plate_{idx}.jpg")
        cv2.imwrite(save_path, plate_img)
        print(f"已保存车牌：{save_path}")

方法二：使用Ultralytics YOLO封装库 #

这种方式无需关注底层预处理和NMS逻辑，代码更简洁高效：

from ultralytics import YOLO
import cv2
import os

# 配置参数
weights_path = "best.pt"
test_img_path = "test_car.jpg"
save_dir = "cropped_plates"
conf_threshold = 0.5

# 创建保存目录
os.makedirs(save_dir, exist_ok=True)

# 加载训练好的YOLOv7模型
model = YOLO(weights_path)

# 执行推理
results = model(test_img_path, conf=conf_threshold)

# 处理检测结果
for result in results:
    # 遍历每个检测框
    for idx, box in enumerate(result.boxes):
        # 获取检测框坐标
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        # 读取原始图像并裁剪车牌
        img_original = cv2.imread(test_img_path)
        plate_img = img_original[y1:y2, x1:x2]
        # 保存车牌
        save_path = os.path.join(save_dir, f"plate_{idx}.jpg")
        cv2.imwrite(save_path, plate_img)
        print(f"已保存车牌：{save_path}")

关键提示 #

• 与YOLOv3的.weights文件不同，YOLOv7的.pt是PyTorch格式的模型文件，需要使用PyTorch兼容的加载方式（上述两种方法均满足）。
• 调整conf_threshold可以控制检测的严格程度：值越高，检测结果越少但精度越高；值越低，检出率越高但可能出现误检。
• 如果检测框位置不准确，可检查训练时的图像尺寸与推理时的new_shape是否匹配，或微调iou_threshold参数。

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