源码调试-带你了解下车牌识别的深度学习模型-LPRNet

# -*- coding: utf-8 -*- # /usr/bin/env/python3  ''' 单图片验证LPRNet模型 '''  from data.load_data import CHARS, CHARS_DICT from PIL import Image, ImageDraw, ImageFont from model.LPRNet import build_lprnet from torch.autograd import Variable import torch.nn.functional as F import torch.nn as nn import numpy as np import argparse import torch import time import cv2 import os  def cv2ImgAddText(img, text, pos, textColor=(255, 0, 0), textSize=12):     if (isinstance(img, np.ndarray)):         img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))     draw = ImageDraw.Draw(img)     fontText = ImageFont.truetype("data/NotoSansCJK-Regular.ttc", textSize, encoding="utf-8")     draw.text(pos, text, textColor, font=fontText)     return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)  def predict_single_image(image_path):     # 硬编码参数     img_size = [94, 24]          # 图像尺寸     lpr_max_len = 8              # 车牌最大长度     dropout_rate = 0             # dropout率     use_cuda = False             # 是否使用CUDA     pretrained_model = './weights/Final_LPRNet_model.pth'  # 预训练模型路径          # 检查模型文件是否存在     if not os.path.exists(pretrained_model):         print(f"[Error] 预训练模型未找到: {pretrained_model}")         return          # 构建模型     lprnet = build_lprnet(lpr_max_len=lpr_max_len, phase=False,                           class_num=len(CHARS), dropout_rate=dropout_rate)     device = torch.device("cuda:0" if use_cuda else "cpu")     lprnet.to(device)          # 加载预训练模型     lprnet.load_state_dict(torch.load(pretrained_model, map_location=device))     lprnet.eval()          # 读取并预处理图像     if not os.path.exists(image_path):         print(f"[Error] 图像未找到: {image_path}")         return          # 使用OpenCV读取图像     img_orig = cv2.imread(image_path)     if img_orig is None:         print(f"[Error] 无法读取图像: {image_path}")         return          # 调整图像大小并归一化     img = cv2.resize(img_orig, (img_size[0], img_size[1]))     img = img.astype('float32')     img -= 127.5     img *= 0.0078125     img = np.transpose(img, (2, 0, 1))  # HWC -> CHW          # 转换为torch tensor并添加batch维度     img = torch.from_numpy(img).unsqueeze(0)     if use_cuda:         img = img.cuda()          # 预测     with torch.no_grad():         prebs = lprnet(img)          # 解码预测结果     prebs = prebs.cpu().numpy()[0]  # 去掉batch维度     preb_label = []     for j in range(prebs.shape[1]):         preb_label.append(np.argmax(prebs[:, j], axis=0))          # 去除重复和空白标签     no_repeat_blank_label = []     pre_c = preb_label[0]     if pre_c != len(CHARS) - 1:  # 忽略空白标签         no_repeat_blank_label.append(pre_c)     for c in preb_label:         if (pre_c == c) or (c == len(CHARS) - 1):             if c == len(CHARS) - 1:                 pre_c = c             continue         no_repeat_blank_label.append(c)         pre_c = c          # 将数字标签转换为字符     plate_number = ''.join([CHARS[i] for i in no_repeat_blank_label])          # 显示结果     result_img = cv2ImgAddText(img_orig, plate_number, (10, 30), textColor=(0, 255, 0), textSize=30)     cv2.imshow('License Plate Recognition', result_img)     print(f"识别结果: {plate_number}")     cv2.waitKey(0)     cv2.destroyAllWindows()  if __name__ == "__main__":     # 直接在这里指定要测试的图像路径     test_image_path = "3.png"  # 替换为您的测试图像路径     predict_single_image(test_image_path)