Spaces:
Runtime error
Runtime error
File size: 4,405 Bytes
2366e36 |
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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmcv.runner import BaseModule, Sequential
import mmocr.utils as utils
from mmocr.models.builder import BACKBONES
from mmocr.models.textrecog.layers import BasicBlock
@BACKBONES.register_module()
class ResNetABI(BaseModule):
"""Implement ResNet backbone for text recognition, modified from `ResNet.
<https://arxiv.org/pdf/1512.03385.pdf>`_ and
`<https://github.com/FangShancheng/ABINet>`_
Args:
in_channels (int): Number of channels of input image tensor.
stem_channels (int): Number of stem channels.
base_channels (int): Number of base channels.
arch_settings (list[int]): List of BasicBlock number for each stage.
strides (Sequence[int]): Strides of the first block of each stage.
out_indices (None | Sequence[int]): Indices of output stages. If not
specified, only the last stage will be returned.
last_stage_pool (bool): If True, add `MaxPool2d` layer to last stage.
"""
def __init__(self,
in_channels=3,
stem_channels=32,
base_channels=32,
arch_settings=[3, 4, 6, 6, 3],
strides=[2, 1, 2, 1, 1],
out_indices=None,
last_stage_pool=False,
init_cfg=[
dict(type='Xavier', layer='Conv2d'),
dict(type='Constant', val=1, layer='BatchNorm2d')
]):
super().__init__(init_cfg=init_cfg)
assert isinstance(in_channels, int)
assert isinstance(stem_channels, int)
assert utils.is_type_list(arch_settings, int)
assert utils.is_type_list(strides, int)
assert len(arch_settings) == len(strides)
assert out_indices is None or isinstance(out_indices, (list, tuple))
assert isinstance(last_stage_pool, bool)
self.out_indices = out_indices
self.last_stage_pool = last_stage_pool
self.block = BasicBlock
self.inplanes = stem_channels
self._make_stem_layer(in_channels, stem_channels)
self.res_layers = []
planes = base_channels
for i, num_blocks in enumerate(arch_settings):
stride = strides[i]
res_layer = self._make_layer(
block=self.block,
inplanes=self.inplanes,
planes=planes,
blocks=num_blocks,
stride=stride)
self.inplanes = planes * self.block.expansion
planes *= 2
layer_name = f'layer{i + 1}'
self.add_module(layer_name, res_layer)
self.res_layers.append(layer_name)
def _make_layer(self, block, inplanes, planes, blocks, stride=1):
layers = []
downsample = None
if stride != 1 or inplanes != planes:
downsample = nn.Sequential(
nn.Conv2d(inplanes, planes, 1, stride, bias=False),
nn.BatchNorm2d(planes),
)
layers.append(
block(
inplanes,
planes,
use_conv1x1=True,
stride=stride,
downsample=downsample))
inplanes = planes
for _ in range(1, blocks):
layers.append(block(inplanes, planes, use_conv1x1=True))
return Sequential(*layers)
def _make_stem_layer(self, in_channels, stem_channels):
self.conv1 = nn.Conv2d(
in_channels, stem_channels, kernel_size=3, stride=1, padding=1)
self.bn1 = nn.BatchNorm2d(stem_channels)
self.relu1 = nn.ReLU(inplace=True)
def forward(self, x):
"""
Args:
x (Tensor): Image tensor of shape :math:`(N, 3, H, W)`.
Returns:
Tensor or list[Tensor]: Feature tensor. Its shape depends on
ResNetABI's config. It can be a list of feature outputs at specific
layers if ``out_indices`` is specified.
"""
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
outs = []
for i, layer_name in enumerate(self.res_layers):
res_layer = getattr(self, layer_name)
x = res_layer(x)
if self.out_indices and i in self.out_indices:
outs.append(x)
return tuple(outs) if self.out_indices else x
|