| """ |
| 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/) |
| """ |
|
|
| import torch.nn as nn |
| import torch |
|
|
| ''' |
| Source - https://github.com/rishikksh20/ResUnet/blob/master/core/res_unet.py |
| An implementation of this paper - https://arxiv.org/pdf/1711.10684.pdf |
| ''' |
|
|
| class ResidualConv(nn.Module): |
| def __init__(self, input_dim, output_dim, stride, padding): |
| super(ResidualConv, self).__init__() |
|
|
| self.conv_block = nn.Sequential( |
| nn.BatchNorm2d(input_dim), |
| nn.ReLU(), |
| nn.Conv2d( |
| input_dim, output_dim, kernel_size=3, stride=stride, padding=padding |
| ), |
| nn.BatchNorm2d(output_dim), |
| nn.ReLU(), |
| nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=1), |
| ) |
| self.conv_skip = nn.Sequential( |
| nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=stride, padding=1), |
| nn.BatchNorm2d(output_dim), |
| ) |
|
|
| def forward(self, x): |
|
|
| return self.conv_block(x) + self.conv_skip(x) |
|
|
|
|
| class Upsample(nn.Module): |
| def __init__(self, input_dim, output_dim, kernel, stride): |
| super(Upsample, self).__init__() |
|
|
| self.upsample = nn.ConvTranspose2d( |
| input_dim, output_dim, kernel_size=kernel, stride=stride |
| ) |
|
|
| def forward(self, x): |
| return self.upsample(x) |
| |
| class ResUnet(nn.Module): |
| def __init__(self, in_channel=1,out_channel=512, filters=[64, 128, 256, 512]): |
| super(ResUnet, self).__init__() |
|
|
| self.input_layer = nn.Sequential( |
| nn.Conv2d(in_channel, filters[0], kernel_size=3, padding=1), |
| nn.BatchNorm2d(filters[0]), |
| nn.ReLU(), |
| nn.Conv2d(filters[0], filters[0], kernel_size=3, padding=1), |
| ) |
| self.input_skip = nn.Sequential( |
| nn.Conv2d(in_channel, filters[0], kernel_size=3, padding=1) |
| ) |
|
|
| self.residual_conv_1 = ResidualConv(filters[0], filters[1], 2, 1) |
| self.residual_conv_2 = ResidualConv(filters[1], filters[2], 2, 1) |
|
|
| self.bridge = ResidualConv(filters[2], filters[3], 2, 1) |
|
|
| self.upsample_1 = Upsample(filters[3], filters[3], 2, 2) |
| self.up_residual_conv1 = ResidualConv(filters[3] + filters[2], filters[2], 1, 1) |
|
|
| self.upsample_2 = Upsample(filters[2], filters[2], 2, 2) |
| self.up_residual_conv2 = ResidualConv(filters[2] + filters[1], filters[1], 1, 1) |
|
|
| self.upsample_3 = Upsample(filters[1], filters[1], 2, 2) |
| self.up_residual_conv3 = ResidualConv(filters[1] + filters[0], filters[0], 1, 1) |
|
|
| self.output_layer = nn.Sequential( |
| nn.Conv2d(filters[0], out_channel, 1, 1), |
| nn.Sigmoid(), |
| ) |
|
|
| def forward(self, x): |
| |
| x1 = self.input_layer(x) + self.input_skip(x) |
| x2 = self.residual_conv_1(x1) |
| x3 = self.residual_conv_2(x2) |
| |
| x4 = self.bridge(x3) |
| |
| x4 = self.upsample_1(x4) |
| x5 = torch.cat([x4, x3], dim=1) |
|
|
| x6 = self.up_residual_conv1(x5) |
|
|
| x6 = self.upsample_2(x6) |
| x7 = torch.cat([x6, x2], dim=1) |
|
|
| x8 = self.up_residual_conv2(x7) |
|
|
| x8 = self.upsample_3(x8) |
| x9 = torch.cat([x8, x1], dim=1) |
|
|
| x10 = self.up_residual_conv3(x9) |
|
|
| output = self.output_layer(x10) |
|
|
| return output |
|
|
| |
| |
| |
| |
|
|
