1. 导入必要的库
python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import numpy as np
import random
import string
import matplotlib.pyplot as plt
from captcha.image import ImageCaptcha
2. 定义参数和字符集
python
# 定义参数
characters = string.digits + string.ascii_uppercase
width, height, n_len, n_class = 170, 80, 4, len(characters) + 1 # 加1是为了空白类别
3. 生成验证码图像
python
# 生成随机验证码字符串和对应的图像
generator = ImageCaptcha(width=width, height=height)
def generate_random_string():
return ''.join(random.choice(characters) for _ in range(n_len))
random_str = generate_random_string()
img = generator.generate_image(random_str)
plt.imshow(img)
plt.title(random_str)
plt.show()
4. 自定义数据集
python
class CaptchaDataset(Dataset):
def __init__(self, characters, width, height, n_len, transform=None):
self.characters = characters
self.width = width
self.height = height
self.n_len = n_len
self.n_class = len(characters) + 1
self.transform = transform
self.generator = ImageCaptcha(width=self.width, height=self.height)
def __len__(self):
return 10000 # 可以根据需要调整
def __getitem__(self, idx):
random_str = generate_random_string()
img = self.generator.generate_image(random_str)
img = np.array(img)
label = [self.characters.find(x) for x in random_str]
if self.transform:
img = self.transform(img)
return img, label
# 定义转换
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
dataset = CaptchaDataset(characters, width, height, n_len, transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
5. 定义模型架构
python
class CaptchaModel(nn.Module):
def __init__(self, n_class):
super(CaptchaModel, self).__init__()
self.conv_layers = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2, 2),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2, 2)
)
self.rnn = nn.GRU(128 * 10, 128, bidirectional=True, batch_first=True)
self.dense = nn.Linear(128 * 2, n_class)
def forward(self, x):
x = self.conv_layers(x)
batch, channel, height, width = x.size()
x = x.permute(0, 3, 1, 2).contiguous().view(batch, width, -1)
x, _ = self.rnn(x)
x = self.dense(x)
return x
model = CaptchaModel(n_class)
6. 定义CTC损失函数和优化器
python
criterion = nn.CTCLoss(blank=len(characters))
optimizer = optim.Adam(model.parameters(), lr=0.001)
7. 训练和评估
python
def train(model, dataloader, criterion, optimizer, epochs=10):
model.train()
for epoch in range(epochs):
epoch_loss = 0更多内容联系1436423940
for imgs, labels in dataloader:
imgs = imgs.permute(0, 3, 1, 2)
labels = torch.tensor(labels, dtype=torch.long)
input_lengths = torch.full((imgs.size(0),), imgs.size(3), dtype=torch.long)
label_lengths = torch.full((imgs.size(0),), n_len, dtype=torch.long)
optimizer.zero_grad()
outputs = model(imgs)
outputs = outputs.log_softmax(2)
loss = criterion(outputs, labels, input_lengths, label_lengths)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
print(f'Epoch {epoch+1}/{epochs}, Loss: {epoch_loss/len(dataloader)}')
train(model, dataloader, criterion, optimizer)
8. 测试模型
python
def decode_predictions(preds, characters):
preds = preds.permute(1, 0, 2).argmax(2)
preds = preds.cpu().detach().numpy()
decoded = []
for pred in preds:
decoded.append(''.join([characters[i] for i in pred if i < len(characters)]))
return decoded
model.eval()
with torch.no_grad():
for imgs, labels in dataloader:
imgs = imgs.permute(0, 3, 1, 2)
outputs = model(imgs)
preds = decode_predictions(outputs, characters)
break
print(preds)