赞
踩
文字识别是AI的一个重要应用场景,文字识别过程一般由图像输入、预处理、文本检测、文本识别、结果输出等环节组成。
其中,文本检测、文本识别是最核心的环节。文本检测方面,在我的 OCR_detection
专栏相关文章中已介绍过了多种基于深度学习的方法(有的还没完成,待整理后都会放入该专栏),可针对各种场景实现对文字的检测,详请见专栏中的相关文章。
在以前的 OCR 任务中,识别过程分为两步:单字切割
和 分类任务
。我们一般都会将一连串文字的文本文件先利用 投影法
切割出单个字体,再送入 CNN 里进行文字分类。但是此法已经有点过时了,现在更流行的是基于深度学习的端到端的文字识别,即我们不需要显式加入文字切割这个环节,而是将文字识别转化为序列学习问题,虽然输入的图像尺度不同,文本长度不同,但是经过 DCNN 和 RNN 后,在输出阶段经过一定的 CTC 翻译转录后,就可以对整个文本图像进行识别,也就是说,文字的切割也被融入到深度学习中去了。
现今基于深度学习的端到端 OCR 技术有两大主流技术:CRNN OCR
和 attention OCR
。其实这两大方法主要区别在于最后的输出层(翻译层),即怎么将网络学习到的序列特征信息转化为最终的识别结果。这两大主流技术在其特征学习阶段都采用了 CNN+RNN
的网络结构,CRNN OCR
在对齐时采取的方式是 CTC
算法,而 attention OCR
采取的方式则是 attention
机制。本部分主要介绍应用更为广泛的 CRNN
算法。
CRNN 模型,即将 CNN 与 RNN 网络结合,共同训练。主要用于在一定程度上实现端到端(end-to-end)地对不定长的文本序列进行识别,不用先对单个文字进行切割,而是将文本识别转化为时序依赖的序列学习问题,就是基于图像的序列识别。(说一定程度是因为虽然输入图像不需要精确给出每个字符的位置信息,但实际上还是需要对原始的图像进行前期的裁切工作)
- 构建 CRNN 输入特征序列;
- 其中还涉及到了 CTC 模块,目的是对其输入输出结果
整个CRNN网络结构包含三部分,从下到上依次为:
- CNN(卷积层):使用深度 CNN,对输入图像提取特征,得到特征图;
- RNN(循环层):使用 双向RNN(BLSTM)对特征序列进行预测,对序列中的每个特征向量进行学习,并输出预测标签(真实值)分布;
- CTC loss(转录层):使用 CTC 损失,把从循环层获取的一系列标签分布转换成最终的标签序列。
这里有一个很精彩的改动,一共有四个最大池化层,但是最后两个池化层的窗口尺寸由 2x2
改为 1x2
,也就是图片的高度减半了四次(除以
2
4
2^4
24),而宽度则只减半了两次(除以
2
2
2^2
22),这是因为文本图像多数都是高较小而宽较长,所以其 feature map 也是这种高小宽长的矩形形状,如果使用 1×2
的池化窗口可以尽量保证不丢失在宽度方向的信息,更适合英文字母识别(比如区分 i
和 l
)。
CRNN 还引入了 Batch Normalization 模块,加速模型收敛,缩短训练过程。
例如:
- 输入图像为灰度图像(单通道);
- 高度为32,这是固定的,图片通过 CNN 后,高度就变为 1,这点很重要;
- 宽度为160,宽度也可以为其他的值,但需要统一,所以输入 CNN 的数据尺寸为 (channel, height, width)=(1, 32, 160)。
- CNN 的输出尺寸为 (512, 1, 40)。即 CNN 最后得到 512 个特征图,每个特征图的高度为 1,宽度为 40。
注意:最后的卷积层是一个
2*2, s=1, p=0
的卷积,此时也是相当于将 feature map 放缩为原来的 1/2,所以整个 CNN 层将图像的 h 放缩为原来的 1 2 4 × 2 = 1 32 \frac{1}{2^4 \times 2}=\frac{1}{32} 24×21=321,所以最后 CNN 输出的 feature map 的高度为1。assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
- 1
在程序中,图像的 h 必须为 16 的整数倍。
assert h == 1, "the height of conv must be 1"
- 1
前向传播时,CNN 得到的 feature map 的 h 必须为 1。
最后 CNN 得到的 feature map 尺度为 512x1x16
不能直接把 CNN 得到的特征图送入 RNN 进行训练的,需要进行一些调整,根据特征图提取 RNN 需要的特征向量序列。
现在需要从 CNN 模型产生的特征图中提取特征向量序列,每一个特征向量(如上图中的一个红色框)在特征图上 按列 从左到右生成,每一列包含 512 维特征,这意味着第 i
个特征向量是所有的特征图第 i
列像素的连接,这些特征向量就构成一个序列。
由于卷积层,最大池化层和激活函数在局部区域上执行,因此它们是平移不变的。因此,特征图的每列(即一个特征向量)对应于原始图像的一个矩形区域(称为感受野),并且这些矩形区域与特征图上从左到右的相应列具有相同的顺序。特征序列中的每个向量关联一个感受野。如下图所示:
这些特征向量序列就作为循环层的输入,每个特征向量作为 RNN 在一个时间步(time step)的输入。
因为 RNN 有梯度消失的问题,不能获取更多上下文信息,所以 CRNN 中使用的是 LSTM
,LSTM 的特殊设计允许它捕获长距离依赖。
LSTM 是单向的,它只使用过去的信息。然而,在基于图像的序列中,两个方向的上下文是相互有用且互补的。将两个 LSTM,一个向前和一个向后组合到一个双向 LSTM 中。此外,可以堆叠多层双向 LSTM,深层结构允许比浅层抽象更高层次的抽象。
这里采用的是两层各 256 单元的双向 LSTM 网络:
通过上面一步,我们得到了 40 个特征向量,每个特征向量长度为 512,在 LSTM 中一个时间步就传入一个特征向量进行分类,这里一共有 40 个时间步。
我们知道一个特征向量就相当于原图中的一个小矩形区域,RNN 的目标就是预测这个矩形区域为哪个字符,即根据输入的特征向量,进行预测,得到所有字符的 softmax 概率分布,这是一个长度为字符类别数的向量,作为 CTC 层的输入。
因为每个时间步都会有一个输入特征向量 x t x_t xt,输出一个所有字符的概率分布 y t y_t yt,所以输出为 40 个长度为字符类别数的向量构成的后验概率矩阵。如下图所示:
然后将这个后验概率矩阵传入转录层。
该部分的源码如下:
self.rnn = nn.Sequential(
BidirectionalLSTM(512, nh, nh),
BidirectionalLSTM(nh, nh, nclass))
然后参数设置如下:
nh = 256
nclass = len(opt.alphabet) + 1
nc = 1
其中 BLSTM 的实现如下:
class BidirectionalLSTM(nn.Module):
def __init__(self, nIn, nHidden, nOut):
super(BidirectionalLSTM, self).__init__()
self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True)
self.embedding = nn.Linear(nHidden * 2, nOut)
def forward(self, input):
recurrent, _ = self.rnn(input)
T, b, h = recurrent.size()
t_rec = recurrent.view(T * b, h)
output = self.embedding(t_rec) # [T * b, nOut]
output = output.view(T, b, -1)
return output
所以第一次 LSTM 得到的 output=[40*256,256]
,然后 view 成 output=[40,256,256]
第二次 LSTM 得到的结果是 output=[40*256,nclass]
,然后 view 成 output=[40,256,nclass]
这算是 CRNN 最难的地方,这一层为转录层,转录是将 RNN 对每个特征向量所做的预测转换成标签序列的过程。数学上,转录是根据每帧预测找到具有最高概率组合的标签序列。
端到端 OCR 识别的难点在于怎么处理不定长序列对齐的问题!OCR 可建模为时序依赖的文本图像问题,然后使用 CTC(Connectionist Temporal Classification, CTC)的损失函数来对 CNN 和 RNN 进行端到端的联合训练。
我们现在要将 RNN 输出的序列翻译成最终的识别结果,RNN 进行时序分类时,不可避免地会出现很多冗余信息,比如一个字母被连续识别两次,这就需要一套去冗余机制。
比如我们要识别上面这个文本,其中 RNN 中有 5 个时间步,理想情况下 t0, t1, t2
时刻都应映射为 “a”
,t3, t4
时刻都应映射为 “b”
,然后将这些字符序列连接起来得到 “aaabb”
,我们再将连续重复的字符合并成一个,那么最终结果为 “ab”
。
这似乎是个比较好的方法,但是存在一个问题,如果是 book,hello 之类的词,合并连续字符后就会得到 bok 和 helo,这显然不行,所以 CTC 有一个 blank
机制来解决这个问题。
我们以 “-”
符号代表 blank
,RNN 输出序列时,在文本标签中的重复的字符之间插入一个 “-”
,比如输出序列为 “bbooo-ookk”
,则最后将被映射为 “book”
,即有 blank
字符隔开的话,连续相同字符就不进行合并。
即对字符序列先删除连续重复字符,然后从路径中删除所有 “-”
字符,这个称为解码过程,而编码则是由神经网络来实现。引入 blank
机制,我们就可以很好地解决重复字符的问题。
相同的文本标签可以有多个不同的字符对齐组合,例如,“aa-b” 和 “aabb” 以及 “-abb” 都代表相同的文本 (“ab”),但是与图像的对齐方式不同。更总结地说,一个文本标签存在一条或多条的路径。
在训练阶段,我们需要根据这些概率分布向量和相应的文本标签得到损失函数,从而训练神经网路模型,下面来看看如何得到损失函数的。
如上图,对于最简单的时序为 2 的字符识别,有两个时间步长 (t0,t1)
和三个可能的字符为 “a”
,“b”
和 “-”
,我们得到两个概率分布向量,如果采取最大概率路径解码的方法,则 “--”
的概率最大,即真实字符为空的概率为 0.6*0.6=0.36。
但是为字符 “a” 的情况有多种对齐组合,“aa”, “a-“ 和 “-a” 都是代表 “a”,所以,输出 “a” 的概率应该为三种之和:
所以 “a” 的概率比空 “-” 的概率高!如果标签文本为 “a”,则通过计算图像中为 “a” 的所有可能的对齐组合(或者路径)的分数之和来计算损失函数。
所以对于 RNN 给定输入概率分布矩阵为 y={y1,y2,…,yT},T是序列长度,最后映射为标签文本l的总概率为:
其中 B(π) 代表从序列到序列的映射函数 B 变换后是文本 l 的所有路径集合,而 π 则是其中的一条路径。每条路径的概率为各个时间步中对应字符的分数的乘积。
我们就是需要训练网络使得这个概率值最大化,类似于普通的分类,CTC 的损失函数定义为概率的负最大似然函数,为了计算方便,对似然函数取对数。
通过对损失函数的计算,就可以对之前的神经网络进行反向传播,神经网络的参数根据所使用的优化器进行更新,从而找到最可能的像素区域对应的字符。
这种通过映射变换和所有可能路径概率之和的方式使得 CTC 不需要对原始的输入字符序列进行准确的切分。
在测试阶段与训练阶段有所不同,我们用训练好的神经网络来识别新的文本图像。这时候我们事先不知道任何文本,如果我们像上面一样将每种可能文本的所有路径计算出来,对于很长的时间步和很长的字符序列来说,这个计算量是非常庞大的,这不是一个可行的方案。
我们知道 RNN 在每一个时间步的输出为所有字符类别的概率分布,即一个包含每个字符分数的向量,我们取其中最大概率的字符作为该时间步的输出字符,然后将所有时间步得到一个字符进行拼接得到一个序列路径,即最大概率路径,再根据上面介绍的合并序列方法得到最终的预测文本结果。
在输出阶段经过 CTC 的翻译,即将网络学习到的序列特征信息转化为最终的识别文本,就可以对整个文本图像进行识别。
比如上面这个图,有 5 个时间步,字符类别有 “a”, “b” and “-” (blank),对于每个时间步的概率分布,我们都取分数最大的字符,所以得到序列路径 “aaa-b”,先移除相邻重复的字符得到 “a-b”,然后去除 blank 字符得到最终结果:“ab”。
预测过程中,先使用标准的 CNN
网络提取文本图像的特征,再利用 BLSTM
将特征向量进行融合以提取字符序列的上下文特征,然后得到每列特征的概率分布,最后通过 CTC
进行预测得到文本序列。
利用 BLSTM 和 CTC 学习到文本图像中的上下文关系,从而有效提升文本识别准确率,使得模型更加鲁棒。
在训练阶段,CRNN 将训练图像统一缩放为 w×32(w×h)
;在测试阶段,针对字符拉伸会导致识别率降低的问题,CRNN保持输入图像尺寸比例,但是图像高度还是必须统一为 32
个像素,卷积特征图的尺寸动态决定 LSTM 的时序长度(时间步长)。
import torch.nn as nn
from collections import OrderedDict
class BidirectionalLSTM(nn.Module):
def __init__(self, nIn, nHidden, nOut):
super(BidirectionalLSTM, self).__init__()
self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True)
self.embedding = nn.Linear(nHidden * 2, nOut)
def forward(self, input):
recurrent, _ = self.rnn(input)
T, b, h = recurrent.size()
t_rec = recurrent.view(T * b, h)
output = self.embedding(t_rec) # [T * b, nOut]
output = output.view(T, b, -1)
return output
class CRNN(nn.Module):
def __init__(self, imgH, nc, nclass, nh, leakyRelu=False):
super(CRNN, self).__init__()
assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
# 1x32x128
self.conv1 = nn.Conv2d(nc, 64, 3, 1, 1)
self.relu1 = nn.ReLU(True)
self.pool1 = nn.MaxPool2d(2, 2)
# 64x16x64
self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
self.relu2 = nn.ReLU(True)
self.pool2 = nn.MaxPool2d(2, 2)
# 128x8x32
self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 1)
self.bn3 = nn.BatchNorm2d(256)
self.relu3_1 = nn.ReLU(True)
self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 1)
self.relu3_2 = nn.ReLU(True)
self.pool3 = nn.MaxPool2d((2, 2), (2, 1), (0, 1))
# 256x4x16
self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 1)
self.bn4 = nn.BatchNorm2d(512)
self.relu4_1 = nn.ReLU(True)
self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 1)
self.relu4_2 = nn.ReLU(True)
self.pool4 = nn.MaxPool2d((2, 2), (2, 1), (0, 1))
# 512x2x16
self.conv5 = nn.Conv2d(512, 512, 2, 1, 0)
self.bn5 = nn.BatchNorm2d(512)
self.relu5 = nn.ReLU(True)
# 512x1x16
self.rnn = nn.Sequential(
BidirectionalLSTM(512, nh, nh),
BidirectionalLSTM(nh, nh, nclass))
def forward(self, input):
# conv features
x = self.pool1(self.relu1(self.conv1(input)))
x = self.pool2(self.relu2(self.conv2(x)))
x = self.pool3(self.relu3_2(self.conv3_2(self.relu3_1(self.bn3(self.conv3_1(x))))))
x = self.pool4(self.relu4_2(self.conv4_2(self.relu4_1(self.bn4(self.conv4_1(x))))))
conv = self.relu5(self.bn5(self.conv5(x)))
# print(conv.size())
b, c, h, w = conv.size()
assert h == 1, "the height of conv must be 1"
conv = conv.squeeze(2)
conv = conv.permute(2, 0, 1) # [w, b, c]
# rnn features
output = self.rnn(conv)
return output
class CRNN_v2(nn.Module):
def __init__(self, imgH, nc, nclass, nh, leakyRelu=False):
super(CRNN_v2, self).__init__()
assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
# 1x32x128
self.conv1_1 = nn.Conv2d(nc, 32, 3, 1, 1)
self.bn1_1 = nn.BatchNorm2d(32)
self.relu1_1 = nn.ReLU(True)
self.conv1_2 = nn.Conv2d(32, 64, 3, 1, 1)
self.bn1_2 = nn.BatchNorm2d(64)
self.relu1_2 = nn.ReLU(True)
self.pool1 = nn.MaxPool2d(2, 2)
# 64x16x64
self.conv2_1 = nn.Conv2d(64, 64, 3, 1, 1)
self.bn2_1 = nn.BatchNorm2d(64)
self.relu2_1 = nn.ReLU(True)
self.conv2_2 = nn.Conv2d(64, 128, 3, 1, 1)
self.bn2_2 = nn.BatchNorm2d(128)
self.relu2_2 = nn.ReLU(True)
self.pool2 = nn.MaxPool2d(2, 2)
# 128x8x32
self.conv3_1 = nn.Conv2d(128, 96, 3, 1, 1)
self.bn3_1 = nn.BatchNorm2d(96)
self.relu3_1 = nn.ReLU(True)
self.conv3_2 = nn.Conv2d(96, 192, 3, 1, 1)
self.bn3_2 = nn.BatchNorm2d(192)
self.relu3_2 = nn.ReLU(True)
self.pool3 = nn.MaxPool2d((2, 2), (2, 1), (0, 1))
# 192x4x32
self.conv4_1 = nn.Conv2d(192, 128, 3, 1, 1)
self.bn4_1 = nn.BatchNorm2d(128)
self.relu4_1 = nn.ReLU(True)
self.conv4_2 = nn.Conv2d(128, 256, 3, 1, 1)
self.bn4_2 = nn.BatchNorm2d(256)
self.relu4_2 = nn.ReLU(True)
self.pool4 = nn.MaxPool2d((2, 2), (2, 1), (0, 1))
# 256x2x32
self.bn5 = nn.BatchNorm2d(256)
# 256x2x32
self.rnn = nn.Sequential(
BidirectionalLSTM(512, nh, nh),
BidirectionalLSTM(nh, nh, nclass))
def forward(self, input):
# conv features
x = self.pool1(self.relu1_2(self.bn1_2(self.conv1_2(self.relu1_1(self.bn1_1(self.conv1_1(input)))))))
x = self.pool2(self.relu2_2(self.bn2_2(self.conv2_2(self.relu2_1(self.bn2_1(self.conv2_1(x)))))))
x = self.pool3(self.relu3_2(self.bn3_2(self.conv3_2(self.relu3_1(self.bn3_1(self.conv3_1(x)))))))
x = self.pool4(self.relu4_2(self.bn4_2(self.conv4_2(self.relu4_1(self.bn4_1(self.conv4_1(x)))))))
conv = self.bn5(x)
# print(conv.size())
b, c, h, w = conv.size()
assert h == 2, "the height of conv must be 2"
conv = conv.reshape([b,c*h,w])
conv = conv.permute(2, 0, 1) # [w, b, c]
# rnn features
output = self.rnn(conv)
return output
def conv3x3(nIn, nOut, stride=1):
# "3x3 convolution with padding"
return nn.Conv2d( nIn, nOut, kernel_size=3, stride=stride, padding=1, bias=False )
class basic_res_block(nn.Module):
def __init__(self, nIn, nOut, stride=1, downsample=None):
super( basic_res_block, self ).__init__()
m = OrderedDict()
m['conv1'] = conv3x3( nIn, nOut, stride )
m['bn1'] = nn.BatchNorm2d( nOut )
m['relu1'] = nn.ReLU( inplace=True )
m['conv2'] = conv3x3( nOut, nOut )
m['bn2'] = nn.BatchNorm2d( nOut )
self.group1 = nn.Sequential( m )
self.relu = nn.Sequential( nn.ReLU( inplace=True ) )
self.downsample = downsample
def forward(self, x):
if self.downsample is not None:
residual = self.downsample( x )
else:
residual = x
out = self.group1( x ) + residual
out = self.relu( out )
return out
class CRNN_res(nn.Module):
def __init__(self, imgH, nc, nclass, nh):
super(CRNN_res, self).__init__()
assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
self.conv1 = nn.Conv2d(nc, 64, 3, 1, 1)
self.relu1 = nn.ReLU(True)
self.res1 = basic_res_block(64, 64)
# 1x32x128
down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=1, stride=2, bias=False),nn.BatchNorm2d(128))
self.res2_1 = basic_res_block( 64, 128, 2, down1 )
self.res2_2 = basic_res_block(128,128)
# 64x16x64
down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=1, stride=2, bias=False),nn.BatchNorm2d(256))
self.res3_1 = basic_res_block(128, 256, 2, down2)
self.res3_2 = basic_res_block(256, 256)
self.res3_3 = basic_res_block(256, 256)
# 128x8x32
down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=1, stride=(2, 1), bias=False),nn.BatchNorm2d(512))
self.res4_1 = basic_res_block(256, 512, (2, 1), down3)
self.res4_2 = basic_res_block(512, 512)
self.res4_3 = basic_res_block(512, 512)
# 256x4x16
self.pool = nn.AvgPool2d((2, 2), (2, 1), (0, 1))
# 512x2x16
self.conv5 = nn.Conv2d(512, 512, 2, 1, 0)
self.bn5 = nn.BatchNorm2d(512)
self.relu5 = nn.ReLU(True)
# 512x1x16
self.rnn = nn.Sequential(
BidirectionalLSTM(512, nh, nh),
BidirectionalLSTM(nh, nh, nclass))
def forward(self, input):
# conv features
x = self.res1(self.relu1(self.conv1(input)))
x = self.res2_2(self.res2_1(x))
x = self.res3_3(self.res3_2(self.res3_1(x)))
x = self.res4_3(self.res4_2(self.res4_1(x)))
x = self.pool(x)
conv = self.relu5(self.bn5(self.conv5(x)))
# print(conv.size())
b, c, h, w = conv.size()
assert h == 1, "the height of conv must be 1"
conv = conv.squeeze(2)
conv = conv.permute(2, 0, 1) # [w, b, c]
# rnn features
output = self.rnn(conv)
return output
if __name__ == '__main__':
pass
- 共约 364 万张图片,按照 99:1 划分成训练集和验证集
- 数据利用中文语料库(新闻 + 文言文),通过字体、大小、灰度、模糊、透视、拉伸等变化随机生成
- 包含汉字、英文字母、数字和标点共 5990 个字符
- 每个样本固定 10 个字符,字符随机截取自语料库中的句子
- 图片分辨率统一为
280x32
关于识别数据集的生成,可参看这个项目:https://github.com/Belval/TextRecognitionDataGenerator
通过使用 Pillow 生成图片和绘上文字,并对图片随机叠加椒盐噪声,以更加贴近现实场景。
# 生成椒盐噪声
def img_salt_pepper_noise(src,percetage):
NoiseImg = src
NoiseNum = int(percetage*src.shape[0]*src.shape[1])
for i in range(NoiseNum):
randX = random.randint(0, src.shape[0]-1)
randY = random.randint(0, src.shape[1]-1)
if random.randint(0,1) == 0:
NoiseImg[randX, randY] = 0
else:
NoiseImg[randX, randY] = 255
return NoiseImg
# 随机生成不定长图片集
def gen_text(cnt):
font_path = '/data/work/tensorflow/fonts/arial.ttf' # 设置文字字体和大小
font_size = 30
font = ImageFont.truetype(font_path,font_size)
for i in range(cnt):
rnd = random.randint(1, 10) # 随机生成1到10位的不定长数字
text = ''
for j in range(rnd):
text = text + DIGITS[random.randint(0, len(DIGITS)-1)]
img = Image.new("RGB", (256,32)) # 生成图片并绘上文字
draw = ImageDraw.Draw(img)
draw.text((1,1), text, font=font, fill='white')
img = np.array(img)
img = img_salt_pepper_noise(img, float(random.randint(1,10)/100.0)) # 随机叠加椒盐噪声并保存图像
cv2.imwrite(data_dir + text + '_' + str(i+1) + '.jpg',img)
由于文字是不定长的,因此,如果读取图片并获取标签,然后将标签存放在一个紧密矩阵中进行向量化,那将会出现大量的零元素,很浪费空间。因此,使用稀疏矩阵对标签进行向量化。所谓“稀疏矩阵”就是矩阵中的零元素远远多于非零元素,采用这种方式存储可有效节约空间。
稀疏矩阵有3个属性,分别是:
将标签转为稀疏矩阵,对标签进行向量化,核心代码如下:
def sparse_tuple_from(sequences, dtype=np.int32):
""" 序列转为稀疏矩阵
input: 序列
output: indices非零坐标点,values数据值,shape稀疏矩阵大小
"""
indices = []
values = []
for n, seq in enumerate(sequences):
indices.extend(zip([n]*len(seq), range(len(seq))))
values.extend(seq)
indices = np.asarray(indices, dtype=np.int64)
values = np.asarray(values, dtype=dtype)
shape = np.asarray([len(sequences), np.asarray(indices).max(0)[1] + 1], dtype=np.int64)
return indices, values, shape
将稀疏矩阵转为标签,用于输出结果,核心代码如下:
def decode_sparse_tensor(sparse_tensor):
""" 稀疏矩阵转为序列
input: 稀疏矩阵
output: 序列
"""
decoded_indexes = list()
current_i = 0
current_seq = []
for offset, i_and_index in enumerate(sparse_tensor[0]):
i = i_and_index[0]
if i != current_i:
decoded_indexes.append(current_seq)
current_i = i
current_seq = list()
current_seq.append(offset)
decoded_indexes.append(current_seq)
result = []
for index in decoded_indexes:
result.append(decode_a_seq(index, sparse_tensor))
return result
def decode_a_seq(indexes, spars_tensor):
""" 序列编码转换 """
decoded = []
for m in indexes:
str = DIGITS[spars_tensor[1][m]]
decoded.append(str)
return decoded
读取图像数据以及进行标签向量化,以便于输入到模型进行训练,核心代码如下:
def get_file_text_array():
""" 将文件和标签读到内存,减少磁盘IO """
file_name_array = []
text_array = []
for parent, dirnames, filenames in os.walk(data_dir):
file_name_array = filenames
for f in file_name_array:
text = f.split('_')[0]
text_array.append(text)
return file_name_array,text_array
def get_next_batch(file_name_array,text_array,batch_size=64):
""" 获取训练的批量数据 """
inputs = np.zeros([batch_size, OUTPUT_SHAPE[1], OUTPUT_SHAPE[0]])
codes = []
# 获取训练样本
for i in range(batch_size):
index = random.randint(0, len(file_name_array)-1)
image = cv2.imread(data_dir + file_name_array[index])
image = cv2.resize(image, (OUTPUT_SHAPE[1], OUTPUT_SHAPE[0]), 3)
image = cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
text = text_array[index]
inputs[i,:] = np.transpose(image.reshape((OUTPUT_SHAPE[0], OUTPUT_SHAPE[1]))) # 矩阵转置
codes.append(list(text)) # 标签转成列表
targets = [np.asarray(i) for i in codes] # 标签转成稀疏矩阵
sparse_targets = sparse_tuple_from(targets)
seq_len = np.ones(inputs.shape[0]) * OUTPUT_SHAPE[1]
return inputs, sparse_targets, seq_len
def get_train_model():
inputs = tf.placeholder(tf.float32, [None, None, OUTPUT_SHAPE[0]])
targets = tf.sparse_placeholder(tf.int32) # 稀疏矩阵
seq_len = tf.placeholder(tf.int32, [None]) # 序列长度 [batch_size,]
# 定义LSTM网络
cell = tf.contrib.rnn.LSTMCell(num_hidden, state_is_tuple=True)
stack = tf.contrib.rnn.MultiRNNCell([cell] * num_layers, state_is_tuple=True)
outputs, _ = tf.nn.dynamic_rnn(cell, inputs, seq_len, dtype=tf.float32)
shape = tf.shape(inputs)
batch_s, max_timesteps = shape[0], shape[1]
outputs = tf.reshape(outputs, [-1, num_hidden])
W = tf.Variable(tf.truncated_normal([num_hidden, num_classes], stddev=0.1), name="W")
b = tf.Variable(tf.constant(0., shape=[num_classes]), name="b")
logits = tf.matmul(outputs, W) + b
logits = tf.reshape(logits, [batch_s, -1, num_classes])
# 转置矩阵
logits = tf.transpose(logits, (1, 0, 2))
return logits, inputs, targets, seq_len, W, b
在训练之前,先定义好准确率评估方法,以便于在训练过程中不断评估模型的准确性,核心代码如下:
def report_accuracy(decoded_list, test_targets):
""" 准确性评估
input: 预测结果序列 decoded_list ,目标序列 test_targets
output: 准确率
"""
original_list = decode_sparse_tensor(test_targets)
detected_list = decode_sparse_tensor(decoded_list)
true_numer = 0
# 预测序列与目标序列的维度不一致,说明有些预测失败,直接返回
if len(original_list) != len(detected_list):
print("len(original_list)", len(original_list), "len(detected_list)", len(detected_list),
" test and detect length desn't match")
return
# 比较预测序列与结果序列是否一致,并统计准确率
print("T/F: original(length) <-------> detectcted(length)")
for idx, number in enumerate(original_list):
detect_number = detected_list[idx]
hit = (number == detect_number)
print(hit, number, "(", len(number), ") <-------> ", detect_number, "(", len(detect_number), ")")
if hit:
true_numer = true_numer + 1
accuracy = true_numer * 1.0 / len(original_list)
print("Test Accuracy:", accuracy)
return accuracy
接着开始对模型进行训练,核心代码如下:
def train():
# 获取训练样本数据
file_name_array, text_array = get_file_text_array()
# 定义学习率
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(INITIAL_LEARNING_RATE,
global_step,
DECAY_STEPS,
LEARNING_RATE_DECAY_FACTOR,
staircase=True)
# 获取网络结构
logits, inputs, targets, seq_len, W, b = get_train_model()
# 设置损失函数
loss = tf.nn.ctc_loss(labels=targets, inputs=logits, sequence_length=seq_len)
cost = tf.reduce_mean(loss)
# 设置优化器
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss, global_step=global_step)
decoded, log_prob = tf.nn.ctc_beam_search_decoder(logits, seq_len, merge_repeated=False)
acc = tf.reduce_mean(tf.edit_distance(tf.cast(decoded[0], tf.int32), targets))
init = tf.global_variables_initializer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session() as session:
session.run(init)
saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
for curr_epoch in range(num_epochs):
train_cost = 0
train_ler = 0
for batch in range(BATCHES):
# 训练模型
train_inputs, train_targets, train_seq_len = get_next_batch(file_name_array, text_array, BATCH_SIZE)
feed = {inputs: train_inputs, targets: train_targets, seq_len: train_seq_len}
b_loss, b_targets, b_logits, b_seq_len, b_cost, steps, _ = session.run(
[loss, targets, logits, seq_len, cost, global_step, optimizer], feed)
# 评估模型
if steps > 0 and steps % REPORT_STEPS == 0:
test_inputs, test_targets, test_seq_len = get_next_batch(file_name_array, text_array, BATCH_SIZE)
test_feed = {inputs: test_inputs,targets: test_targets,seq_len: test_seq_len}
dd, log_probs, accuracy = session.run([decoded[0], log_prob, acc], test_feed)
report_accuracy(dd, test_targets)
# 保存识别模型
save_path = saver.save(session, model_dir + "lstm_ctc_model.ctpk",global_step=steps)
c = b_cost
train_cost += c * BATCH_SIZE
train_cost /= TRAIN_SIZE
# 计算 loss
train_inputs, train_targets, train_seq_len = get_next_batch(file_name_array, text_array, BATCH_SIZE)
val_feed = {inputs: train_inputs,targets: train_targets,seq_len: train_seq_len}
val_cost, val_ler, lr, steps = session.run([cost, acc, learning_rate, global_step], feed_dict=val_feed)
log = "{} Epoch {}/{}, steps = {}, train_cost = {:.3f}, val_cost = {:.3f}"
print(log.format(curr_epoch + 1, num_epochs, steps, train_cost, val_cost))
为了方便其它程序调用 LSTM + CTC 的识别能力,对识别能力进行封装,只需要输入一张图片,即可识别后返回结果。核心代码如下:
def predict(image):
""" LSTM+CTC 文字识别能力封装
input: 图片
output: 识别结果文字
"""
# 获取网络结构
logits, inputs, targets, seq_len, W, b = get_train_model()
decoded, log_prob = tf.nn.ctc_beam_search_decoder(logits, seq_len, merge_repeated=False)
saver = tf.train.Saver()
with tf.Session() as sess:
# 加载模型
saver.restore(sess, tf.train.latest_checkpoint(model_dir))
# 图像预处理
image = cv2.resize(image, (OUTPUT_SHAPE[1], OUTPUT_SHAPE[0]), 3)
image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
pred_inputs = np.zeros([1, OUTPUT_SHAPE[1], OUTPUT_SHAPE[0]])
pred_inputs[0, :] = np.transpose(image.reshape((OUTPUT_SHAPE[0], OUTPUT_SHAPE[1])))
pred_seq_len = np.ones(1) * OUTPUT_SHAPE[1]
# 模型预测
pred_feed = {inputs: pred_inputs,seq_len: pred_seq_len}
dd, log_probs = sess.run([decoded[0], log_prob], pred_feed)
# 识别结果转换
detected_list = decode_sparse_tensor(dd)[0]
detected_text = ''
for d in detected_list:
detected_text = detected_text + d
return detected_text
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。