基于AI的車牌檢測和識別算法
問題描述、應用場景與難點
問題描述
車牌檢測和識別是計算機視覺領域的一個特定任務,主要包含兩個核心步驟:
- 車牌檢測:從圖像中準確定位車牌的位置和區域
- 車牌識別:對檢測到的車牌區域進行字符識別,轉換為文本信息
整個流程需要處理圖像輸入,輸出結構化的車牌文本信息,實現從圖像到文字的轉換。
應用場景
- 智能交通監控系統(違章識別、交通流量統計)
- 停車場自動化管理(入場出庫自動登記、計費)
- 高速公路收費站ETC輔助系統
- 車輛防盜與追蹤系統
- 城市道路規劃與交通流分析
- 小區、園區等封閉區域的車輛管理
問題難點
- 車牌在圖像中占比小,特征不明顯
- 光照條件復雜(強光、逆光、夜間等)
- 車牌存在污損、遮擋、模糊等情況
- 車輛運動造成的圖像模糊
- 不同地區車牌樣式、字符集差異大
- 復雜背景干擾(相似顏色、紋理干擾)
- 車牌可能存在傾斜、變形等情況
PyTorch實現車牌檢測和識別算法
下面實現一個兩階段的車牌檢測與識別系統:首先使用改進的YOLOv5進行車牌檢測,然后使用CNN-LSTM網絡進行字符識別。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
import os
from PIL import Image
import torchvision.transforms as transforms
import matplotlib.pyplot as plt# 1. 車牌檢測模型 - 基于YOLOv5的簡化版本
class YOLOv5LicensePlateDetector(nn.Module):def __init__(self, num_classes=1):super().__init__()# 簡化的YOLOv5骨干網絡self.backbone = nn.Sequential(# 輸入: 3x416x416nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(32),nn.LeakyReLU(0.1),nn.MaxPool2d(2), # 32x208x208nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(64),nn.LeakyReLU(0.1),nn.MaxPool2d(2), # 64x104x104nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(128),nn.LeakyReLU(0.1),nn.Conv2d(128, 64, kernel_size=1, stride=1),nn.BatchNorm2d(64),nn.LeakyReLU(0.1),nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(128),nn.LeakyReLU(0.1),nn.MaxPool2d(2), # 128x52x52nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(256),nn.LeakyReLU(0.1),nn.Conv2d(256, 128, kernel_size=1, stride=1),nn.BatchNorm2d(128),nn.LeakyReLU(0.1),nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(256),nn.LeakyReLU(0.1),nn.MaxPool2d(2), # 256x26x26nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(512),nn.LeakyReLU(0.1),nn.Conv2d(512, 256, kernel_size=1, stride=1),nn.BatchNorm2d(256),nn.LeakyReLU(0.1),nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(512),nn.LeakyReLU(0.1),nn.Conv2d(512, 256, kernel_size=1, stride=1),nn.BatchNorm2d(256),nn.LeakyReLU(0.1),nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(512),nn.LeakyReLU(0.1),)# 檢測頭 - 輸出邊界框和類別# 每個檢測框包含: x, y, w, h, confidence, class_probself.head = nn.Conv2d(512, (5 + num_classes) * 3, kernel_size=1, stride=1)self.num_classes = num_classesdef forward(self, x):x = self.backbone(x)x = self.head(x) # (batch_size, (5 + num_classes)*3, 26, 26)# 調整輸出格式 [batch_size, num_anchors, height, width, 5 + num_classes]batch_size = x.shape[0]num_anchors = 3grid_size = x.shape[2]x = x.view(batch_size, num_anchors, 5 + self.num_classes, grid_size, grid_size)x = x.permute(0, 1, 3, 4, 2).contiguous() # [batch, anchors, grid_h, grid_w, 5 + classes]return x# 2. 車牌識別模型 - CNN + LSTM架構
class LicensePlateRecognizer(nn.Module):def __init__(self, num_chars, img_height=32, img_width=128):super().__init__()self.num_chars = num_chars# CNN特征提取部分self.cnn = nn.Sequential(nn.Conv2d(3, 32, kernel_size=3, padding=1),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2, 2), # 32x16x64nn.Conv2d(32, 64, kernel_size=3, padding=1),nn.BatchNorm2d(64),nn.ReLU(),nn.MaxPool2d(2, 2), # 64x8x32nn.Conv2d(64, 128, kernel_size=3, padding=1),nn.BatchNorm2d(128),nn.ReLU(),nn.MaxPool2d(2, 2), # 128x4x16nn.Conv2d(128, 256, kernel_size=3, padding=1),nn.BatchNorm2d(256),nn.ReLU(),nn.MaxPool2d(2, 2) # 256x2x8)# 計算CNN輸出特征的尺寸self.feature_height = img_height // (2**4) # 32 / 16 = 2self.feature_width = img_width // (2**4) # 128 / 16 = 8# LSTM序列識別部分self.lstm = nn.LSTM(input_size=256 * self.feature_height, # 256 * 2 = 512hidden_size=128,num_layers=2,bidirectional=True,batch_first=True)# 輸出層self.fc = nn.Linear(256, num_chars + 1) # +1 是為了CTC的空白字符def forward(self, x):# x: [batch_size, 3, height, width]batch_size = x.size(0)# CNN特征提取x = self.cnn(x) # [batch_size, 256, 2, 8]# 調整形狀以適應LSTM輸入x = x.permute(0, 3, 1, 2) # [batch_size, width, 256, 2]x = x.view(batch_size, -1, 256 * self.feature_height) # [batch_size, 8, 512]# LSTM處理x, _ = self.lstm(x) # [batch_size, 8, 256] (雙向所以是128*2)# 輸出層x = self.fc(x) # [batch_size, 8, num_chars + 1]# 轉置為CTC Loss所需的格式 [seq_len, batch_size, num_classes]x = x.permute(1, 0, 2) # [8, batch_size, num_chars + 1]return x# 3. 數據集類
class LicensePlateDataset(Dataset):def __init__(self, image_dir, label_file, transform=None):self.image_dir = image_dirself.transform = transformself.samples = []# 讀取標簽文件# 標簽文件格式: 圖像名 x1 y1 x2 y2 車牌字符with open(label_file, 'r', encoding='utf-8') as f:for line in f:parts = line.strip().split()if len(parts) < 6:continueimg_name = parts[0]bbox = tuple(map(int, parts[1:5]))plate_chars = parts[5]self.samples.append((img_name, bbox, plate_chars))# 字符集定義 (以中國車牌為例)self.char_set = "0123456789ABCDEFGHJKLMNPQRSTUVWXYZ京津冀晉蒙遼吉黑滬蘇浙皖閩贛魯豫鄂湘粵桂瓊渝川貴云藏陜甘青寧新"self.char_to_idx = {char: i+1 for i, char in enumerate(self.char_set)} # 0留給空白字符self.idx_to_char = {i+1: char for i, char in enumerate(self.char_set)}self.idx_to_char[0] = "" # 空白字符def __len__(self):return len(self.samples)def __getitem__(self, idx):img_name, bbox, plate_chars = self.samples[idx]img_path = os.path.join(self.image_dir, img_name)# 讀取圖像img = cv2.imread(img_path)img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)# 提取車牌區域x1, y1, x2, y2 = bboxplate_img = img[y1:y2, x1:x2]# 調整大小plate_img = cv2.resize(plate_img, (128, 32))# 轉換字符為索引label = [self.char_to_idx[char] for char in plate_chars if char in self.char_to_idx]# 應用變換if self.transform:plate_img = self.transform(plate_img)# 返回原圖、車牌圖像、邊界框和標簽return {'original_image': transforms.ToTensor()(img),'plate_image': plate_img,'bbox': torch.tensor(bbox, dtype=torch.float32),'label': torch.tensor(label, dtype=torch.long),'label_length': torch.tensor(len(label), dtype=torch.long)}# 4. 訓練函數
def train_model(detector, recognizer, train_loader, val_loader, epochs=10):device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')detector.to(device)recognizer.to(device)# 定義損失函數和優化器det_criterion = nn.MSELoss() # 簡化版,實際YOLO使用更復雜的損失rec_criterion = nn.CTCLoss(blank=0)det_optimizer = torch.optim.Adam(detector.parameters(), lr=1e-4)rec_optimizer = torch.optim.Adam(recognizer.parameters(), lr=1e-4)# 訓練循環for epoch in range(epochs):detector.train()recognizer.train()train_loss = 0.0for batch in train_loader:original_images = batch['original_image'].to(device)plate_images = batch['plate_image'].to(device)bboxes = batch['bbox'].to(device)labels = batch['label']label_lengths = batch['label_length']# 清零梯度det_optimizer.zero_grad()rec_optimizer.zero_grad()# 檢測模型前向傳播det_outputs = detector(original_images)# 簡化的檢測損失計算# 實際中需要根據YOLO的輸出格式計算損失det_loss = det_criterion(det_outputs.mean(dim=[1,2,3]), bboxes)# 識別模型前向傳播rec_outputs = recognizer(plate_images)batch_size = rec_outputs.size(1)input_lengths = torch.full((batch_size,), rec_outputs.size(0), dtype=torch.long)# 計算CTC損失# 需要將標簽展平flat_labels = torch.cat([label for label in labels])rec_loss = rec_criterion(rec_outputs.log_softmax(2), flat_labels, input_lengths, label_lengths)# 總損失total_loss = det_loss + rec_losstotal_loss.backward()# 更新參數det_optimizer.step()rec_optimizer.step()train_loss += total_loss.item()# 計算平均訓練損失avg_train_loss = train_loss / len(train_loader)# 驗證detector.eval()recognizer.eval()val_loss = 0.0with torch.no_grad():for batch in val_loader:original_images = batch['original_image'].to(device)plate_images = batch['plate_image'].to(device)bboxes = batch['bbox'].to(device)labels = batch['label']label_lengths = batch['label_length']# 檢測模型det_outputs = detector(original_images)det_loss = det_criterion(det_outputs.mean(dim=[1,2,3]), bboxes)# 識別模型rec_outputs = recognizer(plate_images)batch_size = rec_outputs.size(1)input_lengths = torch.full((batch_size,), rec_outputs.size(0), dtype=torch.long)flat_labels = torch.cat([label for label in labels])rec_loss = rec_criterion(rec_outputs.log_softmax(2), flat_labels, input_lengths, label_lengths)val_loss += (det_loss + rec_loss).item()avg_val_loss = val_loss / len(val_loader)print(f'Epoch {epoch+1}/{epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}')return detector, recognizer# 5. 推理函數
def predict_license_plate(detector, recognizer, image_path, dataset, threshold=0.5):device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')detector.to(device)recognizer.to(device)detector.eval()recognizer.eval()# 圖像預處理transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])img = cv2.imread(image_path)img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)img_tensor = transform(img_rgb).unsqueeze(0).to(device)# 檢測車牌with torch.no_grad():det_outputs = detector(img_tensor)# 簡化的后處理,實際需要更復雜的解碼# 這里假設我們直接得到了邊界框# 在實際應用中,需要從YOLO輸出中解碼出邊界框和置信度bbox = [50, 50, 200, 100] # 示例值,實際應從det_outputs計算# 提取車牌區域x1, y1, x2, y2 = map(int, bbox)plate_img = img_rgb[y1:y2, x1:x2]plate_img = cv2.resize(plate_img, (128, 32))plate_tensor = transform(plate_img).unsqueeze(0).to(device)# 識別車牌字符with torch.no_grad():rec_outputs = recognizer(plate_tensor)# 應用softmaxprobs = F.softmax(rec_outputs, dim=2)# 取概率最大的字符索引_, preds = torch.max(probs, 2)preds = preds.squeeze().cpu().numpy()# 解碼預測結果(去除空白字符和重復字符)result = []prev = -1for p in preds:if p != prev and p != 0:result.append(dataset.idx_to_char[p])prev = pplate_text = ''.join(result)# 可視化結果plt.figure(figsize=(10, 5))plt.subplot(121)plt.imshow(img_rgb)plt.plot([x1, x2, x2, x1, x1], [y1, y1, y2, y2, y1], 'r-')plt.title('Detected License Plate')plt.subplot(122)plt.imshow(plate_img)plt.title(f'Recognized: {plate_text}')plt.show()return bbox, plate_text# 主函數
def main():# 數據路徑(請替換為實際路徑)train_image_dir = 'train_images/'val_image_dir = 'val_images/'train_label_file = 'train_labels.txt'val_label_file = 'val_labels.txt'# 數據變換transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),transforms.RandomAffine(degrees=5, translate=(0.1, 0.1)), # 數據增強transforms.ColorJitter(brightness=0.2, contrast=0.2)])# 創建數據集和數據加載器train_dataset = LicensePlateDataset(train_image_dir, train_label_file, transform)val_dataset = LicensePlateDataset(val_image_dir, val_label_file, transform)train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4)val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False, num_workers=4)# 初始化模型detector = YOLOv5LicensePlateDetector()recognizer = LicensePlateRecognizer(num_chars=len(train_dataset.char_set))# 訓練模型trained_detector, trained_recognizer = train_model(detector, recognizer, train_loader, val_loader, epochs=10)# 保存模型torch.save(trained_detector.state_dict(), 'license_plate_detector.pth')torch.save(trained_recognizer.state_dict(), 'license_plate_recognizer.pth')# 測試推理test_image_path = 'test_image.jpg' # 替換為測試圖像路徑bbox, plate_text = predict_license_plate(trained_detector, trained_recognizer, test_image_path, train_dataset)print(f'Detected License Plate: {plate_text} at position {bbox}')if __name__ == "__main__":main()
代碼說明
上述實現包含了一個完整的車牌檢測與識別系統,主要分為以下幾個部分:
-
車牌檢測模型:基于簡化的YOLOv5架構,用于從原始圖像中定位車牌位置。模型輸出車牌的邊界框坐標和置信度。
-
車牌識別模型:采用CNN+LSTM架構,CNN用于提取車牌圖像的特征,LSTM用于處理序列信息,實現對車牌字符的識別。使用CTC損失函數處理不定長字符序列的識別問題。
-
數據集類:自定義數據集類用于加載和預處理數據,支持讀取圖像、解析標簽、提取車牌區域和字符編碼等功能。
-
訓練與推理函數:實現了模型的訓練流程和推理功能,支持模型保存和結果可視化。
數據集需求與準備方法
數據集需求
一個高質量的車牌檢測與識別數據集應具備以下特點:
- 數據規模:至少包含10,000張以上的車輛圖像,涵蓋不同場景和條件
- 標注信息:
- 車牌位置的邊界框坐標
- 車牌上的字符內容(精確到每個字符)
- 多樣性:
- 不同類型的車輛(轎車、貨車、摩托車等)
- 不同光照條件(白天、夜晚、陰天、逆光等)
- 不同角度和姿態(正面、側面、傾斜等)
- 不同天氣條件(晴天、雨天、雪天等)
- 不同背景環境(城市道路、高速公路、停車場等)
- 不同的車牌狀態(干凈、污損、遮擋等)
常用公開數據集
- CCPD (Chinese City Parking Dataset):包含大量中國城市停車場的車牌圖像,標注詳細
- ApolloScape:包含各種交通場景的圖像,其中有車牌標注
- CALTECH LPR Dataset:包含美國車牌的數據集
- SSIG-Segmented License Plate Dataset:包含多種國家車牌的數據集
數據集準備步驟
-
數據收集:
- 收集公開數據集
- 自行拍攝補充特定場景數據
- 注意遵守數據隱私法規
-
數據清洗:
- 去除模糊、過暗或過亮的低質量圖像
- 檢查并修正錯誤標注
- 去除重復樣本
-
數據標注:
- 使用標注工具(如LabelImg、VGG Image Annotator等)標注車牌位置
- 標注車牌上的字符內容
- 建立統一的標注格式
-
數據增強:
- 幾何變換:旋轉、縮放、裁剪、平移、翻轉
- 顏色變換:亮度、對比度、飽和度調整
- 添加噪聲、模糊處理模擬真實場景
- 車牌遮擋模擬
-
數據劃分:
- 按7:2:1的比例劃分為訓練集、驗證集和測試集
- 確保各集合的數據分布一致
相關研究最新進展
近年來,車牌檢測與識別領域的研究取得了顯著進展:
-
端到端方法:傳統的兩階段方法(先檢測后識別)逐漸被端到端模型取代,如YOLO-LPR、E2E-LPR等,直接從圖像輸出車牌字符,簡化了流程并提高了精度。
-
Transformer架構應用:基于Transformer的模型(如DETR衍生模型)在車牌檢測任務中表現出色,能夠更好地處理復雜背景和小目標檢測問題。
-
小樣本學習:針對特定區域或特殊車牌類型的數據稀缺問題,小樣本學習方法被引入,通過元學習等技術提高模型的泛化能力。
-
多模態融合:結合可見光圖像和紅外圖像的多模態方法,提高了夜間和惡劣天氣條件下的識別性能。
-
實時性優化:通過模型壓縮、知識蒸餾等技術,使車牌識別模型能夠在嵌入式設備和邊緣計算平臺上實時運行,滿足實際應用需求。
-
魯棒性提升:針對車牌污損、遮擋等問題,研究人員提出了基于生成對抗網絡(GAN)的數據增強方法,以及注意力機制來聚焦關鍵字符區域。
-
跨域適應性:研究如何使模型在不同國家/地區的車牌樣式之間進行遷移,減少對特定數據集的依賴。
這些進展使得車牌檢測與識別系統在實際應用中的準確率和魯棒性不斷提高,為智能交通系統的發展提供了有力支持。