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

Halcom

OCR深度学习模型CRNN+BiLSTM

1. import torch
   
2. import torch.nn as nn
   
3. 
   
4. def conv3x3(in_planes, out_planes, stride=1):
   
5.     """3x3 convolution with padding"""
   
6.     return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
   
7.                    padding=1, bias=False)
   
8. 
   
9. def conv1x1(in_planes, out_planes, stride=1):
   
10.     """1x1 convolution"""
    
11.     return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
    
12. 
    
13. class AsterBlock(nn.Module):
    
14.     def __init__(self, inplanes, planes, stride=1, downsample=None):
    
15.         super(AsterBlock, self).__init__()
    
16.         self.conv1 = conv1x1(inplanes, planes, stride)
    
17.         self.bn1 = nn.BatchNorm2d(planes)
    
18.         self.relu = nn.ReLU(inplace=True)
    
19.         self.conv2 = conv3x3(planes, planes)
    
20.         self.bn2 = nn.BatchNorm2d(planes)
    
21.         self.downsample = downsample
    
22.         self.stride = stride
    
23. 
    
24.     def forward(self, x):
    
25.         residual = x
    
26.         out = self.conv1(x)
    
27.         out = self.bn1(out)
    
28.         out = self.relu(out)
    
29.         out = self.conv2(out)
    
30.         out = self.bn2(out)
    
31. 
    
32.         if self.downsample is not None:
    
33.             residual = self.downsample(x)
    
34.         out += residual
    
35.         out = self.relu(out)
    
36.         return out
    
37. 
    
38. class ResNet_ASTER(nn.Module):
    
39.     """For aster or crnn"""
    
40. 
    
41.     def __init__(self, num_class, with_lstm=False):
    
42.         super(ResNet_ASTER, self).__init__()
    
43.         self.with_lstm = with_lstm
    
44. 
    
45.         in_channels = 1
    
46.         self.layer0 = nn.Sequential(
    
47.             nn.Conv2d(in_channels, 32, kernel_size=(3, 3), stride=1, padding=1, bias=False),
    
48.             nn.BatchNorm2d(32),
    
49.             nn.ReLU(inplace=True))
    
50. 
    
51.         self.inplanes = 32
    
52.         self.layer1 = self._make_layer(32,  3, [2, 2]) # [16]
    
53.         self.layer2 = self._make_layer(64,  4, [2, 2]) # [8]
    
54.         self.layer3 = self._make_layer(128, 6, [2, 1]) # [4]
    
55.         self.layer4 = self._make_layer(256, 6, [2, 1]) # [2]
    
56.         self.layer5 = self._make_layer(512, 3, [2, 1]) # [1]
    
57. 
    
58.         self.output_layer = nn.Linear(512,num_class)
    
59. 
    
60.         if with_lstm:
    
61.             self.rnn = nn.LSTM(512, 256, bidirectional=True, num_layers=2)
    
62.             self.out_planes = 2 * 256
    
63.         else:
    
64.             self.out_planes = 512
    
65. 
    
66.         for m in self.modules():
    
67.             if isinstance(m, nn.Conv2d):
    
68.                 nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
    
69.             elif isinstance(m, nn.BatchNorm2d):
    
70.                 nn.init.constant_(m.weight, 1)
    
71.                 nn.init.constant_(m.bias, 0)
    
72.         nn.init.normal_(self.output_layer.weight,std=0.01)
    
73.         nn.init.constant_(self.output_layer.bias,0)
    
74. 
    
75.     def _make_layer(self, planes, blocks, stride):
    
76.         downsample = None
    
77.         if stride != [1, 1] or self.inplanes != planes:
    
78.             downsample = nn.Sequential(
    
79.                 conv1x1(self.inplanes, planes, stride),
    
80.                 nn.BatchNorm2d(planes))
    
81. 
    
82.         layers = []
    
83.         layers.append(AsterBlock(self.inplanes, planes, stride, downsample))
    
84.         self.inplanes = planes
    
85.         for _ in range(1, blocks):
    
86.             layers.append(AsterBlock(self.inplanes, planes))
    
87.         return nn.Sequential(*layers)
    
88. 
    
89.     def forward(self, x):
    
90.         x0 = self.layer0(x)
    
91.         x1 = self.layer1(x0)
    
92.         x2 = self.layer2(x1)
    
93.         x3 = self.layer3(x2)
    
94.         x4 = self.layer4(x3)
    
95.         x5 = self.layer5(x4)
    
96. 
    
97.         cnn_feat = x5.squeeze(2) # [N, c, w]
    
98.         cnn_feat = cnn_feat.permute(2,0,1) #[T, b, input_size]
    
99.         if self.with_lstm:
    
100.             rnn_feat, _ = self.rnn(cnn_feat)
     
101.             T,b,h = rnn_feat.size()
     
102.             output = rnn_feat.view(T*b,h)
     
103.             output = self.output_layer(output)
     
104.             output = output.view(T,b,-1)
     
105.             output = nn.functional.log_softmax(output,dim=2)
     
106.             return output
     
107.         else:
     
108.             return cnn_feat
     
109. 
     
110. 
     
111. def get_crnn(config):
     
112.     assert config.MODEL.IMAGE_SIZE.H == 32, 'imgH has to be a multiple of 32'
     
113.     return ResNet_ASTER(config.MODEL.NUM_CLASSES + 1, True)
     
114. 
     
115. if __name__ == "__main__":
     
116.     from torchsummary import summary
     
117. 
     
118.     model = ResNet_ASTER(35, True)
     
119.     model.eval()
     
120. 
     
121.     summary(model, (3, 32, 288))

复制代码