『带你学AI』一文带你搞懂OCR识别算法CRNN:解析+源码

浪子不回头ぞ 提交于 2021-01-24 14:03:26

来源 | 极简AI

头图 | 视觉中国

本文搬运自小宋的小伙伴:https://blog.csdn.net/libo1004,欢迎大家关注文章链接:https://blog.csdn.net/libo1004/article/details/111595054

前言

文字识别是AI的一个重要应用场景,文字识别过程一般由图像输入、预处理、文本检测、文本识别、结果输出等环节组成。

其中,文本检测、文本识别是最核心的环节。文本检测方面,在我的 OCR_detection 专栏相关文章中已介绍过了多种基于深度学习的方法(有的还没完成,待整理后都会放入该专栏),可针对各种场景实现对文字的检测,详请见专栏中的相关文章。

在以前的 OCR 任务中,识别过程分为两步:单字切割 和 分类任务。我们一般都会将一连串文字的文本文件先利用 投影法 切割出单个字体,再送入 CNN 里进行文字分类。但是此法已经有点过时了,现在更流行的是基于深度学习的端到端的文字识别,即我们不需要显式加入文字切割这个环节,而是将文字识别转化为序列学习问题,虽然输入的图像尺度不同,文本长度不同,但是经过 DCNN 和 RNN 后,在输出阶段经过一定的 CTC 翻译转录后,就可以对整个文本图像进行识别,也就是说,文字的切割也被融入到深度学习中去了。

现今基于深度学习的端到端 OCR 技术有两大主流技术:CRNN OCR 和 attention OCR。其实这两大方法主要区别在于最后的输出层(翻译层),即怎么将网络学习到的序列特征信息转化为最终的识别结果。这两大主流技术在其特征学习阶段都采用了 CNN+RNN 的网络结构,CRNN OCR 在对齐时采取的方式是 CTC 算法,而 attention OCR 采取的方式则是 attention 机制。本部分主要介绍应用更为广泛的 CRNN 算法。

CRNN 介绍

CRNN 模型,即将 CNN 与 RNN 网络结合,共同训练。主要用于在一定程度上实现端到端(end-to-end)地对不定长的文本序列进行识别,不用先对单个文字进行切割,而是将文本识别转化为时序依赖的序列学习问题,就是基于图像的序列识别。(说一定程度是因为虽然输入图像不需要精确给出每个字符的位置信息,但实际上还是需要对原始的图像进行前期的裁切工作)

  1. 构建 CRNN 输入特征序列;

  2. 其中还涉及到了 CTC 模块,目的是对其输入输出结果

整个CRNN网络结构包含三部分,从下到上依次为:

  1. CNN(卷积层):使用深度 CNN,对输入图像提取特征,得到特征图;

  2. RNN(循环层):使用 双向RNN(BLSTM)对特征序列进行预测,对序列中的每个特征向量进行学习,并输出预测标签(真实值)分布;

  3. CTC loss(转录层):使用 CTC 损失,把从循环层获取的一系列标签分布转换成最终的标签序列。

CRNN 网络结构

1.CNN

这里有一个很精彩的改动,一共有四个最大池化层,但是最后两个池化层的窗口尺寸由 2x2 改为 1x2,也就是图片的高度减半了四次(除以 2 4 2^424),而宽度则只减半了两次(除以 2 2 2^222),这是因为文本图像多数都是高较小而宽较长,所以其 feature map 也是这种高小宽长的矩形形状,如果使用 1×2 的池化窗口可以尽量保证不丢失在宽度方向的信息,更适合英文字母识别(比如区分 i 和 l)。

CRNN 还引入了 Batch Normalization 模块,加速模型收敛,缩短训练过程。

例如:

  1. 输入图像为灰度图像(单通道);

  2. 高度为32,这是固定的,图片通过 CNN 后,高度就变为 1,这点很重要;

  3. 宽度为160,宽度也可以为其他的值,但需要统一,所以输入 CNN 的数据尺寸为 (channel, height, width)=(1, 32, 160)。

  4. 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,所以最后 CNN 输出的 feature map 的高度为1。assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
在程序中,图像的 h 必须为 16 的整数倍。assert h == 1, "the height of conv must be 1"
前向传播时,CNN 得到的 feature map 的 h 必须为 1。
最后 CNN 得到的 feature map 尺度为 512x1x16


2.Map-to-Sequence

不能直接把 CNN 得到的特征图送入 RNN 进行训练的,需要进行一些调整,根据特征图提取 RNN 需要的特征向量序列。

现在需要从 CNN 模型产生的特征图中提取特征向量序列,每一个特征向量(如上图中的一个红色框)在特征图上 按列 从左到右生成,每一列包含 512 维特征,这意味着第 i 个特征向量是所有的特征图第 i 列像素的连接,这些特征向量就构成一个序列。

由于卷积层,最大池化层和激活函数在局部区域上执行,因此它们是平移不变的。因此,特征图的每列(即一个特征向量)对应于原始图像的一个矩形区域(称为感受野),并且这些矩形区域与特征图上从左到右的相应列具有相同的顺序。特征序列中的每个向量关联一个感受野。如下图所示:

