OCR深度学习模型CRNN+BiLSTM 模型1

Halcom

OCR深度学习模型CRNN+BiLSTM：

1. import torch.nn as nn
   
2. import torch.nn.functional as F
   
3. 
   
4. class BidirectionalLSTM(nn.Module):
   
5.     # Inputs hidden units Out
   
6.     def __init__(self, nIn, nHidden, nOut):
   
7.         super(BidirectionalLSTM, self).__init__()
   
8. 
   
9.         self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True)
   
10.         self.embedding = nn.Linear(nHidden * 2, nOut)
    
11. 
    
12.     def forward(self, input):
    
13.         recurrent, _ = self.rnn(input)
    
14.         T, b, h = recurrent.size()
    
15.         t_rec = recurrent.view(T * b, h)
    
16. 
    
17.         output = self.embedding(t_rec)  # [T * b, nOut]
    
18.         output = output.view(T, b, -1)
    
19. 
    
20.         return output
    
21. 
    
22. class CRNN(nn.Module):
    
23.     def __init__(self, imgH, nc, nclass, nh, n_rnn=2, leakyRelu=False):
    
24.         super(CRNN, self).__init__()
    
25.         assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
    
26. 
    
27.         ks = [3, 3, 3, 3, 3, 3, 2]
    
28.         ps = [1, 1, 1, 1, 1, 1, 0]
    
29.         ss = [1, 1, 1, 1, 1, 1, 1]
    
30.         nm = [64, 128, 256, 256, 512, 512, 512]
    
31. 
    
32.         cnn = nn.Sequential()
    
33. 
    
34.         def convRelu(i, batchNormalization=False):
    
35.             nIn = nc if i == 0 else nm[i - 1]
    
36.             nOut = nm[i]
    
37.             cnn.add_module('conv{0}'.format(i),
    
38.                            nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i]))
    
39.             if batchNormalization:
    
40.                 cnn.add_module('batchnorm{0}'.format(i), nn.BatchNorm2d(nOut))
    
41.             if leakyRelu:
    
42.                 cnn.add_module('relu{0}'.format(i),
    
43.                                nn.LeakyReLU(0.2, inplace=True))
    
44.             else:
    
45.                 cnn.add_module('relu{0}'.format(i), nn.ReLU(True))
    
46. 
    
47.         convRelu(0)
    
48.         cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2))  # 64x16x64
    
49.         convRelu(1)
    
50.         cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2))  # 128x8x32
    
51.         convRelu(2, True)
    
52.         convRelu(3)
    
53.         cnn.add_module('pooling{0}'.format(2),
    
54.                        nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 256x4x16
    
55.         convRelu(4, True)
    
56.         convRelu(5)
    
57.         cnn.add_module('pooling{0}'.format(3),
    
58.                        nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 512x2x16
    
59.         convRelu(6, True)  # 512x1x16
    
60. 
    
61.         self.cnn = cnn
    
62.         self.rnn = nn.Sequential(
    
63.             BidirectionalLSTM(512, nh, nh),
    
64.             BidirectionalLSTM(nh, nh, nclass))
    
65. 
    
66.     def forward(self, input):
    
67. 
    
68.         # conv features
    
69.         conv = self.cnn(input)
    
70.         b, c, h, w = conv.size()
    
71.         # print(conv.size())
    
72.         assert h == 1, "the height of conv must be 1"
    
73.         conv = conv.squeeze(2) # b *512 * width
    
74.         conv = conv.permute(2, 0, 1)  # [w, b, c]
    
75.         output = F.log_softmax(self.rnn(conv), dim=2)
    
76. 
    
77.         return output
    
78. 
    
79. def weights_init(m):
    
80.     classname = m.__class__.__name__
    
81.     if classname.find('Conv') != -1:
    
82.         m.weight.data.normal_(0.0, 0.02)
    
83.     elif classname.find('BatchNorm') != -1:
    
84.         m.weight.data.normal_(1.0, 0.02)
    
85.         m.bias.data.fill_(0)
    
86. 
    
87. def get_crnn(config):
    
88. 
    
89.     model = CRNN(config.MODEL.IMAGE_SIZE.H, 1, config.MODEL.NUM_CLASSES + 1, config.MODEL.NUM_HIDDEN)
    
90.     model.apply(weights_init)
    
91. 
    
92.     return model

复制代码