OCR深度学习模型CRNN

Halcom

OCR深度学习模型CRNN：CNN + RNN + CTCLoss

1. import torch
   
2. import torch.nn as nn
   
3. from torch.autograd import Variable
   
4. import torchvision.models as models
   
5. import string
   
6. import numpy as np
   
7. 
   
8. class CRNN(nn.Module):
   
9.     def __init__(self,
   
10.                  abc=string.digits,
    
11.                  backend='resnet18',
    
12.                  rnn_hidden_size=128,
    
13.                  rnn_num_layers=2,
    
14.                  rnn_dropout=0,
    
15.                  seq_proj=[0, 0]):
    
16.         super(CRNN, self).__init__()
    
17. 
    
18.         self.abc = abc
    
19.         self.num_classes = len(self.abc)
    
20. 
    
21.         self.feature_extractor = getattr(models, backend)(pretrained=True)
    
22.         self.cnn = nn.Sequential(
    
23.             self.feature_extractor.conv1,
    
24.             self.feature_extractor.bn1,
    
25.             self.feature_extractor.relu,
    
26.             self.feature_extractor.maxpool,
    
27.             self.feature_extractor.layer1,
    
28.             self.feature_extractor.layer2,
    
29.             self.feature_extractor.layer3,
    
30.             self.feature_extractor.layer4
    
31.         )
    
32. 
    
33.         self.fully_conv = seq_proj[0] == 0
    
34.         if not self.fully_conv:
    
35.             self.proj = nn.Conv2d(seq_proj[0], seq_proj[1], kernel_size=1)
    
36. 
    
37.         self.rnn_hidden_size = rnn_hidden_size
    
38.         self.rnn_num_layers = rnn_num_layers
    
39.         self.rnn = nn.GRU(self.get_block_size(self.cnn),
    
40.                           rnn_hidden_size, rnn_num_layers,
    
41.                           batch_first=False,
    
42.                           dropout=rnn_dropout, bidirectional=True)
    
43.         self.linear = nn.Linear(rnn_hidden_size * 2, self.num_classes + 1)
    
44.         self.softmax = nn.Softmax(dim=2)
    
45. 
    
46.     def forward(self, x, decode=False):
    
47.         hidden = self.init_hidden(x.size(0), next(self.parameters()).is_cuda)
    
48.         features = self.cnn(x)
    
49.         features = self.features_to_sequence(features)
    
50.         seq, hidden = self.rnn(features, hidden)
    
51.         seq = self.linear(seq)
    
52.         if not self.training:
    
53.             seq = self.softmax(seq)
    
54.             if decode:
    
55.                 seq = self.decode(seq)
    
56.         return seq
    
57. 
    
58.     def init_hidden(self, batch_size, gpu=False):
    
59.         h0 = Variable(torch.zeros( self.rnn_num_layers * 2,
    
60.                                    batch_size,
    
61.                                    self.rnn_hidden_size))
    
62.         if gpu:
    
63.             h0 = h0.cuda()
    
64.         return h0
    
65. 
    
66.     def features_to_sequence(self, features):
    
67.         b, c, h, w = features.size()
    
68.         assert h == 1, "the height of out must be 1"
    
69.         if not self.fully_conv:
    
70.             features = features.permute(0, 3, 2, 1)
    
71.             features = self.proj(features)
    
72.             features = features.permute(1, 0, 2, 3)
    
73.         else:
    
74.             features = features.permute(3, 0, 2, 1)
    
75.         features = features.squeeze(2)
    
76.         return features
    
77. 
    
78.     def get_block_size(self, layer):
    
79.         return layer[-1][-1].bn2.weight.size()[0]
    
80. 
    
81.     def pred_to_string(self, pred):
    
82.         seq = []
    
83.         for i in range(pred.shape[0]):
    
84.             label = np.argmax(pred[i])
    
85.             seq.append(label - 1)
    
86.         out = []
    
87.         for i in range(len(seq)):
    
88.             if len(out) == 0:
    
89.                 if seq[i] != -1:
    
90.                     out.append(seq[i])
    
91.             else:
    
92.                 if seq[i] != -1 and seq[i] != seq[i - 1]:
    
93.                     out.append(seq[i])
    
94.         out = ''.join(self.abc[i] for i in out)
    
95.         return out
    
96. 
    
97.     def decode(self, pred):
    
98.         pred = pred.permute(1, 0, 2).cpu().data.numpy()
    
99.         seq = []
    
100.         for i in range(pred.shape[0]):
     
101.             seq.append(self.pred_to_string(pred[i]))
     
102.         return seq
     

复制代码