这些特征向量序列就作为循环层的输入,每个特征向量作为 RNN 在一个时间步(time step)的输入。

3.RNN

因为 RNN 有梯度消失的问题,不能获取更多上下文信息,所以 CRNN 中使用的是 LSTM,LSTM 的特殊设计允许它捕获长距离依赖。

LSTM 是单向的,它只使用过去的信息。然而,在基于图像的序列中,两个方向的上下文是相互有用且互补的。将两个 LSTM,一个向前和一个向后组合到一个双向 LSTM 中。此外,可以堆叠多层双向 LSTM,深层结构允许比浅层抽象更高层次的抽象。

这里采用的是两层各 256 单元的双向 LSTM 网络:

通过上面一步,我们得到了 40 个特征向量,每个特征向量长度为 512,在 LSTM 中一个时间步就传入一个特征向量进行分类,这里一共有 40 个时间步。

我们知道一个特征向量就相当于原图中的一个小矩形区域,RNN 的目标就是预测这个矩形区域为哪个字符,即根据输入的特征向量,进行预测,得到所有字符的 softmax 概率分布,这是一个长度为字符类别数的向量,作为 CTC 层的输入。

因为每个时间步都会有一个输入特征向量 x t x_txt,输出一个所有字符的概率分布 y t y_tyt,所以输出为 40 个长度为字符类别数的向量构成的后验概率矩阵。如下图所示:

然后将这个后验概率矩阵传入转录层。

  1. 该部分的源码如下:

    self.rnn = nn.Sequential(
        BidirectionalLSTM(512, nh, nh),
        BidirectionalLSTM(nh, nh, nclass)

    )

  2. 然后参数设置如下:

    nh = 256
    nclass = len(opt.alphabet) + 1
    nc = 1


  3. 其中 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














  4. 所以第一次 LSTM 得到的 output=[40*256,256],然后 view 成 output=[40,256,256]

  5. 第二次 LSTM 得到的结果是 output=[40*256,nclass],然后 view 成 output=[40,256,nclass]

4.CTC Loss

这算是 CRNN 最难的地方,这一层为转录层,转录是将 RNN 对每个特征向量所做的预测转换成标签序列的过程。数学上,转录是根据每帧预测找到具有最高概率组合的标签序列。

端到端 OCR 识别的难点在于怎么处理不定长序列对齐的问题!OCR 可建模为时序依赖的文本图像问题,然后使用 CTC(Connectionist Temporal Classification, CTC)的损失函数来对 CNN 和 RNN 进行端到端的联合训练。

4.1 序列合并机制

我们现在要将 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”),但是与图像的对齐方式不同。更总结地说,一个文本标签存在一条或多条的路径。

4.2 训练阶段

在训练阶段,我们需要根据这些概率分布向量和相应的文本标签得到损失函数,从而训练神经网路模型,下面来看看如何得到损失函数的。

如上图,对于最简单的时序为 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 不需要对原始的输入字符序列进行准确的切分。

4.3 测试阶段

在测试阶段与训练阶段有所不同,我们用训练好的神经网络来识别新的文本图像。这时候我们事先不知道任何文本,如果我们像上面一样将每种可能文本的所有路径计算出来,对于很长的时间步和很长的字符序列来说,这个计算量是非常庞大的,这不是一个可行的方案。

我们知道 RNN 在每一个时间步的输出为所有字符类别的概率分布,即一个包含每个字符分数的向量,我们取其中最大概率的字符作为该时间步的输出字符,然后将所有时间步得到一个字符进行拼接得到一个序列路径,即最大概率路径,再根据上面介绍的合并序列方法得到最终的预测文本结果。

在输出阶段经过 CTC 的翻译,即将网络学习到的序列特征信息转化为最终的识别文本,就可以对整个文本图像进行识别。

比如上面这个图,有 5 个时间步,字符类别有 “a”, “b” and “-” (blank),对于每个时间步的概率分布,我们都取分数最大的字符,所以得到序列路径 “aaa-b”,先移除相邻重复的字符得到 “a-b”,然后去除 blank 字符得到最终结果:“ab”。

CRNN 小结

预测过程中,先使用标准的 CNN 网络提取文本图像的特征,再利用 BLSTM 将特征向量进行融合以提取字符序列的上下文特征,然后得到每列特征的概率分布,最后通过 CTC 进行预测得到文本序列。

利用 BLSTM 和 CTC 学习到文本图像中的上下文关系,从而有效提升文本识别准确率,使得模型更加鲁棒。

在训练阶段,CRNN 将训练图像统一缩放为 w×32(w×h);在测试阶段,针对字符拉伸会导致识别率降低的问题,CRNN保持输入图像尺寸比例,但是图像高度还是必须统一为 32 个像素,卷积特征图的尺寸动态决定 LSTM 的时序长度(时间步长)。

CRNN 网络模型搭建

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

参考链接

  1. https://blog.csdn.net/wa1tzy/article/details/107357911

  2. https://blog.csdn.net/qq_24819773/article/details/104605994

  3. https://mp.weixin.qq.com/s/p801KZ5kv5aYnLvlahFlnA

更多精彩推荐
标签
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!