Bengali CRNN architecture

bengali/shobdo_bengali.py

@@ -1,52 +1,39 @@
-
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 class ConvBnRelu(nn.Module):
-    def __init__(self, in_ch, out_ch, k=3, s=1, p=1):
+    def __init__(self, i, o, k=3, s=1, p=1):
         super().__init__()
-        self.block = nn.Sequential(
-            nn.Conv2d(in_ch, out_ch, k, s, p, bias=False),
-            nn.BatchNorm2d(out_ch),
-            nn.ReLU(inplace=True)
-        )
-    def forward(self, x): return self.block(x)
+        self.b = nn.Sequential(nn.Conv2d(i,o,k,s,p,bias=False),nn.BatchNorm2d(o),nn.ReLU(inplace=True))
+    def forward(self, x): return self.b(x)
 
 class LightCNN(nn.Module):
     def __init__(self):
         super().__init__()
-        self.b1 = nn.Sequential(ConvBnRelu(1,32),ConvBnRelu(32,32),nn.MaxPool2d(2,2))
-        self.b2 = nn.Sequential(ConvBnRelu(32,64),ConvBnRelu(64,64),nn.MaxPool2d(2,2))
-        self.b3 = nn.Sequential(ConvBnRelu(64,128),ConvBnRelu(128,128),nn.MaxPool2d((2,1)))
-        self.b4 = nn.Sequential(ConvBnRelu(128,256),ConvBnRelu(256,256),nn.MaxPool2d((2,1)))
-        self.b5 = nn.Sequential(ConvBnRelu(256,256),ConvBnRelu(256,256))
-        self.pool = nn.AdaptiveAvgPool2d((1, None))
+        self.b1=nn.Sequential(ConvBnRelu(1,32),ConvBnRelu(32,32),nn.MaxPool2d(2,2))
+        self.b2=nn.Sequential(ConvBnRelu(32,64),ConvBnRelu(64,64),nn.MaxPool2d(2,2))
+        self.b3=nn.Sequential(ConvBnRelu(64,128),ConvBnRelu(128,128),nn.MaxPool2d((2,1)))
+        self.b4=nn.Sequential(ConvBnRelu(128,256),ConvBnRelu(256,256),nn.MaxPool2d((2,1)))
+        self.b5=nn.Sequential(ConvBnRelu(256,256),ConvBnRelu(256,256))
+        self.pool=nn.AdaptiveAvgPool2d((1,None))
     def forward(self, x):
-        for b in [self.b1,self.b2,self.b3,self.b4,self.b5]: x = b(x)
+        for b in [self.b1,self.b2,self.b3,self.b4,self.b5]: x=b(x)
         return self.pool(x).squeeze(2)
 
-class BidirectionalLSTM(nn.Module):
-    def __init__(self, input_size, hidden_size, output_size):
+class BiLSTMBlock(nn.Module):
+    def __init__(self, i, h, o):
         super().__init__()
-        self.rnn = nn.LSTM(input_size, hidden_size, bidirectional=True, batch_first=True)
-        self.linear = nn.Linear(hidden_size*2, output_size)
+        self.rnn=nn.LSTM(i,h,bidirectional=True,batch_first=True)
+        self.fc=nn.Linear(h*2,o)
     def forward(self, x):
-        out, _ = self.rnn(x)
-        return self.linear(out)
+        o,_=self.rnn(x); return self.fc(o)
 
 class Model(nn.Module):
-    """EasyOCR requires the class to be named Model"""
-    def __init__(self, input_channel, output_channel, hidden_size, num_class):
+    def __init__(self, input_channel=1, output_channel=256, hidden_size=256, num_class=149):
         super().__init__()
-        self.cnn = LightCNN()
-        self.rnn = nn.Sequential(
-            BidirectionalLSTM(256, hidden_size, hidden_size),
-            BidirectionalLSTM(hidden_size, hidden_size, num_class)
-        )
+        self.cnn=LightCNN()
+        self.rnn=nn.Sequential(BiLSTMBlock(256,hidden_size,hidden_size),BiLSTMBlock(hidden_size,hidden_size,num_class))
     def forward(self, x):
-        f = self.cnn(x)
-        f = f.permute(0,2,1)
-        o = self.rnn(f)
-        o = o.permute(1,0,2)
-        return F.log_softmax(o, dim=2)
+        f=self.cnn(x).permute(0,2,1)
+        o=self.rnn(f).permute(1,0,2)
+        return F.log_softmax(o,dim=2)