import torch
import torch.nn as nn
import torch.optim as optim
import cv2
import os
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

class LicensePlateDataset(Dataset):
def init(self, image_folder, labels_file):
self.image_folder = image_folder
self.labels = []
self.char_to_index = {}
self.index_to_char = {}

    # Загрузка данных из текстового файла
    with open(labels_file, "r", encoding="utf-8") as f:
        for line in f:
            img_path, label = line.strip().split(maxsplit=1)  # Разделение на путь к изображению и метку
            self.labels.append((img_path, label))
            for char in label:  # Создание набора символов
                if char not in self.char_to_index:
                    index = len(self.char_to_index)
                    self.char_to_index[char] = index
                    self.index_to_char[index] = char

    self.num_classes = len(self.char_to_index) + 1  # +1 для пустой метки

def __len__(self):
    return len(self.labels)

def __getitem__(self, idx):
    img_path, label = self.labels[idx]
    image = cv2.imread(os.path.join(self.image_folder, img_path), cv2.IMREAD_GRAYSCALE)
    if image is None:
        print(f"⚠️ Изображение не найдено: {img_path}")
        image = np.zeros((32, 128), dtype=np.float32)  # Заполнитель для отсутствующего изображения
    else:
        image = cv2.resize(image, (128, 32)) / 255.0  # Нормализация

    image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)  # Форма: (1, 32, 128)
    target = torch.tensor([self.char_to_index[char] for char in label], dtype=torch.long)

    return image, target

Пользовательская функция сборки

def collate_fn(batch):
images, targets = zip(*batch) # Разгруппировка батча
images = torch.stack(images) # Объединение изображений
targets = pad_sequence(targets, batch_first=True, padding_value=0) # Дополнение целей
return images, targets # Возврат изображений и дополненных целей

Загрузка набора данных

train_dataset = LicensePlateDataset(“/content/drive/MyDrive/images”, “/content/Lebal.txt”)
num_classes = train_dataset.num_classes
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)

Определение модели CRNN

class CRNN(nn.Module):
def init(self, num_classes):
super(CRNN, self).init()
self.cnn = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),

        nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2),

        nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2)
    )

    self.lstm_input_size = 256 * 4  # Регулируется на основе выхода CNN
    self.lstm = nn.LSTM(self.lstm_input_size, 256, bidirectional=True, batch_first=True)
    self.fc = nn.Linear(512, num_classes)  # Выход BiLSTM равен 512

def forward(self, x):
    x = self.cnn(x)  # Форма: (batch_size, 256, 4, 16)
    x = x.permute(0, 2, 3, 1)  # (batch_size, 4, 16, 256)
    batch_size, height, width, channels = x.shape
    x = x.reshape(batch_size, width, height * channels)  # (batch_size, 16, 1024)
    x, _ = self.lstm(x)  # (batch_size, 16, 512)

    x = self.fc(x)  # (batch_size, 16, num_classes)
    return x

Настройка обучения

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model = CRNN(num_classes).to(device)
criterion = nn.CTCLoss(blank=num_classes – 1, zero_infinity=True)
optimizer = optim.Adam(model.parameters(), lr=0.0001)
num_epochs = 50

def init_weights(m):
if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight, mode=”fan_out”, nonlinearity=’relu’)
if m.bias is not None:
nn.init.constant_(m.bias, 0)

model.apply(init_weights)

Цикл обучения

for epoch in range(num_epochs):
model.train()
total_loss = 0
for images, targets in train_loader:
images = images.to(device)
targets = targets.to(device) # Перемещение целей на то же устройство
optimizer.zero_grad()

    outputs = model(images)

    input_lengths = torch.full((outputs.size(0),), outputs.size(1), dtype=torch.long).to(device)
    target_lengths = torch.tensor([target.size(0) for target in targets], dtype=torch.long).to(device)

    # Вычисление потерь
    loss = criterion(outputs.permute(1, 0, 2), targets, input_lengths, target_lengths)

    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5)
    optimizer.step()
    total_loss += loss.item()

print(f"📌 Эпоха {epoch + 1}, Потери: {total_loss:.4f}")

Сохранить модель

torch.save(model.state_dict(), “crnn_model.pth”)