File size: 3,861 Bytes
77f8d5f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
"""
Paper: "UTRNet: High-Resolution Urdu Text Recognition In Printed Documents" presented at ICDAR 2023
Authors: Abdur Rahman, Arjun Ghosh, Chetan Arora
GitHub Repository: https://github.com/abdur75648/UTRNet-High-Resolution-Urdu-Text-Recognition
Project Website: https://abdur75648.github.io/UTRNet/
Copyright (c) 2023-present: This work is licensed under the Creative Commons Attribution-NonCommercial
4.0 International License (http://creativecommons.org/licenses/by-nc/4.0/)
"""
from torch import nn
import torch.nn.functional as F
# For Gated RCNN
class GRCL(nn.Module):
def __init__(self, input_channel, output_channel, num_iteration, kernel_size, pad):
super(GRCL, self).__init__()
self.wgf_u = nn.Conv2d(input_channel, output_channel, 1, 1, 0, bias=False)
self.wgr_x = nn.Conv2d(output_channel, output_channel, 1, 1, 0, bias=False)
self.wf_u = nn.Conv2d(input_channel, output_channel, kernel_size, 1, pad, bias=False)
self.wr_x = nn.Conv2d(output_channel, output_channel, kernel_size, 1, pad, bias=False)
self.BN_x_init = nn.BatchNorm2d(output_channel)
self.num_iteration = num_iteration
self.GRCL = [GRCL_unit(output_channel) for _ in range(num_iteration)]
self.GRCL = nn.Sequential(*self.GRCL)
def forward(self, input):
""" The input of GRCL is consistant over time t, which is denoted by u(0)
thus wgf_u / wf_u is also consistant over time t.
"""
wgf_u = self.wgf_u(input)
wf_u = self.wf_u(input)
x = F.relu(self.BN_x_init(wf_u))
for i in range(self.num_iteration):
x = self.GRCL[i](wgf_u, self.wgr_x(x), wf_u, self.wr_x(x))
return x
class GRCL_unit(nn.Module):
def __init__(self, output_channel):
super(GRCL_unit, self).__init__()
self.BN_gfu = nn.BatchNorm2d(output_channel)
self.BN_grx = nn.BatchNorm2d(output_channel)
self.BN_fu = nn.BatchNorm2d(output_channel)
self.BN_rx = nn.BatchNorm2d(output_channel)
self.BN_Gx = nn.BatchNorm2d(output_channel)
def forward(self, wgf_u, wgr_x, wf_u, wr_x):
G_first_term = self.BN_gfu(wgf_u)
G_second_term = self.BN_grx(wgr_x)
G = F.sigmoid(G_first_term + G_second_term)
x_first_term = self.BN_fu(wf_u)
x_second_term = self.BN_Gx(self.BN_rx(wr_x) * G)
x = F.relu(x_first_term + x_second_term)
return x
class RCNN(nn.Module):
""" FeatureExtractor of GRCNN (https://papers.nips.cc/paper/6637-gated-recurrent-convolution-neural-network-for-ocr.pdf) """
def __init__(self, input_channel=1, output_channel=512):
super(RCNN, self).__init__()
self.output_channel = [int(output_channel / 8), int(output_channel / 4),
int(output_channel / 2), output_channel] # [64, 128, 256, 512]
self.ConvNet = nn.Sequential(
nn.Conv2d(input_channel, self.output_channel[0], 3, 1, 1), nn.ReLU(True),
nn.MaxPool2d(2, 2), # 64 x 16 x 50
GRCL(self.output_channel[0], self.output_channel[0], num_iteration=5, kernel_size=3, pad=1),
nn.MaxPool2d(2, 2), # 64 x 8 x 25
GRCL(self.output_channel[0], self.output_channel[1], num_iteration=5, kernel_size=3, pad=1),
nn.MaxPool2d(2, (2, 1), (0, 1)), # 128 x 4 x 26
GRCL(self.output_channel[1], self.output_channel[2], num_iteration=5, kernel_size=3, pad=1),
nn.MaxPool2d(2, (2, 1), (0, 1)), # 256 x 2 x 27
nn.Conv2d(self.output_channel[2], self.output_channel[3], 2, 1, 0, bias=False),
nn.BatchNorm2d(self.output_channel[3]), nn.ReLU(True)) # 512 x 1 x 26
def forward(self, input):
return self.ConvNet(input)
# import torch
# x = torch.randn(1, 1, 32, 400)
# net = RCNN()
# out = net(x)
# print(out.shape) |