网络上分享到的验证码数据,只有1040张图片。
制作dataloaderimport os
import cv2 as cv
import torch
from torch.utils.data import Dataset
class CodeDataset(Dataset):
def __init__(self, path='', transform=None):#, transform=None
self.path = path
self.transform = transform
self.imgs = []
for root, dirs, files in os.walk(path):
for file in files:
self.imgs.append(file)
self.n_samples = self.imgs.__len__()
self.labels = [i.split('.')[0] for i in self.imgs]
unique_set=set(char for label in self.labels for char in label)
self.map=sorted(list(unique_set))
self.characters = ['-']+sorted(list(unique_set))
self.myclass_len=len(self.characters)
self.char_to_num = dict((c, i) for i, c in enumerate(self.characters))
self.num_to_char = dict((i, c) for i, c in enumerate(self.characters))
def __getitem__(self, index):
img_path = os.path.join(self.path, self.imgs[index])
to_img = cv.imread(img_path)
if self.transform:
to_img = self.transform(to_img)
label = [self.char_to_num[i] for i in self.labels[index]]
return to_img, torch.tensor(label,dtype=torch.int32)
def __len__(self):
return self.n_samples
if __name__ == '__main__':
dataset = CodeDataset(path=r'C:UsersAdministratorDesktopWangTingfive_numbers_verification_code')
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0)
- 继承torch.utils.data中的Dataset。在__init__函数中读取图片地址和标签数据。标签编码成为数值。标签编码为数值一定要排序,不然每次重新读取标签数据就会导致字符和对应标签位置不对应,预测结果混乱。
- 编码要把空白字符带上,且一定要放在第一位,因为之后pytorch自带的ctcloss函数默认空白字符在第一位blank=0,而且beam search decoder的时候,使用的fast_ctc_decode也是默认空白字符是第一个位置上。我把空白字符用'-'表示。
- 可以使用opencv对图像做处理,不过pytorch自带的transfrom使用起来更加方便。
- transforms.ToPILImage(),不转换为PIL会报错
def __init__(self, path='', transform=None):
pass
def __getitem__(self, index):
if self.transform:
to_img = self.transform(img)
my_transforms = transforms.Compose([
transforms.ToPILImage(), # 不转换为PIL会报错
transforms.Resize([50, 50]),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5,), std=(0.5, 0.5, 0.5,))
])
dataset = CodeDataset(path=r'C:UsersAdministratorDesktopWangTingfive_numbers_verification_code',
transform=my_transforms) # , transform=transforms.ToTensor()
- 标签转化为整型的tensor就好了torch.tensor(label,dtype=torch.int32)
class DigitsModel(nn.Module):
def __init__(self,char_kinds):
super(DigitsModel, self).__init__()
self.embed = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, stride=(2, 1), padding=(1, 1), bias=False),
nn.ELU(),
nn.Conv2d(16, 32, kernel_size=3, stride=(1, 1), padding=(1, 1), bias=False),
nn.ELU(),
nn.Conv2d(32, 64, kernel_size=3, stride=(2, 2), padding=(1, 1), bias=False),
nn.ELU(),
nn.Conv2d(64, 32, kernel_size=3, stride=(1, 1), padding=(1, 1), bias=False),
nn.ELU()
)
self.rnns = nn.ModuleList()
dim = 416
hidden_size = 256
self.rnns.append(nn.LSTM(input_size=dim, hidden_size=hidden_size))
self.output_layer = nn.Linear(in_features=hidden_size, out_features=char_kinds) # len(DIGITS_MAP) + 1
def forward(self, features):
# features: n, 1, h, w
embedding = self.embed(features)
n, c, h, w = embedding.size()
embedding = embedding.view(n, c * h, w).permute(2, 0, 1)
# embed: t, n, f
h = embedding
for l in self.rnns:
h, _ = l(h)
logits = self.output_layer(h)
# lengths = torch.zeros((n,)).fill_(w)
return logits
- 使用pycharm调试model的参数,比如代码中的dim和hidden_size,虽然有公式可以让你计算,但是公式还要上网搜,pycharm打断点查看非常方便,而且可以看源码,比如nn.CTCLoss这个函数的源码,如果只用anaconda就不方便了。
- 这里需要注意,logits是要传入ctcloss函数的所以他的shape要符合要求。
参数:
blank:空白标签,默认是0 。
reduction:指明对结果使用的reduction.‘mean’代表输出的损失将处以目标长度,然后取整个批次的平均值。
zero_infinity:是否把无限损失值和相关梯度恢复为0.无限损失主要发生在输入太短而不能对其目标值得时候。
参考文章
除了“参考文章”中的内容,再补充一下,输入内容都有两种形式,带批次和不带批次的形式。一般都用带批次的这种形式。
训练模型- 注意1,2,3的顺序,我看到过其他顺序的。
- fast_ctc_decoder除了beam_serach还有viterbi_search。一些参考文章中提到的ctcloss已经跟不上python或pytorch更新速度,建议使用fast_ctc_loss
训练结果
训练20个epoch后模型拟合度非常高 ,但是因为数据量很少,泛化必然不佳。使用另一种替代方案pytesseract,发现该开源工具效果也很差,识别不到数据。这是因为验证码都几乎都添加噪声,所以对于特定验证码需要特定的预处理去除噪声。
项目代码我的gitee,安全审核两天后可见



