openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/modeling/encoders/rec_mv1_enhance.py

@@ -0,0 +1,254 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from openrec.modeling.common import Activation
+
+
+class ConvBNLayer(nn.Module):
+
+    def __init__(
+        self,
+        num_channels,
+        filter_size,
+        num_filters,
+        stride,
+        padding,
+        num_groups=1,
+        act='hard_swish',
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.act = act
+        self._conv = nn.Conv2d(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=padding,
+            groups=num_groups,
+            bias=False,
+        )
+
+        self._batch_norm = nn.BatchNorm2d(num_filters, )
+        if self.act is not None:
+            self._act = Activation(act_type=act, inplace=True)
+
+    def forward(self, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        if self.act is not None:
+            y = self._act(y)
+        return y
+
+
+class DepthwiseSeparable(nn.Module):
+
+    def __init__(
+        self,
+        num_channels,
+        num_filters1,
+        num_filters2,
+        num_groups,
+        stride,
+        scale,
+        dw_size=3,
+        padding=1,
+        use_se=False,
+    ):
+        super(DepthwiseSeparable, self).__init__()
+        self._depthwise_conv = ConvBNLayer(
+            num_channels=num_channels,
+            num_filters=int(num_filters1 * scale),
+            filter_size=dw_size,
+            stride=stride,
+            padding=padding,
+            num_groups=int(num_groups * scale),
+        )
+        self._se = None
+        if use_se:
+            self._se = SEModule(int(num_filters1 * scale))
+        self._pointwise_conv = ConvBNLayer(
+            num_channels=int(num_filters1 * scale),
+            filter_size=1,
+            num_filters=int(num_filters2 * scale),
+            stride=1,
+            padding=0,
+        )
+
+    def forward(self, inputs):
+        y = self._depthwise_conv(inputs)
+        if self._se is not None:
+            y = self._se(y)
+        y = self._pointwise_conv(y)
+        return y
+
+
+class MobileNetV1Enhance(nn.Module):
+
+    def __init__(self,
+                 in_channels=3,
+                 scale=0.5,
+                 last_conv_stride=1,
+                 last_pool_type='max',
+                 **kwargs):
+        super().__init__()
+        self.scale = scale
+        self.block_list = []
+
+        self.conv1 = ConvBNLayer(
+            num_channels=in_channels,
+            filter_size=3,
+            num_filters=int(32 * scale),
+            stride=2,
+            padding=1,
+        )
+
+        conv2_1 = DepthwiseSeparable(
+            num_channels=int(32 * scale),
+            num_filters1=32,
+            num_filters2=64,
+            num_groups=32,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv2_1)
+
+        conv2_2 = DepthwiseSeparable(
+            num_channels=int(64 * scale),
+            num_filters1=64,
+            num_filters2=128,
+            num_groups=64,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv2_2)
+
+        conv3_1 = DepthwiseSeparable(
+            num_channels=int(128 * scale),
+            num_filters1=128,
+            num_filters2=128,
+            num_groups=128,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv3_1)
+
+        conv3_2 = DepthwiseSeparable(
+            num_channels=int(128 * scale),
+            num_filters1=128,
+            num_filters2=256,
+            num_groups=128,
+            stride=(2, 1),
+            scale=scale,
+        )
+        self.block_list.append(conv3_2)
+
+        conv4_1 = DepthwiseSeparable(
+            num_channels=int(256 * scale),
+            num_filters1=256,
+            num_filters2=256,
+            num_groups=256,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv4_1)
+
+        conv4_2 = DepthwiseSeparable(
+            num_channels=int(256 * scale),
+            num_filters1=256,
+            num_filters2=512,
+            num_groups=256,
+            stride=(2, 1),
+            scale=scale,
+        )
+        self.block_list.append(conv4_2)
+
+        for _ in range(5):
+            conv5 = DepthwiseSeparable(
+                num_channels=int(512 * scale),
+                num_filters1=512,
+                num_filters2=512,
+                num_groups=512,
+                stride=1,
+                dw_size=5,
+                padding=2,
+                scale=scale,
+                use_se=False,
+            )
+            self.block_list.append(conv5)
+
+        conv5_6 = DepthwiseSeparable(
+            num_channels=int(512 * scale),
+            num_filters1=512,
+            num_filters2=1024,
+            num_groups=512,
+            stride=(2, 1),
+            dw_size=5,
+            padding=2,
+            scale=scale,
+            use_se=True,
+        )
+        self.block_list.append(conv5_6)
+
+        conv6 = DepthwiseSeparable(
+            num_channels=int(1024 * scale),
+            num_filters1=1024,
+            num_filters2=1024,
+            num_groups=1024,
+            stride=last_conv_stride,
+            dw_size=5,
+            padding=2,
+            use_se=True,
+            scale=scale,
+        )
+        self.block_list.append(conv6)
+
+        self.block_list = nn.Sequential(*self.block_list)
+        if last_pool_type == 'avg':
+            self.pool = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+        else:
+            self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.out_channels = int(1024 * scale)
+
+    def forward(self, inputs):
+        y = self.conv1(inputs)
+        y = self.block_list(y)
+        y = self.pool(y)
+        return y
+
+
+def hardsigmoid(x):
+    return F.relu6(x + 3.0, inplace=True) / 6.0
+
+
+class SEModule(nn.Module):
+
+    def __init__(self, channel, reduction=4):
+        super(SEModule, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv1 = nn.Conv2d(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+
+    def forward(self, inputs):
+        outputs = self.avg_pool(inputs)
+        outputs = self.conv1(outputs)
+        outputs = F.relu(outputs)
+        outputs = self.conv2(outputs)
+        outputs = hardsigmoid(outputs)
+        x = torch.mul(inputs, outputs)
+
+        return x

openocr/openrec/modeling/encoders/rec_nrtr_mtb.py

@@ -0,0 +1,37 @@
+import torch
+from torch import nn
+
+
+class MTB(nn.Module):
+
+    def __init__(self, cnn_num, in_channels):
+        super(MTB, self).__init__()
+        self.block = nn.Sequential()
+        self.out_channels = in_channels
+        self.cnn_num = cnn_num
+        if self.cnn_num == 2:
+            for i in range(self.cnn_num):
+                self.block.add_module(
+                    'conv_{}'.format(i),
+                    nn.Conv2d(
+                        in_channels=in_channels if i == 0 else 32 *
+                        (2**(i - 1)),
+                        out_channels=32 * (2**i),
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                    ),
+                )
+                self.block.add_module('relu_{}'.format(i), nn.ReLU())
+                self.block.add_module('bn_{}'.format(i),
+                                      nn.BatchNorm2d(32 * (2**i)))
+
+    def forward(self, images):
+        x = self.block(images)
+        if self.cnn_num == 2:
+            # (b, w, h, c)
+            x = x.permute(0, 3, 2, 1)
+            x_shape = x.shape
+            x = torch.reshape(
+                x, (x_shape[0], x_shape[1], x_shape[2] * x_shape[3]))
+        return x

openocr/openrec/modeling/encoders/rec_resnet_31.py

@@ -0,0 +1,213 @@
+import torch.nn as nn
+
+__all__ = ['ResNet31']
+
+
+def conv3x3(in_channel, out_channel, stride=1):
+    return nn.Conv2d(in_channel,
+                     out_channel,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=1,
+                     bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_channels, channels, stride=1, downsample=False):
+        super().__init__()
+        self.conv1 = conv3x3(in_channels, channels, stride)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.relu = nn.ReLU()
+        self.conv2 = conv3x3(channels, channels)
+        self.bn2 = nn.BatchNorm2d(channels)
+        self.downsample = downsample
+        if downsample:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_channels,
+                          channels * self.expansion,
+                          1,
+                          stride,
+                          bias=False),
+                nn.BatchNorm2d(channels * self.expansion),
+            )
+        else:
+            self.downsample = nn.Sequential()
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet31(nn.Module):
+    """
+    Args:
+        in_channels (int): Number of channels of input image tensor.
+        layers (list[int]): List of BasicBlock number for each stage.
+        channels (list[int]): List of out_channels of Conv2d layer.
+        out_indices (None | Sequence[int]): Indices of output stages.
+        last_stage_pool (bool): If True, add `MaxPool2d` layer to last stage.
+    """
+
+    def __init__(
+        self,
+        in_channels=3,
+        layers=[1, 2, 5, 3],
+        channels=[64, 128, 256, 256, 512, 512, 512],
+        out_indices=None,
+        last_stage_pool=False,
+    ):
+        super(ResNet31, self).__init__()
+        assert isinstance(in_channels, int)
+        assert isinstance(last_stage_pool, bool)
+
+        self.out_indices = out_indices
+        self.last_stage_pool = last_stage_pool
+
+        # conv 1 (Conv Conv)
+        self.conv1_1 = nn.Conv2d(in_channels,
+                                 channels[0],
+                                 kernel_size=3,
+                                 stride=1,
+                                 padding=1)
+        self.bn1_1 = nn.BatchNorm2d(channels[0])
+        self.relu1_1 = nn.ReLU(inplace=True)
+
+        self.conv1_2 = nn.Conv2d(channels[0],
+                                 channels[1],
+                                 kernel_size=3,
+                                 stride=1,
+                                 padding=1)
+        self.bn1_2 = nn.BatchNorm2d(channels[1])
+        self.relu1_2 = nn.ReLU(inplace=True)
+
+        # conv 2 (Max-pooling, Residual block, Conv)
+        self.pool2 = nn.MaxPool2d(kernel_size=2,
+                                  stride=2,
+                                  padding=0,
+                                  ceil_mode=True)
+        self.block2 = self._make_layer(channels[1], channels[2], layers[0])
+        self.conv2 = nn.Conv2d(channels[2],
+                               channels[2],
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.bn2 = nn.BatchNorm2d(channels[2])
+        self.relu2 = nn.ReLU(inplace=True)
+
+        # conv 3 (Max-pooling, Residual block, Conv)
+        self.pool3 = nn.MaxPool2d(kernel_size=2,
+                                  stride=2,
+                                  padding=0,
+                                  ceil_mode=True)
+        self.block3 = self._make_layer(channels[2], channels[3], layers[1])
+        self.conv3 = nn.Conv2d(channels[3],
+                               channels[3],
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.bn3 = nn.BatchNorm2d(channels[3])
+        self.relu3 = nn.ReLU(inplace=True)
+
+        # conv 4 (Max-pooling, Residual block, Conv)
+        self.pool4 = nn.MaxPool2d(kernel_size=(2, 1),
+                                  stride=(2, 1),
+                                  padding=0,
+                                  ceil_mode=True)
+        self.block4 = self._make_layer(channels[3], channels[4], layers[2])
+        self.conv4 = nn.Conv2d(channels[4],
+                               channels[4],
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.bn4 = nn.BatchNorm2d(channels[4])
+        self.relu4 = nn.ReLU(inplace=True)
+
+        # conv 5 ((Max-pooling), Residual block, Conv)
+        self.pool5 = None
+        if self.last_stage_pool:
+            self.pool5 = nn.MaxPool2d(kernel_size=2,
+                                      stride=2,
+                                      padding=0,
+                                      ceil_mode=True)
+        self.block5 = self._make_layer(channels[4], channels[5], layers[3])
+        self.conv5 = nn.Conv2d(channels[5],
+                               channels[5],
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.bn5 = nn.BatchNorm2d(channels[5])
+        self.relu5 = nn.ReLU(inplace=True)
+
+        self.out_channels = channels[-1]
+
+    def _make_layer(self, input_channels, output_channels, blocks):
+        layers = []
+        for _ in range(blocks):
+            downsample = None
+            if input_channels != output_channels:
+                downsample = nn.Sequential(
+                    nn.Conv2d(
+                        input_channels,
+                        output_channels,
+                        kernel_size=1,
+                        stride=1,
+                        bias=False,
+                    ),
+                    nn.BatchNorm2d(output_channels),
+                )
+
+            layers.append(
+                BasicBlock(input_channels,
+                           output_channels,
+                           downsample=downsample))
+            input_channels = output_channels
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1_1(x)
+        x = self.bn1_1(x)
+        x = self.relu1_1(x)
+
+        x = self.conv1_2(x)
+        x = self.bn1_2(x)
+        x = self.relu1_2(x)
+
+        outs = []
+        for i in range(4):
+            layer_index = i + 2
+            pool_layer = getattr(self, 'pool{}'.format(layer_index))
+            block_layer = getattr(self, 'block{}'.format(layer_index))
+            conv_layer = getattr(self, 'conv{}'.format(layer_index))
+            bn_layer = getattr(self, 'bn{}'.format(layer_index))
+            relu_layer = getattr(self, 'relu{}'.format(layer_index))
+
+            if pool_layer is not None:
+                x = pool_layer(x)
+            x = block_layer(x)
+            x = conv_layer(x)
+            x = bn_layer(x)
+            x = relu_layer(x)
+
+            outs.append(x)
+
+        if self.out_indices is not None:
+            return tuple([outs[i] for i in self.out_indices])
+
+        return x

openocr/openrec/modeling/encoders/rec_resnet_45.py

@@ -0,0 +1,183 @@
+import math
+
+import numpy as np
+import torch.nn as nn
+
+from openrec.modeling.common import Block
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=1,
+                     stride=stride,
+                     bias=False)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding."""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=1,
+                     bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet45(nn.Module):
+
+    def __init__(
+        self,
+        in_channels=3,
+        block=BasicBlock,
+        layers=[3, 4, 6, 6, 3],
+        strides=[2, 1, 2, 1, 1],
+        last_stage=False,
+        out_channels=256,
+        trans_layer=0,
+        out_dim=384,
+        feat2d=True,
+        return_list=False,
+    ):
+        super(ResNet45, self).__init__()
+        self.inplanes = 32
+        self.conv1 = nn.Conv2d(in_channels,
+                               32,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.layer1 = self._make_layer(block, 32, layers[0], stride=strides[0])
+        self.layer2 = self._make_layer(block, 64, layers[1], stride=strides[1])
+        self.layer3 = self._make_layer(block,
+                                       128,
+                                       layers[2],
+                                       stride=strides[2])
+        self.layer4 = self._make_layer(block,
+                                       256,
+                                       layers[3],
+                                       stride=strides[3])
+        self.layer5 = self._make_layer(block,
+                                       512,
+                                       layers[4],
+                                       stride=strides[4])
+        self.out_channels = 512
+        self.feat2d = feat2d
+        self.return_list = return_list
+        if trans_layer > 0:
+            dpr = np.linspace(0, 0.1, trans_layer)
+            blocks = [nn.Linear(512, out_dim)] + [
+                Block(dim=out_dim,
+                      num_heads=out_dim // 32,
+                      mlp_ratio=4.0,
+                      qkv_bias=False,
+                      drop_path=dpr[i]) for i in range(trans_layer)
+            ]
+            self.trans_blocks = nn.Sequential(*blocks)
+            dim = out_dim
+            self.out_channels = out_dim
+        else:
+            self.trans_blocks = None
+            dim = 512
+        self.last_stage = last_stage
+        if last_stage:
+            self.out_channels = out_channels
+            self.last_conv = nn.Linear(dim, self.out_channels, bias=False)
+            self.hardswish = nn.Hardswish()
+            self.dropout = nn.Dropout(p=0.1)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2.0 / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, mean=0, std=0.02)
+                if isinstance(m, nn.Linear) and m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+        x2 = self.layer2(x)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+        x5 = self.layer5(x4)
+
+        if self.return_list:
+            return [x2, x3, x4, x5]
+        x = x5
+        if self.trans_blocks is not None:
+            B, C, H, W = x.shape
+            x = self.trans_blocks(x.flatten(2, 3).transpose(1, 2))
+            x = x.transpose(1, 2).reshape(B, -1, H, W)
+
+        if self.last_stage:
+            x = x.mean(2).transpose(1, 2)
+            x = self.last_conv(x)
+            x = self.hardswish(x)
+            x = self.dropout(x)
+        elif not self.feat2d:
+            x = x.flatten(2).transpose(1, 2)
+        return x