引言
- CRNN是经典的文本识别算法,这里主要用来夯实基础,掌握CRNN基本原理以及PyTorch实现。
基本原理
- CRNN采取的架构是CNN+RNN+CTC,
- CNN:使用深度CNN,对输入图像提取特征,得到特征图
- RNN:使用双向RNN对特征序列进行预测,对序列中每个特征向量进行学习,并输出预测标签
- CTC:使用CTC损失,把从循环层获取的一系列标签分布转换为最终的标签序列
核心代码实现(可直接复制运行)
torch.nn.CTCLoss()
的输入必须要经过logsoftmax
函数的- 以下代码包括:训练和损失,推理和解码4个部分
import string
import torch
from torch import nn
import torch.nn.functional as F
class CRNN(nn.Module):
def __init__(self, img_height, input_channel, n_class, hidden_size):
super().__init__()
if img_height % 16 != 0:
raise ValueError('img_height has to be a multiple of 16')
kernel_size = [3, 3, 3, 3, 3, 3, 2]
padding_size = [1, 1, 1, 1, 1, 1, 0]
stride = [1, 1, 1, 1, 1, 1, 1]
channel = [64, 128, 256, 256, 512, 512, 512]
def conv_relu(i, batchNormalization=False):
in_channels = input_channel if i == 0 else channel[i - 1]
out_channels = channel[i]
cnn.add_module(f'conv{i}',
nn.Conv2d(in_channels, out_channels,
kernel_size[i],
stride[i],
padding_size[i]))
if batchNormalization:
cnn.add_module(f'batchnorm{i}', nn.BatchNorm2d(out_channels))
cnn.add_module(f'relu{i}', nn.ReLU(True))
cnn = nn.Sequential()
conv_relu(0)
cnn.add_module('pooling0', nn.MaxPool2d(2, 2))
conv_relu(1)
cnn.add_module('pooling1', nn.MaxPool2d(2, 2))
conv_relu(2, True)
conv_relu(3)
cnn.add_module('pooling2',
nn.MaxPool2d(kernel_size=(2, 2),
stride=(2, 1),
padding=(0, 1)))
conv_relu(4, True)
conv_relu(5)
cnn.add_module('pooling3',
nn.MaxPool2d((2, 2), (2, 1), (0, 1)))
conv_relu(6, True)
self.cnn = cnn
self.rnn = nn.Sequential(
BidirectionalLSTM(512, hidden_size, hidden_size),
BidirectionalLSTM(hidden_size, hidden_size, n_class)
)
def forward(self, x):
cnn_feature = self.cnn(x)
h = cnn_feature.size()[2]
if h != 1:
raise ValueError("the height of cnn_feature must be 1")
cnn_feature = cnn_feature.squeeze(2)
cnn_feature = cnn_feature.permute(2, 0, 1)
output = self.rnn(cnn_feature)
x = F.log_softmax(x, dim=2)
return output
class BidirectionalLSTM(nn.Module):
def __init__(self, input_size, hidden_size, out_feature):
super().__init__()
self.rnn = nn.LSTM(input_size, hidden_size, bidirectional=True)
self.embedding = nn.Linear(hidden_size * 2, out_feature)
def forward(self, x):
recurrent, _ = self.rnn(x)
T, b, h = recurrent.size()
t_rec = recurrent.view(T * b, h)
output = self.embedding(t_rec)
output = output.view(T, b, -1)
return output
def decode(preds, preds_length):
length = preds_length[0]
char_list = []
for i in range(length):
if preds[i] != 0 and (not (i > 0 and preds[i - 1] == preds[i])):
char_list.append(alphabet[preds[i] - 1])
return ''.join(char_list)
if __name__ == '__main__':
alphabet = ['blank'] + list(string.ascii_lowercase)
num_classes = len(alphabet)
img = torch.randn((1, 1, 32, 320))
ctc_loss = nn.CTCLoss()
crnn = CRNN(32, 1, num_classes, 256)
preds = crnn(img)
_, infer_preds = preds.max(2)
infer_preds = infer_preds.transpose(1, 0).contiguous().view(-1)
preds_len = torch.IntTensor([infer_preds.shape[0]])
text = decode(infer_preds, preds_len)
print(text)
min_seq_length = 10
max_seq_length = 30
batch_size = img.shape[0]
time_step = preds.shape[0]
input_length = torch.IntTensor([time_step] * batch_size)
target = torch.randint(low=1, high=num_classes,
size=(batch_size, max_seq_length),
dtype=torch.long)
target_length = torch.randint(low=min_seq_length,
high=max_seq_length,
size=(batch_size,), dtype=torch.long)
loss = ctc_loss(preds, target, input_length, target_length)
print(preds.shape)
参考博客