OCR识别网络CRNN理解与Pytorch实现-编程知识

CRNN是2015年的论文“An End-to-End Trainable Neural Network for Image-based Sequence Recognition and Its Application to Scene Text Recognition”提出的图像字符识别网络，也是目前工业界使用较为广泛的一个OCR网络。论文地址：https://arxiv.org/abs/1507.05717

1. 网络结构

CRNN是一个端到端可训练的网络，并且可处理任意长度的字符序列。CRNN得名于Convolutional Recurrent Neural Network，从名称即可看出，该网络包含了卷积网络和递归网络。实际上，CRNN由三部分组成，分别是卷积层部分（Convolutional layers）、递归层部分（Recurrent Layers）和转录层部分（Transcription Layers），如下图所示：

其中，卷积层的作用是从输入图像中提取特征，递归层则对卷积层输出的feature maps进行预测，最后，转录层将递归层的预测结果翻译成文字标签序列。CNN和RNN可由同一个损失函数进行联合训练。

在图像输入CRNN之前，需要缩放到指定高度height，宽度无限制。卷积层输出的feature maps在送入RNN之前，从左到右生成一个feature vector序列，第i个feature vector为feature maps第i列的元素的级联。这样做的好处是，每个feature vector代表了原图像上一个矩形区域的特征（感受野），使得网络能够预测不同长度的字符序列。

RNN网络的优势在于，它能够有效利用序列的上下文信息进行预测，比分别预测单个字符有更好的精确度和稳定性，同时，它对输入序列的长度无限制，比单纯使用CNN网络更加灵活。

由于传统RNN存在梯度爆炸和梯度消失问题，因此，在这篇文章中，作者采用了LTSM（Long-Short Term Memory）来克服该问题。一个LSTM包含一个记忆单元（Memory Cell）和三个乘法门（Multiplicative gates），分别为输入门（input gate）、输出门（output gate）和遗忘门（forget gate），如下图所示：

由于基于图像的文字识别具有较强的前向和后向上下文信息，因此，使用双向LSTM（bidirectional LSTM）是一个合适的选择。

转录层将RNN层的预测结果（用概率表示）映射到字符序列。在实践中，存在两种转录模式，分别是基于词典的转录，和无词典转录。在基于词典的模式中，会选择词典中最高概率的标签进行预测；而无词典模式，预测则是在无任何词典的情况下进行的。

CRNN的具体网络结构及配置如下：

2. 代码实现

网上找到一个CRNN的Pytorch实现，亲测好用，代码链接：CRNN Pytorch

网络定义：

import torch.nn as nnclass CRNN(nn.Module):def __init__(self, imgH, nc, nclass, nh, n_rnn=2, leakyRelu=False):super(CRNN, self).__init__()assert imgH % 16 == 0, 'imgH has to be a multiple of 16'ks = [3, 3, 3, 3, 3, 3, 2]ps = [1, 1, 1, 1, 1, 1, 0]ss = [1, 1, 1, 1, 1, 1, 1]nm = [64, 128, 256, 256, 512, 512, 512]cnn = nn.Sequential()def convRelu(i, batchNormalization=False):nIn = nc if i == 0 else nm[i - 1]nOut = nm[i]cnn.add_module('conv{0}'.format(i),nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i]))if batchNormalization:cnn.add_module('batchnorm{0}'.format(i), nn.BatchNorm2d(nOut))if leakyRelu:cnn.add_module('relu{0}'.format(i),nn.LeakyReLU(0.2, inplace=True))else:cnn.add_module('relu{0}'.format(i), nn.ReLU(True))convRelu(0)cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2))  # 64x16x64convRelu(1)cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2))  # 128x8x32convRelu(2, True)convRelu(3)cnn.add_module('pooling{0}'.format(2),nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 256x4x16convRelu(4, True)convRelu(5)cnn.add_module('pooling{0}'.format(3),nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 512x2x16convRelu(6, True)  # 512x1x16self.cnn = cnnself.rnn = nn.Sequential(BidirectionalLSTM(512, nh, nh),BidirectionalLSTM(nh, nh, nclass))def forward(self, input):# conv featuresconv = self.cnn(input)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 featuresoutput = self.rnn(conv)return output

其中，Bidirectional LSTM的定义如下：

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

Demo是基于字典的转录方式，可以识别0~9的1-0个数字，以及a~z的26个字母。

import torch
from torch.autograd import Variable
import utils
import dataset
from PIL import Imageimport models.crnn as crnnmodel_path = './data/crnn.pth'  # 与训练模型路径
img_path = './data/demo.png'    # 测试图片路径
alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'    # 字典model = crnn.CRNN(32, 1, 37, 256)
if torch.cuda.is_available():model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict(torch.load(model_path))converter = utils.strLabelConverter(alphabet)    # 定义字典转录函数transformer = dataset.resizeNormalize((100, 32))   # 图像预处理函数
image = Image.open(img_path).convert('L')
image = transformer(image)
if torch.cuda.is_available():image = image.cuda()
print('image size: ', image.shape)
image = image.view(1, *image.size())
image = Variable(image)model.eval()
preds = model(image)    # CRNN预测_, preds = preds.max(2)    # 找到最大概率所对应的index
preds = preds.transpose(1, 0).contiguous().view(-1)preds_size = Variable(torch.IntTensor([preds.size(0)]))
raw_pred = converter.decode(preds.data, preds_size.data, raw=True)   # 逐一输出预测字符，如a-----v--a-i-l-a-bb-l-e---
sim_pred = converter.decode(preds.data, preds_size.data, raw=False)  # 输出最终预测结果，如available
print('%-20s => %-20s' % (raw_pred, sim_pred))

demo执行结果：a-----v--a-i-l-a-bb-l-e--- => available