diff --git a/python/demo.py b/python/demo.py index c4cbda4..ab4fdd7 100644 --- a/python/demo.py +++ b/python/demo.py @@ -18,7 +18,7 @@ from rapidocr import RapidOCR, VisRes # engine = engine = RapidOCR( # params={"Global.with_onnx": True, "EngineConfig.onnxruntime.use_cuda": True} # ) -engine = RapidOCR(params={"Global.with_openvino": True, "Global.lang": "ch"}) +engine = RapidOCR(params={"Global.with_torch": True, "Global.lang": "ch"}) vis = VisRes() image_path = "tests/test_files/ch_en_num.jpg" diff --git a/python/rapidocr/inference_engine/torch.py b/python/rapidocr/inference_engine/torch.py index 500de16..bd5756f 100644 --- a/python/rapidocr/inference_engine/torch.py +++ b/python/rapidocr/inference_engine/torch.py @@ -2,44 +2,39 @@ # @Author: SWHL # @Contact: liekkaskono@163.com from pathlib import Path -from typing import Dict, Optional, Union +from typing import Optional import numpy as np import torch -import yaml +from omegaconf import OmegaConf + +from ..networks.architectures.base_model import BaseModel +from ..utils.logger import Logger +from .base import InferSession root_dir = Path(__file__).resolve().parent.parent -DEFAULT_CFG_PATH = root_dir / "arch_config.yaml" +DEFAULT_CFG_PATH = root_dir / "networks" / "arch_config.yaml" -def read_yaml(yaml_path: Union[str, Path]) -> Dict[str, Dict]: - with open(yaml_path, "rb") as f: - data = yaml.load(f, Loader=yaml.Loader) - return data - - -from rapidocr_torch.modeling.architectures.base_model import BaseModel - -from ..utils.logger import Logger - - -class TorchInferSession: +class TorchInferSession(InferSession): def __init__(self, config, mode: Optional[str] = None) -> None: self.logger = Logger(logger_name=__name__).get_log() - all_arch_config = read_yaml(DEFAULT_CFG_PATH) + all_arch_config = OmegaConf.load(DEFAULT_CFG_PATH) self.mode = mode model_path = Path(config["model_path"]) self._verify_model(model_path) + file_name = model_path.stem if file_name not in all_arch_config: - raise ValueError(f"architecture {file_name} is not in config.yaml") + raise ValueError(f"architecture {file_name} is not in arch_config.yaml") arch_config = all_arch_config[file_name] self.predictor = BaseModel(arch_config) self.predictor.load_state_dict(torch.load(model_path, weights_only=True)) self.predictor.eval() + self.use_gpu = False if config["use_cuda"]: self.predictor.cuda() @@ -50,17 +45,13 @@ class TorchInferSession: inp = torch.from_numpy(img) if self.use_gpu: inp = inp.cuda() + # 适配跟onnx对齐取值逻辑 outputs = self.predictor(inp).unsqueeze(0) return outputs.cpu().numpy() - @staticmethod - def _verify_model(model_path): - model_path = Path(model_path) - if not model_path.exists(): - raise FileNotFoundError(f"{model_path} does not exists.") - if not model_path.is_file(): - raise FileExistsError(f"{model_path} is not a file.") + def have_key(self, key: str = "character") -> bool: + return False class TorchInferError(Exception): diff --git a/python/rapidocr/networks/__init__.py b/python/rapidocr/networks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/rapidocr/networks/arch_config.yaml b/python/rapidocr/networks/arch_config.yaml new file mode 100644 index 0000000..0879157 --- /dev/null +++ b/python/rapidocr/networks/arch_config.yaml @@ -0,0 +1,98 @@ +ch_ptocr_mobile_v2.0_cls_infer: + model_type: cls + algorithm: CLS + Transform: + Backbone: + name: MobileNetV3 + scale: 0.35 + model_name: small + Neck: + Head: + name: ClsHead + class_dim: 2 + +ch_PP-OCRv4_det_infer: + model_type: det + algorithm: DB + Transform: null + Backbone: + name: PPLCNetV3 + scale: 0.75 + det: True + Neck: + name: RSEFPN + out_channels: 96 + shortcut: True + Head: + name: DBHead + k: 50 + + +ch_PP-OCRv4_det_server_infer: + model_type: det + algorithm: DB + Transform: null + Backbone: + name: PPHGNet_small + det: True + Neck: + name: LKPAN + out_channels: 256 + intracl: true + Head: + name: PFHeadLocal + k: 50 + mode: "large" + + +ch_PP-OCRv4_rec_infer: + model_type: rec + algorithm: SVTR_LCNet + Transform: + Backbone: + name: PPLCNetV3 + scale: 0.95 + Head: + name: MultiHead + out_channels_list: + CTCLabelDecode: 6625 #'blank' + ...(6623) + ' ' + head_list: + - CTCHead: + Neck: + name: svtr + dims: 120 + depth: 2 + hidden_dims: 120 + kernel_size: [ 1, 3 ] + use_guide: True + Head: + fc_decay: 0.00001 + - NRTRHead: + nrtr_dim: 384 + max_text_length: 25 + + +ch_PP-OCRv4_rec_server_infer: + model_type: rec + algorithm: SVTR_HGNet + Transform: + Backbone: + name: PPHGNet_small + Head: + name: MultiHead + out_channels_list: + CTCLabelDecode: 6625 #'blank' + ...(6623) + ' ' + head_list: + - CTCHead: + Neck: + name: svtr + dims: 120 + depth: 2 + hidden_dims: 120 + kernel_size: [ 1, 3 ] + use_guide: True + Head: + fc_decay: 0.00001 + - NRTRHead: + nrtr_dim: 384 + max_text_length: 25 \ No newline at end of file diff --git a/python/rapidocr/networks/architectures/__init__.py b/python/rapidocr/networks/architectures/__init__.py new file mode 100644 index 0000000..7ad5eb4 --- /dev/null +++ b/python/rapidocr/networks/architectures/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy + +__all__ = ["build_model"] + + +def build_model(config, **kwargs): + from .base_model import BaseModel + + config = copy.deepcopy(config) + module_class = BaseModel(config, **kwargs) + return module_class diff --git a/python/rapidocr/networks/architectures/base_model.py b/python/rapidocr/networks/architectures/base_model.py new file mode 100644 index 0000000..e7f7ce4 --- /dev/null +++ b/python/rapidocr/networks/architectures/base_model.py @@ -0,0 +1,105 @@ +from torch import nn + +from ..backbones import build_backbone +from ..heads import build_head +from ..necks import build_neck + + +class BaseModel(nn.Module): + def __init__(self, config, **kwargs): + """ + the module for OCR. + args: + config (dict): the super parameters for module. + """ + super(BaseModel, self).__init__() + + in_channels = config.get("in_channels", 3) + model_type = config["model_type"] + # build backbone, backbone is need for del, rec and cls + if "Backbone" not in config or config["Backbone"] is None: + self.use_backbone = False + else: + self.use_backbone = True + config["Backbone"]["in_channels"] = in_channels + self.backbone = build_backbone(config["Backbone"], model_type) + in_channels = self.backbone.out_channels + + # build neck + # for rec, neck can be cnn,rnn or reshape(None) + # for det, neck can be FPN, BIFPN and so on. + # for cls, neck should be none + if "Neck" not in config or config["Neck"] is None: + self.use_neck = False + else: + self.use_neck = True + config["Neck"]["in_channels"] = in_channels + self.neck = build_neck(config["Neck"]) + in_channels = self.neck.out_channels + + # # build head, head is need for det, rec and cls + if "Head" not in config or config["Head"] is None: + self.use_head = False + else: + self.use_head = True + config["Head"]["in_channels"] = in_channels + self.head = build_head(config["Head"], **kwargs) + + self.return_all_feats = config.get("return_all_feats", False) + + self._initialize_weights() + + def _initialize_weights(self): + # weight initialization + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x): + y = dict() + if self.use_backbone: + x = self.backbone(x) + if isinstance(x, dict): + y.update(x) + else: + y["backbone_out"] = x + final_name = "backbone_out" + if self.use_neck: + x = self.neck(x) + if isinstance(x, dict): + y.update(x) + else: + y["neck_out"] = x + final_name = "neck_out" + if self.use_head: + x = self.head(x) + # for multi head, save ctc neck out for udml + if isinstance(x, dict) and "ctc_nect" in x.keys(): + y["neck_out"] = x["ctc_neck"] + y["head_out"] = x + elif isinstance(x, dict): + y.update(x) + else: + y["head_out"] = x + if self.return_all_feats: + if self.training: + return y + elif isinstance(x, dict): + return x + else: + return {final_name: x} + else: + return x diff --git a/python/rapidocr/networks/backbones/__init__.py b/python/rapidocr/networks/backbones/__init__.py new file mode 100644 index 0000000..8e7c5b3 --- /dev/null +++ b/python/rapidocr/networks/backbones/__init__.py @@ -0,0 +1,61 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__all__ = ["build_backbone"] + + +def build_backbone(config, model_type): + if model_type == "det": + from .det_mobilenet_v3 import MobileNetV3 + from .rec_hgnet import PPHGNet_small + from .rec_lcnetv3 import PPLCNetV3 + + support_dict = [ + "MobileNetV3", + "ResNet", + "ResNet_vd", + "ResNet_SAST", + "PPLCNetV3", + "PPHGNet_small", + ] + elif model_type == "rec" or model_type == "cls": + from .rec_hgnet import PPHGNet_small + from .rec_lcnetv3 import PPLCNetV3 + from .rec_mobilenet_v3 import MobileNetV3 + from .rec_svtrnet import SVTRNet + + support_dict = [ + "MobileNetV1Enhance", + "MobileNetV3", + "ResNet", + "ResNetFPN", + "MTB", + "ResNet31", + "SVTRNet", + "ViTSTR", + "DenseNet", + "PPLCNetV3", + "PPHGNet_small", + ] + else: + raise NotImplementedError + + module_name = config.pop("name") + assert module_name in support_dict, Exception( + "when model typs is {}, backbone only support {}".format( + model_type, support_dict + ) + ) + module_class = eval(module_name)(**config) + return module_class diff --git a/python/rapidocr/networks/backbones/det_mobilenet_v3.py b/python/rapidocr/networks/backbones/det_mobilenet_v3.py new file mode 100644 index 0000000..0351159 --- /dev/null +++ b/python/rapidocr/networks/backbones/det_mobilenet_v3.py @@ -0,0 +1,269 @@ +from torch import nn + +from ..common import Activation + + +def make_divisible(v, divisor=8, min_value=None): + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + if new_v < 0.9 * v: + new_v += divisor + return new_v + + +class ConvBNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + groups=1, + if_act=True, + act=None, + name=None, + ): + super(ConvBNLayer, self).__init__() + self.if_act = if_act + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + bias=False, + ) + + self.bn = nn.BatchNorm2d( + out_channels, + ) + if self.if_act: + self.act = Activation(act_type=act, inplace=True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + if self.if_act: + x = self.act(x) + return x + + +class SEModule(nn.Module): + def __init__(self, in_channels, reduction=4, name=""): + super(SEModule, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.conv1 = nn.Conv2d( + in_channels=in_channels, + out_channels=in_channels // reduction, + kernel_size=1, + stride=1, + padding=0, + bias=True, + ) + self.relu1 = Activation(act_type="relu", inplace=True) + self.conv2 = nn.Conv2d( + in_channels=in_channels // reduction, + out_channels=in_channels, + kernel_size=1, + stride=1, + padding=0, + bias=True, + ) + self.hard_sigmoid = Activation(act_type="hard_sigmoid", inplace=True) + + def forward(self, inputs): + outputs = self.avg_pool(inputs) + outputs = self.conv1(outputs) + outputs = self.relu1(outputs) + outputs = self.conv2(outputs) + outputs = self.hard_sigmoid(outputs) + outputs = inputs * outputs + return outputs + + +class ResidualUnit(nn.Module): + def __init__( + self, + in_channels, + mid_channels, + out_channels, + kernel_size, + stride, + use_se, + act=None, + name="", + ): + super(ResidualUnit, self).__init__() + self.if_shortcut = stride == 1 and in_channels == out_channels + self.if_se = use_se + + self.expand_conv = ConvBNLayer( + in_channels=in_channels, + out_channels=mid_channels, + kernel_size=1, + stride=1, + padding=0, + if_act=True, + act=act, + name=name + "_expand", + ) + self.bottleneck_conv = ConvBNLayer( + in_channels=mid_channels, + out_channels=mid_channels, + kernel_size=kernel_size, + stride=stride, + padding=int((kernel_size - 1) // 2), + groups=mid_channels, + if_act=True, + act=act, + name=name + "_depthwise", + ) + if self.if_se: + self.mid_se = SEModule(mid_channels, name=name + "_se") + self.linear_conv = ConvBNLayer( + in_channels=mid_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + padding=0, + if_act=False, + act=None, + name=name + "_linear", + ) + + def forward(self, inputs): + x = self.expand_conv(inputs) + x = self.bottleneck_conv(x) + if self.if_se: + x = self.mid_se(x) + x = self.linear_conv(x) + if self.if_shortcut: + x = inputs + x + return x + + +class MobileNetV3(nn.Module): + def __init__( + self, in_channels=3, model_name="large", scale=0.5, disable_se=False, **kwargs + ): + """ + the MobilenetV3 backbone network for detection module. + Args: + params(dict): the super parameters for build network + """ + super(MobileNetV3, self).__init__() + + self.disable_se = disable_se + + if model_name == "large": + cfg = [ + # k, exp, c, se, nl, s, + [3, 16, 16, False, "relu", 1], + [3, 64, 24, False, "relu", 2], + [3, 72, 24, False, "relu", 1], + [5, 72, 40, True, "relu", 2], + [5, 120, 40, True, "relu", 1], + [5, 120, 40, True, "relu", 1], + [3, 240, 80, False, "hard_swish", 2], + [3, 200, 80, False, "hard_swish", 1], + [3, 184, 80, False, "hard_swish", 1], + [3, 184, 80, False, "hard_swish", 1], + [3, 480, 112, True, "hard_swish", 1], + [3, 672, 112, True, "hard_swish", 1], + [5, 672, 160, True, "hard_swish", 2], + [5, 960, 160, True, "hard_swish", 1], + [5, 960, 160, True, "hard_swish", 1], + ] + cls_ch_squeeze = 960 + elif model_name == "small": + cfg = [ + # k, exp, c, se, nl, s, + [3, 16, 16, True, "relu", 2], + [3, 72, 24, False, "relu", 2], + [3, 88, 24, False, "relu", 1], + [5, 96, 40, True, "hard_swish", 2], + [5, 240, 40, True, "hard_swish", 1], + [5, 240, 40, True, "hard_swish", 1], + [5, 120, 48, True, "hard_swish", 1], + [5, 144, 48, True, "hard_swish", 1], + [5, 288, 96, True, "hard_swish", 2], + [5, 576, 96, True, "hard_swish", 1], + [5, 576, 96, True, "hard_swish", 1], + ] + cls_ch_squeeze = 576 + else: + raise NotImplementedError( + "mode[" + model_name + "_model] is not implemented!" + ) + + supported_scale = [0.35, 0.5, 0.75, 1.0, 1.25] + assert ( + scale in supported_scale + ), "supported scale are {} but input scale is {}".format(supported_scale, scale) + inplanes = 16 + # conv1 + self.conv = ConvBNLayer( + in_channels=in_channels, + out_channels=make_divisible(inplanes * scale), + kernel_size=3, + stride=2, + padding=1, + groups=1, + if_act=True, + act="hard_swish", + name="conv1", + ) + + self.stages = nn.ModuleList() + self.out_channels = [] + block_list = [] + i = 0 + inplanes = make_divisible(inplanes * scale) + for k, exp, c, se, nl, s in cfg: + se = se and not self.disable_se + if s == 2 and i > 2: + self.out_channels.append(inplanes) + self.stages.append(nn.Sequential(*block_list)) + block_list = [] + block_list.append( + ResidualUnit( + in_channels=inplanes, + mid_channels=make_divisible(scale * exp), + out_channels=make_divisible(scale * c), + kernel_size=k, + stride=s, + use_se=se, + act=nl, + name="conv" + str(i + 2), + ) + ) + inplanes = make_divisible(scale * c) + i += 1 + block_list.append( + ConvBNLayer( + in_channels=inplanes, + out_channels=make_divisible(scale * cls_ch_squeeze), + kernel_size=1, + stride=1, + padding=0, + groups=1, + if_act=True, + act="hard_swish", + name="conv_last", + ) + ) + self.stages.append(nn.Sequential(*block_list)) + self.out_channels.append(make_divisible(scale * cls_ch_squeeze)) + # for i, stage in enumerate(self.stages): + # self.add_sublayer(sublayer=stage, name="stage{}".format(i)) + + def forward(self, x): + x = self.conv(x) + out_list = [] + for stage in self.stages: + x = stage(x) + out_list.append(x) + return out_list diff --git a/python/rapidocr/networks/backbones/rec_hgnet.py b/python/rapidocr/networks/backbones/rec_hgnet.py new file mode 100644 index 0000000..c1515a7 --- /dev/null +++ b/python/rapidocr/networks/backbones/rec_hgnet.py @@ -0,0 +1,290 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +class ConvBNAct(nn.Module): + def __init__( + self, in_channels, out_channels, kernel_size, stride, groups=1, use_act=True + ): + super().__init__() + self.use_act = use_act + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding=(kernel_size - 1) // 2, + groups=groups, + bias=False, + ) + self.bn = nn.BatchNorm2d(out_channels) + if self.use_act: + self.act = nn.ReLU() + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + if self.use_act: + x = self.act(x) + return x + + +class ESEModule(nn.Module): + def __init__(self, channels): + super().__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.conv = nn.Conv2d( + in_channels=channels, + out_channels=channels, + kernel_size=1, + stride=1, + padding=0, + ) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + identity = x + x = self.avg_pool(x) + x = self.conv(x) + x = self.sigmoid(x) + return x * identity + + +class HG_Block(nn.Module): + def __init__( + self, + in_channels, + mid_channels, + out_channels, + layer_num, + identity=False, + ): + super().__init__() + self.identity = identity + + self.layers = nn.ModuleList() + self.layers.append( + ConvBNAct( + in_channels=in_channels, + out_channels=mid_channels, + kernel_size=3, + stride=1, + ) + ) + for _ in range(layer_num - 1): + self.layers.append( + ConvBNAct( + in_channels=mid_channels, + out_channels=mid_channels, + kernel_size=3, + stride=1, + ) + ) + + # feature aggregation + total_channels = in_channels + layer_num * mid_channels + self.aggregation_conv = ConvBNAct( + in_channels=total_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + ) + self.att = ESEModule(out_channels) + + def forward(self, x): + identity = x + output = [] + output.append(x) + for layer in self.layers: + x = layer(x) + output.append(x) + x = torch.cat(output, dim=1) + x = self.aggregation_conv(x) + x = self.att(x) + if self.identity: + x += identity + return x + + +class HG_Stage(nn.Module): + def __init__( + self, + in_channels, + mid_channels, + out_channels, + block_num, + layer_num, + downsample=True, + stride=[2, 1], + ): + super().__init__() + self.downsample = downsample + if downsample: + self.downsample = ConvBNAct( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=3, + stride=stride, + groups=in_channels, + use_act=False, + ) + + blocks_list = [] + blocks_list.append( + HG_Block(in_channels, mid_channels, out_channels, layer_num, identity=False) + ) + for _ in range(block_num - 1): + blocks_list.append( + HG_Block( + out_channels, mid_channels, out_channels, layer_num, identity=True + ) + ) + self.blocks = nn.Sequential(*blocks_list) + + def forward(self, x): + if self.downsample: + x = self.downsample(x) + x = self.blocks(x) + return x + + +class PPHGNet(nn.Module): + """ + PPHGNet + Args: + stem_channels: list. Stem channel list of PPHGNet. + stage_config: dict. The configuration of each stage of PPHGNet. such as the number of channels, stride, etc. + layer_num: int. Number of layers of HG_Block. + use_last_conv: boolean. Whether to use a 1x1 convolutional layer before the classification layer. + class_expand: int=2048. Number of channels for the last 1x1 convolutional layer. + dropout_prob: float. Parameters of dropout, 0.0 means dropout is not used. + class_num: int=1000. The number of classes. + Returns: + model: nn.Layer. Specific PPHGNet model depends on args. + """ + + def __init__( + self, + stem_channels, + stage_config, + layer_num, + in_channels=3, + det=False, + out_indices=None, + ): + super().__init__() + self.det = det + self.out_indices = out_indices if out_indices is not None else [0, 1, 2, 3] + + # stem + stem_channels.insert(0, in_channels) + self.stem = nn.Sequential( + *[ + ConvBNAct( + in_channels=stem_channels[i], + out_channels=stem_channels[i + 1], + kernel_size=3, + stride=2 if i == 0 else 1, + ) + for i in range(len(stem_channels) - 1) + ] + ) + + if self.det: + self.pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + # stages + self.stages = nn.ModuleList() + self.out_channels = [] + for block_id, k in enumerate(stage_config): + ( + in_channels, + mid_channels, + out_channels, + block_num, + downsample, + stride, + ) = stage_config[k] + self.stages.append( + HG_Stage( + in_channels, + mid_channels, + out_channels, + block_num, + layer_num, + downsample, + stride, + ) + ) + if block_id in self.out_indices: + self.out_channels.append(out_channels) + + if not self.det: + self.out_channels = stage_config["stage4"][2] + + self._init_weights() + + def _init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.zeros_(m.bias) + + def forward(self, x): + x = self.stem(x) + if self.det: + x = self.pool(x) + + out = [] + for i, stage in enumerate(self.stages): + x = stage(x) + if self.det and i in self.out_indices: + out.append(x) + if self.det: + return out + + if self.training: + x = F.adaptive_avg_pool2d(x, [1, 40]) + else: + x = F.avg_pool2d(x, [3, 2]) + return x + + +def PPHGNet_small(pretrained=False, use_ssld=False, det=False, **kwargs): + """ + PPHGNet_small + Args: + pretrained: bool=False or str. If `True` load pretrained parameters, `False` otherwise. + If str, means the path of the pretrained model. + use_ssld: bool=False. Whether using distillation pretrained model when pretrained=True. + Returns: + model: nn.Layer. Specific `PPHGNet_small` model depends on args. + """ + stage_config_det = { + # in_channels, mid_channels, out_channels, blocks, downsample + "stage1": [128, 128, 256, 1, False, 2], + "stage2": [256, 160, 512, 1, True, 2], + "stage3": [512, 192, 768, 2, True, 2], + "stage4": [768, 224, 1024, 1, True, 2], + } + + stage_config_rec = { + # in_channels, mid_channels, out_channels, blocks, downsample + "stage1": [128, 128, 256, 1, True, [2, 1]], + "stage2": [256, 160, 512, 1, True, [1, 2]], + "stage3": [512, 192, 768, 2, True, [2, 1]], + "stage4": [768, 224, 1024, 1, True, [2, 1]], + } + + model = PPHGNet( + stem_channels=[64, 64, 128], + stage_config=stage_config_det if det else stage_config_rec, + layer_num=6, + det=det, + **kwargs + ) + return model diff --git a/python/rapidocr/networks/backbones/rec_lcnetv3.py b/python/rapidocr/networks/backbones/rec_lcnetv3.py new file mode 100644 index 0000000..e2bd457 --- /dev/null +++ b/python/rapidocr/networks/backbones/rec_lcnetv3.py @@ -0,0 +1,516 @@ +# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import, division, print_function + +import torch +import torch.nn.functional as F +from torch import nn + +from ..common import Activation + +NET_CONFIG_det = { + "blocks2": + # k, in_c, out_c, s, use_se + [[3, 16, 32, 1, False]], + "blocks3": [[3, 32, 64, 2, False], [3, 64, 64, 1, False]], + "blocks4": [[3, 64, 128, 2, False], [3, 128, 128, 1, False]], + "blocks5": [ + [3, 128, 256, 2, False], + [5, 256, 256, 1, False], + [5, 256, 256, 1, False], + [5, 256, 256, 1, False], + [5, 256, 256, 1, False], + ], + "blocks6": [ + [5, 256, 512, 2, True], + [5, 512, 512, 1, True], + [5, 512, 512, 1, False], + [5, 512, 512, 1, False], + ], +} + +NET_CONFIG_rec = { + "blocks2": + # k, in_c, out_c, s, use_se + [[3, 16, 32, 1, False]], + "blocks3": [[3, 32, 64, 1, False], [3, 64, 64, 1, False]], + "blocks4": [[3, 64, 128, (2, 1), False], [3, 128, 128, 1, False]], + "blocks5": [ + [3, 128, 256, (1, 2), False], + [5, 256, 256, 1, False], + [5, 256, 256, 1, False], + [5, 256, 256, 1, False], + [5, 256, 256, 1, False], + ], + "blocks6": [ + [5, 256, 512, (2, 1), True], + [5, 512, 512, 1, True], + [5, 512, 512, (2, 1), False], + [5, 512, 512, 1, False], + ], +} + + +def make_divisible(v, divisor=16, min_value=None): + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + if new_v < 0.9 * v: + new_v += divisor + return new_v + + +class LearnableAffineBlock(nn.Module): + def __init__(self, scale_value=1.0, bias_value=0.0, lr_mult=1.0, lab_lr=0.1): + super().__init__() + self.scale = nn.Parameter(torch.Tensor([scale_value])) + self.bias = nn.Parameter(torch.Tensor([bias_value])) + + def forward(self, x): + return self.scale * x + self.bias + + +class ConvBNLayer(nn.Module): + def __init__( + self, in_channels, out_channels, kernel_size, stride, groups=1, lr_mult=1.0 + ): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=(kernel_size - 1) // 2, + groups=groups, + bias=False, + ) + + self.bn = nn.BatchNorm2d( + out_channels, + ) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + return x + + +class Act(nn.Module): + def __init__(self, act="hswish", lr_mult=1.0, lab_lr=0.1): + super().__init__() + if act == "hswish": + self.act = nn.Hardswish(inplace=True) + else: + assert act == "relu" + self.act = Activation(act) + self.lab = LearnableAffineBlock(lr_mult=lr_mult, lab_lr=lab_lr) + + def forward(self, x): + return self.lab(self.act(x)) + + +class LearnableRepLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + groups=1, + num_conv_branches=1, + lr_mult=1.0, + lab_lr=0.1, + ): + super().__init__() + self.is_repped = False + self.groups = groups + self.stride = stride + self.kernel_size = kernel_size + self.in_channels = in_channels + self.out_channels = out_channels + self.num_conv_branches = num_conv_branches + self.padding = (kernel_size - 1) // 2 + + self.identity = ( + nn.BatchNorm2d( + num_features=in_channels, + ) + if out_channels == in_channels and stride == 1 + else None + ) + + self.conv_kxk = nn.ModuleList( + [ + ConvBNLayer( + in_channels, + out_channels, + kernel_size, + stride, + groups=groups, + lr_mult=lr_mult, + ) + for _ in range(self.num_conv_branches) + ] + ) + + self.conv_1x1 = ( + ConvBNLayer( + in_channels, out_channels, 1, stride, groups=groups, lr_mult=lr_mult + ) + if kernel_size > 1 + else None + ) + + self.lab = LearnableAffineBlock(lr_mult=lr_mult, lab_lr=lab_lr) + self.act = Act(lr_mult=lr_mult, lab_lr=lab_lr) + + def forward(self, x): + # for export + if self.is_repped: + out = self.lab(self.reparam_conv(x)) + if self.stride != 2: + out = self.act(out) + return out + + out = 0 + if self.identity is not None: + out += self.identity(x) + + if self.conv_1x1 is not None: + out += self.conv_1x1(x) + + for conv in self.conv_kxk: + out += conv(x) + + out = self.lab(out) + if self.stride != 2: + out = self.act(out) + return out + + def rep(self): + if self.is_repped: + return + kernel, bias = self._get_kernel_bias() + self.reparam_conv = nn.Conv2d( + in_channels=self.in_channels, + out_channels=self.out_channels, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + groups=self.groups, + ) + self.reparam_conv.weight.data = kernel + self.reparam_conv.bias.data = bias + self.is_repped = True + + def _pad_kernel_1x1_to_kxk(self, kernel1x1, pad): + if not isinstance(kernel1x1, torch.Tensor): + return 0 + else: + return nn.functional.pad(kernel1x1, [pad, pad, pad, pad]) + + def _get_kernel_bias(self): + kernel_conv_1x1, bias_conv_1x1 = self._fuse_bn_tensor(self.conv_1x1) + kernel_conv_1x1 = self._pad_kernel_1x1_to_kxk( + kernel_conv_1x1, self.kernel_size // 2 + ) + + kernel_identity, bias_identity = self._fuse_bn_tensor(self.identity) + + kernel_conv_kxk = 0 + bias_conv_kxk = 0 + for conv in self.conv_kxk: + kernel, bias = self._fuse_bn_tensor(conv) + kernel_conv_kxk += kernel + bias_conv_kxk += bias + + kernel_reparam = kernel_conv_kxk + kernel_conv_1x1 + kernel_identity + bias_reparam = bias_conv_kxk + bias_conv_1x1 + bias_identity + return kernel_reparam, bias_reparam + + def _fuse_bn_tensor(self, branch): + if not branch: + return 0, 0 + elif isinstance(branch, ConvBNLayer): + kernel = branch.conv.weight + running_mean = branch.bn._mean + running_var = branch.bn._variance + gamma = branch.bn.weight + beta = branch.bn.bias + eps = branch.bn._epsilon + else: + assert isinstance(branch, nn.BatchNorm2d) + if not hasattr(self, "id_tensor"): + input_dim = self.in_channels // self.groups + kernel_value = torch.zeros( + (self.in_channels, input_dim, self.kernel_size, self.kernel_size), + dtype=branch.weight.dtype, + ) + for i in range(self.in_channels): + kernel_value[ + i, i % input_dim, self.kernel_size // 2, self.kernel_size // 2 + ] = 1 + self.id_tensor = kernel_value + kernel = self.id_tensor + running_mean = branch._mean + running_var = branch._variance + gamma = branch.weight + beta = branch.bias + eps = branch._epsilon + std = (running_var + eps).sqrt() + t = (gamma / std).reshape((-1, 1, 1, 1)) + return kernel * t, beta - running_mean * gamma / std + + +class SELayer(nn.Module): + def __init__(self, channel, reduction=4, lr_mult=1.0): + super().__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, + ) + self.relu = nn.ReLU() + self.conv2 = nn.Conv2d( + in_channels=channel // reduction, + out_channels=channel, + kernel_size=1, + stride=1, + padding=0, + ) + self.hardsigmoid = nn.Hardsigmoid(inplace=True) + + def forward(self, x): + identity = x + x = self.avg_pool(x) + x = self.conv1(x) + x = self.relu(x) + x = self.conv2(x) + x = self.hardsigmoid(x) + x = identity * x + return x + + +class LCNetV3Block(nn.Module): + def __init__( + self, + in_channels, + out_channels, + stride, + dw_size, + use_se=False, + conv_kxk_num=4, + lr_mult=1.0, + lab_lr=0.1, + ): + super().__init__() + self.use_se = use_se + self.dw_conv = LearnableRepLayer( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=dw_size, + stride=stride, + groups=in_channels, + num_conv_branches=conv_kxk_num, + lr_mult=lr_mult, + lab_lr=lab_lr, + ) + if use_se: + self.se = SELayer(in_channels, lr_mult=lr_mult) + self.pw_conv = LearnableRepLayer( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + num_conv_branches=conv_kxk_num, + lr_mult=lr_mult, + lab_lr=lab_lr, + ) + + def forward(self, x): + x = self.dw_conv(x) + if self.use_se: + x = self.se(x) + x = self.pw_conv(x) + return x + + +class PPLCNetV3(nn.Module): + def __init__( + self, + scale=1.0, + conv_kxk_num=4, + lr_mult_list=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + lab_lr=0.1, + det=False, + **kwargs + ): + super().__init__() + self.scale = scale + self.lr_mult_list = lr_mult_list + self.det = det + + self.net_config = NET_CONFIG_det if self.det else NET_CONFIG_rec + + assert isinstance( + self.lr_mult_list, (list, tuple) + ), "lr_mult_list should be in (list, tuple) but got {}".format( + type(self.lr_mult_list) + ) + assert ( + len(self.lr_mult_list) == 6 + ), "lr_mult_list length should be 6 but got {}".format(len(self.lr_mult_list)) + + self.conv1 = ConvBNLayer( + in_channels=3, + out_channels=make_divisible(16 * scale), + kernel_size=3, + stride=2, + lr_mult=self.lr_mult_list[0], + ) + + self.blocks2 = nn.Sequential( + *[ + LCNetV3Block( + in_channels=make_divisible(in_c * scale), + out_channels=make_divisible(out_c * scale), + dw_size=k, + stride=s, + use_se=se, + conv_kxk_num=conv_kxk_num, + lr_mult=self.lr_mult_list[1], + lab_lr=lab_lr, + ) + for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks2"]) + ] + ) + + self.blocks3 = nn.Sequential( + *[ + LCNetV3Block( + in_channels=make_divisible(in_c * scale), + out_channels=make_divisible(out_c * scale), + dw_size=k, + stride=s, + use_se=se, + conv_kxk_num=conv_kxk_num, + lr_mult=self.lr_mult_list[2], + lab_lr=lab_lr, + ) + for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks3"]) + ] + ) + + self.blocks4 = nn.Sequential( + *[ + LCNetV3Block( + in_channels=make_divisible(in_c * scale), + out_channels=make_divisible(out_c * scale), + dw_size=k, + stride=s, + use_se=se, + conv_kxk_num=conv_kxk_num, + lr_mult=self.lr_mult_list[3], + lab_lr=lab_lr, + ) + for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks4"]) + ] + ) + + self.blocks5 = nn.Sequential( + *[ + LCNetV3Block( + in_channels=make_divisible(in_c * scale), + out_channels=make_divisible(out_c * scale), + dw_size=k, + stride=s, + use_se=se, + conv_kxk_num=conv_kxk_num, + lr_mult=self.lr_mult_list[4], + lab_lr=lab_lr, + ) + for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks5"]) + ] + ) + + self.blocks6 = nn.Sequential( + *[ + LCNetV3Block( + in_channels=make_divisible(in_c * scale), + out_channels=make_divisible(out_c * scale), + dw_size=k, + stride=s, + use_se=se, + conv_kxk_num=conv_kxk_num, + lr_mult=self.lr_mult_list[5], + lab_lr=lab_lr, + ) + for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks6"]) + ] + ) + self.out_channels = make_divisible(512 * scale) + + if self.det: + mv_c = [16, 24, 56, 480] + self.out_channels = [ + make_divisible(self.net_config["blocks3"][-1][2] * scale), + make_divisible(self.net_config["blocks4"][-1][2] * scale), + make_divisible(self.net_config["blocks5"][-1][2] * scale), + make_divisible(self.net_config["blocks6"][-1][2] * scale), + ] + + self.layer_list = nn.ModuleList( + [ + nn.Conv2d(self.out_channels[0], int(mv_c[0] * scale), 1, 1, 0), + nn.Conv2d(self.out_channels[1], int(mv_c[1] * scale), 1, 1, 0), + nn.Conv2d(self.out_channels[2], int(mv_c[2] * scale), 1, 1, 0), + nn.Conv2d(self.out_channels[3], int(mv_c[3] * scale), 1, 1, 0), + ] + ) + self.out_channels = [ + int(mv_c[0] * scale), + int(mv_c[1] * scale), + int(mv_c[2] * scale), + int(mv_c[3] * scale), + ] + + def forward(self, x): + out_list = [] + x = self.conv1(x) + x = self.blocks2(x) + x = self.blocks3(x) + out_list.append(x) + x = self.blocks4(x) + out_list.append(x) + x = self.blocks5(x) + out_list.append(x) + x = self.blocks6(x) + out_list.append(x) + + if self.det: + out_list[0] = self.layer_list[0](out_list[0]) + out_list[1] = self.layer_list[1](out_list[1]) + out_list[2] = self.layer_list[2](out_list[2]) + out_list[3] = self.layer_list[3](out_list[3]) + return out_list + + if self.training: + x = F.adaptive_avg_pool2d(x, [1, 40]) + else: + x = F.avg_pool2d(x, [3, 2]) + return x diff --git a/python/rapidocr/networks/backbones/rec_mobilenet_v3.py b/python/rapidocr/networks/backbones/rec_mobilenet_v3.py new file mode 100644 index 0000000..d284a6d --- /dev/null +++ b/python/rapidocr/networks/backbones/rec_mobilenet_v3.py @@ -0,0 +1,136 @@ +from torch import nn + +from .det_mobilenet_v3 import ConvBNLayer, ResidualUnit, make_divisible + + +class MobileNetV3(nn.Module): + def __init__( + self, + in_channels=3, + model_name="small", + scale=0.5, + large_stride=None, + small_stride=None, + **kwargs + ): + super(MobileNetV3, self).__init__() + if small_stride is None: + small_stride = [2, 2, 2, 2] + if large_stride is None: + large_stride = [1, 2, 2, 2] + + assert isinstance( + large_stride, list + ), "large_stride type must " "be list but got {}".format(type(large_stride)) + assert isinstance( + small_stride, list + ), "small_stride type must " "be list but got {}".format(type(small_stride)) + assert ( + len(large_stride) == 4 + ), "large_stride length must be " "4 but got {}".format(len(large_stride)) + assert ( + len(small_stride) == 4 + ), "small_stride length must be " "4 but got {}".format(len(small_stride)) + + if model_name == "large": + cfg = [ + # k, exp, c, se, nl, s, + [3, 16, 16, False, "relu", large_stride[0]], + [3, 64, 24, False, "relu", (large_stride[1], 1)], + [3, 72, 24, False, "relu", 1], + [5, 72, 40, True, "relu", (large_stride[2], 1)], + [5, 120, 40, True, "relu", 1], + [5, 120, 40, True, "relu", 1], + [3, 240, 80, False, "hard_swish", 1], + [3, 200, 80, False, "hard_swish", 1], + [3, 184, 80, False, "hard_swish", 1], + [3, 184, 80, False, "hard_swish", 1], + [3, 480, 112, True, "hard_swish", 1], + [3, 672, 112, True, "hard_swish", 1], + [5, 672, 160, True, "hard_swish", (large_stride[3], 1)], + [5, 960, 160, True, "hard_swish", 1], + [5, 960, 160, True, "hard_swish", 1], + ] + cls_ch_squeeze = 960 + elif model_name == "small": + cfg = [ + # k, exp, c, se, nl, s, + [3, 16, 16, True, "relu", (small_stride[0], 1)], + [3, 72, 24, False, "relu", (small_stride[1], 1)], + [3, 88, 24, False, "relu", 1], + [5, 96, 40, True, "hard_swish", (small_stride[2], 1)], + [5, 240, 40, True, "hard_swish", 1], + [5, 240, 40, True, "hard_swish", 1], + [5, 120, 48, True, "hard_swish", 1], + [5, 144, 48, True, "hard_swish", 1], + [5, 288, 96, True, "hard_swish", (small_stride[3], 1)], + [5, 576, 96, True, "hard_swish", 1], + [5, 576, 96, True, "hard_swish", 1], + ] + cls_ch_squeeze = 576 + else: + raise NotImplementedError( + "mode[" + model_name + "_model] is not implemented!" + ) + + supported_scale = [0.35, 0.5, 0.75, 1.0, 1.25] + assert ( + scale in supported_scale + ), "supported scales are {} but input scale is {}".format( + supported_scale, scale + ) + + inplanes = 16 + # conv1 + self.conv1 = ConvBNLayer( + in_channels=in_channels, + out_channels=make_divisible(inplanes * scale), + kernel_size=3, + stride=2, + padding=1, + groups=1, + if_act=True, + act="hard_swish", + name="conv1", + ) + i = 0 + block_list = [] + inplanes = make_divisible(inplanes * scale) + for k, exp, c, se, nl, s in cfg: + block_list.append( + ResidualUnit( + in_channels=inplanes, + mid_channels=make_divisible(scale * exp), + out_channels=make_divisible(scale * c), + kernel_size=k, + stride=s, + use_se=se, + act=nl, + name="conv" + str(i + 2), + ) + ) + inplanes = make_divisible(scale * c) + i += 1 + self.blocks = nn.Sequential(*block_list) + + self.conv2 = ConvBNLayer( + in_channels=inplanes, + out_channels=make_divisible(scale * cls_ch_squeeze), + kernel_size=1, + stride=1, + padding=0, + groups=1, + if_act=True, + act="hard_swish", + name="conv_last", + ) + + self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0) + self.out_channels = make_divisible(scale * cls_ch_squeeze) + + def forward(self, x): + x = self.conv1(x) + x = self.blocks(x) + x = self.conv2(x) + x = self.pool(x) + return x diff --git a/python/rapidocr/networks/backbones/rec_svtrnet.py b/python/rapidocr/networks/backbones/rec_svtrnet.py new file mode 100644 index 0000000..3a11773 --- /dev/null +++ b/python/rapidocr/networks/backbones/rec_svtrnet.py @@ -0,0 +1,638 @@ +import numpy as np +import torch +from torch import nn + +from ..common import Activation + + +def drop_path(x, drop_prob=0.0, training=False): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... + """ + if drop_prob == 0.0 or not training: + return x + keep_prob = torch.as_tensor(1 - drop_prob) + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype) + random_tensor = torch.floor(random_tensor) # binarize + output = x.divide(keep_prob) * random_tensor + return output + + +class ConvBNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=0, + bias_attr=False, + groups=1, + act="gelu", + ): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + bias=bias_attr, + ) + self.norm = nn.BatchNorm2d(out_channels) + self.act = Activation(act_type=act, inplace=True) + + def forward(self, inputs): + out = self.conv(inputs) + out = self.norm(out) + out = self.act(out) + return out + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + +class Identity(nn.Module): + def __init__(self): + super(Identity, self).__init__() + + def forward(self, input): + return input + + +class Mlp(nn.Module): + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer="gelu", + drop=0.0, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = Activation(act_type=act_layer, inplace=True) + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class ConvMixer(nn.Module): + def __init__( + self, + dim, + num_heads=8, + HW=[8, 25], + local_k=[3, 3], + ): + super().__init__() + self.HW = HW + self.dim = dim + self.local_mixer = nn.Conv2d( + dim, + dim, + local_k, + 1, + [local_k[0] // 2, local_k[1] // 2], + groups=num_heads, + ) + + def forward(self, x): + h = self.HW[0] + w = self.HW[1] + x = x.transpose([0, 2, 1]).reshape([0, self.dim, h, w]) + x = self.local_mixer(x) + x = x.flatten(2).permute(0, 2, 1) + return x + + +class Attention(nn.Module): + def __init__( + self, + dim, + num_heads=8, + mixer="Global", + HW=[8, 25], + local_k=[7, 11], + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + self.HW = HW + if HW is not None: + H = HW[0] + W = HW[1] + self.N = H * W + self.C = dim + if mixer == "Local" and HW is not None: + hk = local_k[0] + wk = local_k[1] + mask = torch.ones(H * W, H + hk - 1, W + wk - 1, dtype=torch.float32) + for h in range(0, H): + for w in range(0, W): + mask[h * W + w, h : h + hk, w : w + wk] = 0.0 + mask_paddle = mask[:, hk // 2 : H + hk // 2, wk // 2 : W + wk // 2].flatten( + 1 + ) + mask_inf = torch.full( + [H * W, H * W], fill_value=float("-Inf"), dtype=torch.float32 + ) + mask = torch.where(mask_paddle < 1, mask_paddle, mask_inf) + self.mask = mask.unsqueeze(0).unsqueeze(1) + # self.mask = mask[None, None, :] + self.mixer = mixer + + def forward(self, x): + if self.HW is not None: + N = self.N + C = self.C + else: + _, N, C = x.shape + qkv = self.qkv(x) + qkv = qkv.reshape((-1, N, 3, self.num_heads, C // self.num_heads)).permute( + 2, 0, 3, 1, 4 + ) + q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] + + attn = q.matmul(k.permute(0, 1, 3, 2)) + if self.mixer == "Local": + attn += self.mask + attn = nn.functional.softmax(attn, dim=-1) + attn = self.attn_drop(attn) + + x = (attn.matmul(v)).permute(0, 2, 1, 3).reshape((-1, N, C)) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class Block(nn.Module): + def __init__( + self, + dim, + num_heads, + mixer="Global", + local_mixer=[7, 11], + HW=None, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer="gelu", + norm_layer="nn.LayerNorm", + epsilon=1e-6, + prenorm=True, + ): + super().__init__() + if isinstance(norm_layer, str): + self.norm1 = eval(norm_layer)(dim, eps=epsilon) + else: + self.norm1 = norm_layer(dim) + if mixer == "Global" or mixer == "Local": + self.mixer = Attention( + dim, + num_heads=num_heads, + mixer=mixer, + HW=HW, + local_k=local_mixer, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + elif mixer == "Conv": + self.mixer = ConvMixer(dim, num_heads=num_heads, HW=HW, local_k=local_mixer) + else: + raise TypeError("The mixer must be one of [Global, Local, Conv]") + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity() + if isinstance(norm_layer, str): + self.norm2 = eval(norm_layer)(dim, eps=epsilon) + else: + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp_ratio = mlp_ratio + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + ) + self.prenorm = prenorm + + def forward(self, x): + if self.prenorm: + x = self.norm1(x + self.drop_path(self.mixer(x))) + x = self.norm2(x + self.drop_path(self.mlp(x))) + else: + x = x + self.drop_path(self.mixer(self.norm1(x))) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding""" + + def __init__( + self, + img_size=[32, 100], + in_channels=3, + embed_dim=768, + sub_num=2, + patch_size=[4, 4], + mode="pope", + ): + super().__init__() + num_patches = (img_size[1] // (2**sub_num)) * (img_size[0] // (2**sub_num)) + self.img_size = img_size + self.num_patches = num_patches + self.embed_dim = embed_dim + self.norm = None + if mode == "pope": + if sub_num == 2: + self.proj = nn.Sequential( + ConvBNLayer( + in_channels=in_channels, + out_channels=embed_dim // 2, + kernel_size=3, + stride=2, + padding=1, + act="gelu", + bias_attr=True, + ), + ConvBNLayer( + in_channels=embed_dim // 2, + out_channels=embed_dim, + kernel_size=3, + stride=2, + padding=1, + act="gelu", + bias_attr=True, + ), + ) + if sub_num == 3: + self.proj = nn.Sequential( + ConvBNLayer( + in_channels=in_channels, + out_channels=embed_dim // 4, + kernel_size=3, + stride=2, + padding=1, + act="gelu", + bias_attr=True, + ), + ConvBNLayer( + in_channels=embed_dim // 4, + out_channels=embed_dim // 2, + kernel_size=3, + stride=2, + padding=1, + act="gelu", + bias_attr=True, + ), + ConvBNLayer( + in_channels=embed_dim // 2, + out_channels=embed_dim, + kernel_size=3, + stride=2, + padding=1, + act="gelu", + bias_attr=True, + ), + ) + elif mode == "linear": + self.proj = nn.Conv2d( + 1, embed_dim, kernel_size=patch_size, stride=patch_size + ) + self.num_patches = ( + img_size[0] // patch_size[0] * img_size[1] // patch_size[1] + ) + + def forward(self, x): + B, C, H, W = x.shape + assert ( + H == self.img_size[0] and W == self.img_size[1] + ), "Input image size ({}*{}) doesn't match model ({}*{}).".format( + H, W, self.img_size[0], self.img_size[1] + ) + x = self.proj(x).flatten(2).permute(0, 2, 1) + return x + + +class SubSample(nn.Module): + def __init__( + self, + in_channels, + out_channels, + types="Pool", + stride=[2, 1], + sub_norm="nn.LayerNorm", + act=None, + ): + super().__init__() + self.types = types + if types == "Pool": + self.avgpool = nn.AvgPool2d( + kernel_size=[3, 5], stride=stride, padding=[1, 2] + ) + self.maxpool = nn.MaxPool2d( + kernel_size=[3, 5], stride=stride, padding=[1, 2] + ) + self.proj = nn.Linear(in_channels, out_channels) + else: + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size=3, + stride=stride, + padding=1, + ) + self.norm = eval(sub_norm)(out_channels) + if act is not None: + self.act = act() + else: + self.act = None + + def forward(self, x): + if self.types == "Pool": + x1 = self.avgpool(x) + x2 = self.maxpool(x) + x = (x1 + x2) * 0.5 + out = self.proj(x.flatten(2).permute(0, 2, 1)) + else: + x = self.conv(x) + out = x.flatten(2).permute(0, 2, 1) + out = self.norm(out) + if self.act is not None: + out = self.act(out) + + return out + + +class SVTRNet(nn.Module): + def __init__( + self, + img_size=[32, 100], + in_channels=3, + embed_dim=[64, 128, 256], + depth=[3, 6, 3], + num_heads=[2, 4, 8], + mixer=["Local"] * 6 + ["Global"] * 6, # Local atten, Global atten, Conv + local_mixer=[[7, 11], [7, 11], [7, 11]], + patch_merging="Conv", # Conv, Pool, None + mlp_ratio=4, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + last_drop=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + norm_layer="nn.LayerNorm", + sub_norm="nn.LayerNorm", + epsilon=1e-6, + out_channels=192, + out_char_num=25, + block_unit="Block", + act="gelu", + last_stage=True, + sub_num=2, + prenorm=True, + use_lenhead=False, + **kwargs + ): + super().__init__() + self.img_size = img_size + self.embed_dim = embed_dim + self.out_channels = out_channels + self.prenorm = prenorm + patch_merging = ( + None + if patch_merging != "Conv" and patch_merging != "Pool" + else patch_merging + ) + self.patch_embed = PatchEmbed( + img_size=img_size, + in_channels=in_channels, + embed_dim=embed_dim[0], + sub_num=sub_num, + ) + num_patches = self.patch_embed.num_patches + self.HW = [img_size[0] // (2**sub_num), img_size[1] // (2**sub_num)] + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim[0])) + self.pos_drop = nn.Dropout(p=drop_rate) + Block_unit = eval(block_unit) + + dpr = np.linspace(0, drop_path_rate, sum(depth)) + self.blocks1 = nn.ModuleList( + [ + Block_unit( + dim=embed_dim[0], + num_heads=num_heads[0], + mixer=mixer[0 : depth[0]][i], + HW=self.HW, + local_mixer=local_mixer[0], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer=act, + attn_drop=attn_drop_rate, + drop_path=dpr[0 : depth[0]][i], + norm_layer=norm_layer, + epsilon=epsilon, + prenorm=prenorm, + ) + for i in range(depth[0]) + ] + ) + if patch_merging is not None: + self.sub_sample1 = SubSample( + embed_dim[0], + embed_dim[1], + sub_norm=sub_norm, + stride=[2, 1], + types=patch_merging, + ) + HW = [self.HW[0] // 2, self.HW[1]] + else: + HW = self.HW + self.patch_merging = patch_merging + self.blocks2 = nn.ModuleList( + [ + Block_unit( + dim=embed_dim[1], + num_heads=num_heads[1], + mixer=mixer[depth[0] : depth[0] + depth[1]][i], + HW=HW, + local_mixer=local_mixer[1], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer=act, + attn_drop=attn_drop_rate, + drop_path=dpr[depth[0] : depth[0] + depth[1]][i], + norm_layer=norm_layer, + epsilon=epsilon, + prenorm=prenorm, + ) + for i in range(depth[1]) + ] + ) + if patch_merging is not None: + self.sub_sample2 = SubSample( + embed_dim[1], + embed_dim[2], + sub_norm=sub_norm, + stride=[2, 1], + types=patch_merging, + ) + HW = [self.HW[0] // 4, self.HW[1]] + else: + HW = self.HW + self.blocks3 = nn.ModuleList( + [ + Block_unit( + dim=embed_dim[2], + num_heads=num_heads[2], + mixer=mixer[depth[0] + depth[1] :][i], + HW=HW, + local_mixer=local_mixer[2], + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer=act, + attn_drop=attn_drop_rate, + drop_path=dpr[depth[0] + depth[1] :][i], + norm_layer=norm_layer, + epsilon=epsilon, + prenorm=prenorm, + ) + for i in range(depth[2]) + ] + ) + self.last_stage = last_stage + if last_stage: + self.avg_pool = nn.AdaptiveAvgPool2d([1, out_char_num]) + self.last_conv = nn.Conv2d( + in_channels=embed_dim[2], + out_channels=self.out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=False, + ) + self.hardswish = Activation("hard_swish", inplace=True) # nn.Hardswish() + # self.dropout = nn.Dropout(p=last_drop, mode="downscale_in_infer") + self.dropout = nn.Dropout(p=last_drop) + if not prenorm: + self.norm = eval(norm_layer)(embed_dim[-1], eps=epsilon) + self.use_lenhead = use_lenhead + if use_lenhead: + self.len_conv = nn.Linear(embed_dim[2], self.out_channels) + self.hardswish_len = Activation( + "hard_swish", inplace=True + ) # nn.Hardswish() + self.dropout_len = nn.Dropout(p=last_drop) + + torch.nn.init.xavier_normal_(self.pos_embed) + self.apply(self._init_weights) + + def _init_weights(self, m): + # weight initialization + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + + def forward_features(self, x): + x = self.patch_embed(x) + x = x + self.pos_embed + x = self.pos_drop(x) + for blk in self.blocks1: + x = blk(x) + if self.patch_merging is not None: + x = self.sub_sample1( + x.permute(0, 2, 1).reshape( + [-1, self.embed_dim[0], self.HW[0], self.HW[1]] + ) + ) + for blk in self.blocks2: + x = blk(x) + if self.patch_merging is not None: + x = self.sub_sample2( + x.permute(0, 2, 1).reshape( + [-1, self.embed_dim[1], self.HW[0] // 2, self.HW[1]] + ) + ) + for blk in self.blocks3: + x = blk(x) + if not self.prenorm: + x = self.norm(x) + return x + + def forward(self, x): + x = self.forward_features(x) + if self.use_lenhead: + len_x = self.len_conv(x.mean(1)) + len_x = self.dropout_len(self.hardswish_len(len_x)) + if self.last_stage: + if self.patch_merging is not None: + h = self.HW[0] // 4 + else: + h = self.HW[0] + x = self.avg_pool( + x.permute(0, 2, 1).reshape([-1, self.embed_dim[2], h, self.HW[1]]) + ) + x = self.last_conv(x) + x = self.hardswish(x) + x = self.dropout(x) + if self.use_lenhead: + return x, len_x + return x diff --git a/python/rapidocr/networks/common.py b/python/rapidocr/networks/common.py new file mode 100644 index 0000000..ec1b30c --- /dev/null +++ b/python/rapidocr/networks/common.py @@ -0,0 +1,76 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +class Hswish(nn.Module): + def __init__(self, inplace=True): + super(Hswish, self).__init__() + self.inplace = inplace + + def forward(self, x): + return x * F.relu6(x + 3.0, inplace=self.inplace) / 6.0 + + +# out = max(0, min(1, slop*x+offset)) +# paddle.fluid.layers.hard_sigmoid(x, slope=0.2, offset=0.5, name=None) +class Hsigmoid(nn.Module): + def __init__(self, inplace=True): + super(Hsigmoid, self).__init__() + self.inplace = inplace + + def forward(self, x): + # torch: F.relu6(x + 3., inplace=self.inplace) / 6. + # paddle: F.relu6(1.2 * x + 3., inplace=self.inplace) / 6. + return F.relu6(1.2 * x + 3.0, inplace=self.inplace) / 6.0 + + +class GELU(nn.Module): + def __init__(self, inplace=True): + super(GELU, self).__init__() + self.inplace = inplace + + def forward(self, x): + return torch.nn.functional.gelu(x) + + +class Swish(nn.Module): + def __init__(self, inplace=True): + super(Swish, self).__init__() + self.inplace = inplace + + def forward(self, x): + if self.inplace: + x.mul_(torch.sigmoid(x)) + return x + else: + return x * torch.sigmoid(x) + + +class Activation(nn.Module): + def __init__(self, act_type, inplace=True): + super(Activation, self).__init__() + act_type = act_type.lower() + if act_type == "relu": + self.act = nn.ReLU(inplace=inplace) + elif act_type == "relu6": + self.act = nn.ReLU6(inplace=inplace) + elif act_type == "sigmoid": + raise NotImplementedError + elif act_type == "hard_sigmoid": + self.act = Hsigmoid( + inplace + ) # nn.Hardsigmoid(inplace=inplace)#Hsigmoid(inplace)# + elif act_type == "hard_swish" or act_type == "hswish": + self.act = Hswish(inplace=inplace) + elif act_type == "leakyrelu": + self.act = nn.LeakyReLU(inplace=inplace) + elif act_type == "gelu": + self.act = GELU(inplace=inplace) + elif act_type == "swish": + self.act = Swish(inplace=inplace) + else: + raise NotImplementedError + + def forward(self, inputs): + return self.act(inputs) diff --git a/python/rapidocr/networks/heads/__init__.py b/python/rapidocr/networks/heads/__init__.py new file mode 100644 index 0000000..00428c4 --- /dev/null +++ b/python/rapidocr/networks/heads/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__all__ = ["build_head"] + + +def build_head(config, **kwargs): + # det head + from .det_db_head import DBHead, PFHeadLocal + + # rec head + from .rec_ctc_head import CTCHead + from .rec_multi_head import MultiHead + + # cls head + from .cls_head import ClsHead + + support_dict = [ + "DBHead", + "CTCHead", + "ClsHead", + "MultiHead", + "PFHeadLocal", + ] + + module_name = config.pop("name") + char_num = config.pop("char_num", 6625) + assert module_name in support_dict, Exception( + "head only support {}".format(support_dict) + ) + module_class = eval(module_name)(**config, **kwargs) + return module_class diff --git a/python/rapidocr/networks/heads/cls_head.py b/python/rapidocr/networks/heads/cls_head.py new file mode 100644 index 0000000..9353b9e --- /dev/null +++ b/python/rapidocr/networks/heads/cls_head.py @@ -0,0 +1,23 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +class ClsHead(nn.Module): + """ + Class orientation + Args: + params(dict): super parameters for build Class network + """ + + def __init__(self, in_channels, class_dim, **kwargs): + super(ClsHead, self).__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Linear(in_channels, class_dim, bias=True) + + def forward(self, x): + x = self.pool(x) + x = torch.reshape(x, shape=[x.shape[0], x.shape[1]]) + x = self.fc(x) + x = F.softmax(x, dim=1) + return x diff --git a/python/rapidocr/networks/heads/det_db_head.py b/python/rapidocr/networks/heads/det_db_head.py new file mode 100644 index 0000000..2de9982 --- /dev/null +++ b/python/rapidocr/networks/heads/det_db_head.py @@ -0,0 +1,131 @@ +import torch +import torch.nn.functional as F +from torch import nn + +from ..backbones.det_mobilenet_v3 import ConvBNLayer +from ..common import Activation + + +class Head(nn.Module): + def __init__(self, in_channels, **kwargs): + super(Head, self).__init__() + self.conv1 = nn.Conv2d( + in_channels=in_channels, + out_channels=in_channels // 4, + kernel_size=3, + padding=1, + bias=False, + ) + self.conv_bn1 = nn.BatchNorm2d(in_channels // 4) + self.relu1 = Activation(act_type="relu") + + self.conv2 = nn.ConvTranspose2d( + in_channels=in_channels // 4, + out_channels=in_channels // 4, + kernel_size=2, + stride=2, + ) + self.conv_bn2 = nn.BatchNorm2d(in_channels // 4) + self.relu2 = Activation(act_type="relu") + + self.conv3 = nn.ConvTranspose2d( + in_channels=in_channels // 4, out_channels=1, kernel_size=2, stride=2 + ) + + def forward(self, x, return_f=False): + x = self.conv1(x) + x = self.conv_bn1(x) + x = self.relu1(x) + x = self.conv2(x) + x = self.conv_bn2(x) + x = self.relu2(x) + if return_f is True: + f = x + x = self.conv3(x) + x = torch.sigmoid(x) + if return_f is True: + return x, f + return x + + +class DBHead(nn.Module): + """ + Differentiable Binarization (DB) for text detection: + see https://arxiv.org/abs/1911.08947 + args: + params(dict): super parameters for build DB network + """ + + def __init__(self, in_channels, k=50, **kwargs): + super(DBHead, self).__init__() + self.k = k + binarize_name_list = [ + "conv2d_56", + "batch_norm_47", + "conv2d_transpose_0", + "batch_norm_48", + "conv2d_transpose_1", + "binarize", + ] + thresh_name_list = [ + "conv2d_57", + "batch_norm_49", + "conv2d_transpose_2", + "batch_norm_50", + "conv2d_transpose_3", + "thresh", + ] + self.binarize = Head(in_channels, **kwargs) # binarize_name_list) + self.thresh = Head(in_channels, **kwargs) # thresh_name_list) + + def step_function(self, x, y): + return torch.reciprocal(1 + torch.exp(-self.k * (x - y))) + + def forward(self, x): + # shrink_maps = self.binarize(x) + return self.binarize(x) + # if not self.training: + # return {'maps': shrink_maps} + + # threshold_maps = self.thresh(x) + # binary_maps = self.step_function(shrink_maps, threshold_maps) + # y = torch.cat([shrink_maps, threshold_maps, binary_maps], dim=1) + # return y + + +class LocalModule(nn.Module): + def __init__(self, in_c, mid_c, use_distance=True): + super(self.__class__, self).__init__() + self.last_3 = ConvBNLayer(in_c + 1, mid_c, 3, 1, 1, act="relu") + self.last_1 = nn.Conv2d(mid_c, 1, 1, 1, 0) + + def forward(self, x, init_map, distance_map): + outf = torch.cat([init_map, x], dim=1) + # last Conv + out = self.last_1(self.last_3(outf)) + return out + + +class PFHeadLocal(DBHead): + def __init__(self, in_channels, k=50, mode="small", **kwargs): + super(PFHeadLocal, self).__init__(in_channels, k, **kwargs) + self.mode = mode + + self.up_conv = nn.Upsample(scale_factor=2, mode="nearest") + if self.mode == "large": + self.cbn_layer = LocalModule(in_channels // 4, in_channels // 4) + elif self.mode == "small": + self.cbn_layer = LocalModule(in_channels // 4, in_channels // 8) + + def forward(self, x, targets=None): + shrink_maps, f = self.binarize(x, return_f=True) + base_maps = shrink_maps + cbn_maps = self.cbn_layer(self.up_conv(f), shrink_maps, None) + cbn_maps = F.sigmoid(cbn_maps) + return 0.5 * (base_maps + cbn_maps) + # return {'maps': 0.5 * (base_maps + cbn_maps), 'cbn_maps': cbn_maps} + + # threshold_maps = self.thresh(x) + # binary_maps = self.step_function(shrink_maps, threshold_maps) + # y = torch.cat([cbn_maps, threshold_maps, binary_maps], dim=1) + # return y diff --git a/python/rapidocr/networks/heads/rec_ctc_head.py b/python/rapidocr/networks/heads/rec_ctc_head.py new file mode 100644 index 0000000..42e2fab --- /dev/null +++ b/python/rapidocr/networks/heads/rec_ctc_head.py @@ -0,0 +1,54 @@ +import torch.nn.functional as F +from torch import nn + + +class CTCHead(nn.Module): + def __init__( + self, + in_channels, + out_channels=6625, + fc_decay=0.0004, + mid_channels=None, + return_feats=False, + **kwargs + ): + super(CTCHead, self).__init__() + if mid_channels is None: + self.fc = nn.Linear( + in_channels, + out_channels, + bias=True, + ) + else: + self.fc1 = nn.Linear( + in_channels, + mid_channels, + bias=True, + ) + self.fc2 = nn.Linear( + mid_channels, + out_channels, + bias=True, + ) + + self.out_channels = out_channels + self.mid_channels = mid_channels + self.return_feats = return_feats + + def forward(self, x, labels=None): + if self.mid_channels is None: + predicts = self.fc(x) + else: + x = self.fc1(x) + predicts = self.fc2(x) + + if self.return_feats: + result = (x, predicts) + else: + result = predicts + + if not self.training: + predicts = F.softmax(predicts, dim=2) + result = predicts + + return result diff --git a/python/rapidocr/networks/heads/rec_multi_head.py b/python/rapidocr/networks/heads/rec_multi_head.py new file mode 100644 index 0000000..b69371b --- /dev/null +++ b/python/rapidocr/networks/heads/rec_multi_head.py @@ -0,0 +1,58 @@ +from torch import nn + +from ..necks.rnn import Im2Seq, SequenceEncoder + + +class FCTranspose(nn.Module): + def __init__(self, in_channels, out_channels, only_transpose=False): + super().__init__() + self.only_transpose = only_transpose + if not self.only_transpose: + self.fc = nn.Linear(in_channels, out_channels, bias=False) + + def forward(self, x): + if self.only_transpose: + return x.permute([0, 2, 1]) + else: + return self.fc(x.permute([0, 2, 1])) + + +class MultiHead(nn.Module): + def __init__(self, in_channels, out_channels_list, **kwargs): + super().__init__() + self.head_list = kwargs.pop("head_list") + + self.gtc_head = "sar" + assert len(self.head_list) >= 2 + for idx, head_name in enumerate(self.head_list): + name = list(head_name)[0] + if name == "SARHead": + pass + + elif name == "NRTRHead": + pass + elif name == "CTCHead": + # ctc neck + self.encoder_reshape = Im2Seq(in_channels) + neck_args = self.head_list[idx][name]["Neck"] + encoder_type = neck_args.pop("name") + self.ctc_encoder = SequenceEncoder( + in_channels=in_channels, encoder_type=encoder_type, **neck_args + ) + # ctc head + head_args = self.head_list[idx][name].get("Head", {}) + if head_args is None: + head_args = {} + self.ctc_head = eval(name)( + in_channels=self.ctc_encoder.out_channels, + out_channels=out_channels_list["CTCLabelDecode"], + **head_args + ) + else: + raise NotImplementedError( + "{} is not supported in MultiHead yet".format(name) + ) + + def forward(self, x, data=None): + ctc_encoder = self.ctc_encoder(x) + return self.ctc_head(ctc_encoder) diff --git a/python/rapidocr/networks/necks/__init__.py b/python/rapidocr/networks/necks/__init__.py new file mode 100644 index 0000000..178ba76 --- /dev/null +++ b/python/rapidocr/networks/necks/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__all__ = ["build_neck"] + + +def build_neck(config): + from .db_fpn import DBFPN, LKPAN, RSEFPN + + support_dict = ["DBFPN", "SequenceEncoder", "RSEFPN", "LKPAN"] + + module_name = config.pop("name") + assert module_name in support_dict, Exception( + "neck only support {}".format(support_dict) + ) + module_class = eval(module_name)(**config) + return module_class diff --git a/python/rapidocr/networks/necks/db_fpn.py b/python/rapidocr/networks/necks/db_fpn.py new file mode 100644 index 0000000..9c8460a --- /dev/null +++ b/python/rapidocr/networks/necks/db_fpn.py @@ -0,0 +1,456 @@ +import torch +import torch.nn.functional as F +from torch import nn + +from ..backbones.det_mobilenet_v3 import SEModule +from ..necks.intracl import IntraCLBlock + + +def hard_swish(x, inplace=True): + return x * F.relu6(x + 3.0, inplace=inplace) / 6.0 + + +class DSConv(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + padding, + stride=1, + groups=None, + if_act=True, + act="relu", + **kwargs + ): + super(DSConv, self).__init__() + if groups == None: + groups = in_channels + self.if_act = if_act + self.act = act + self.conv1 = nn.Conv2d( + in_channels=in_channels, + out_channels=in_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + groups=groups, + bias=False, + ) + + self.bn1 = nn.BatchNorm2d(in_channels) + + self.conv2 = nn.Conv2d( + in_channels=in_channels, + out_channels=int(in_channels * 4), + kernel_size=1, + stride=1, + bias=False, + ) + + self.bn2 = nn.BatchNorm2d(int(in_channels * 4)) + + self.conv3 = nn.Conv2d( + in_channels=int(in_channels * 4), + out_channels=out_channels, + kernel_size=1, + stride=1, + bias=False, + ) + self._c = [in_channels, out_channels] + if in_channels != out_channels: + self.conv_end = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + bias=False, + ) + + def forward(self, inputs): + x = self.conv1(inputs) + x = self.bn1(x) + + x = self.conv2(x) + x = self.bn2(x) + if self.if_act: + if self.act == "relu": + x = F.relu(x) + elif self.act == "hardswish": + x = hard_swish(x) + else: + print( + "The activation function({}) is selected incorrectly.".format( + self.act + ) + ) + exit() + + x = self.conv3(x) + if self._c[0] != self._c[1]: + x = x + self.conv_end(inputs) + return x + + +class DBFPN(nn.Module): + def __init__(self, in_channels, out_channels, use_asf=False, **kwargs): + super(DBFPN, self).__init__() + self.out_channels = out_channels + self.use_asf = use_asf + + self.in2_conv = nn.Conv2d( + in_channels=in_channels[0], + out_channels=self.out_channels, + kernel_size=1, + bias=False, + ) + self.in3_conv = nn.Conv2d( + in_channels=in_channels[1], + out_channels=self.out_channels, + kernel_size=1, + bias=False, + ) + self.in4_conv = nn.Conv2d( + in_channels=in_channels[2], + out_channels=self.out_channels, + kernel_size=1, + bias=False, + ) + self.in5_conv = nn.Conv2d( + in_channels=in_channels[3], + out_channels=self.out_channels, + kernel_size=1, + bias=False, + ) + self.p5_conv = nn.Conv2d( + in_channels=self.out_channels, + out_channels=self.out_channels // 4, + kernel_size=3, + padding=1, + bias=False, + ) + self.p4_conv = nn.Conv2d( + in_channels=self.out_channels, + out_channels=self.out_channels // 4, + kernel_size=3, + padding=1, + bias=False, + ) + self.p3_conv = nn.Conv2d( + in_channels=self.out_channels, + out_channels=self.out_channels // 4, + kernel_size=3, + padding=1, + bias=False, + ) + self.p2_conv = nn.Conv2d( + in_channels=self.out_channels, + out_channels=self.out_channels // 4, + kernel_size=3, + padding=1, + bias=False, + ) + + if self.use_asf is True: + self.asf = ASFBlock(self.out_channels, self.out_channels // 4) + + def forward(self, x): + c2, c3, c4, c5 = x + + in5 = self.in5_conv(c5) + in4 = self.in4_conv(c4) + in3 = self.in3_conv(c3) + in2 = self.in2_conv(c2) + + out4 = in4 + F.interpolate( + in5, + scale_factor=2, + mode="nearest", + ) # align_mode=1) # 1/16 + out3 = in3 + F.interpolate( + out4, + scale_factor=2, + mode="nearest", + ) # align_mode=1) # 1/8 + out2 = in2 + F.interpolate( + out3, + scale_factor=2, + mode="nearest", + ) # align_mode=1) # 1/4 + + p5 = self.p5_conv(in5) + p4 = self.p4_conv(out4) + p3 = self.p3_conv(out3) + p2 = self.p2_conv(out2) + p5 = F.interpolate( + p5, + scale_factor=8, + mode="nearest", + ) # align_mode=1) + p4 = F.interpolate( + p4, + scale_factor=4, + mode="nearest", + ) # align_mode=1) + p3 = F.interpolate( + p3, + scale_factor=2, + mode="nearest", + ) # align_mode=1) + + fuse = torch.cat([p5, p4, p3, p2], dim=1) + + if self.use_asf is True: + fuse = self.asf(fuse, [p5, p4, p3, p2]) + + return fuse + + +class RSELayer(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, shortcut=True): + super(RSELayer, self).__init__() + self.out_channels = out_channels + self.in_conv = nn.Conv2d( + in_channels=in_channels, + out_channels=self.out_channels, + kernel_size=kernel_size, + padding=int(kernel_size // 2), + bias=False, + ) + self.se_block = SEModule(self.out_channels) + self.shortcut = shortcut + + def forward(self, ins): + x = self.in_conv(ins) + if self.shortcut: + out = x + self.se_block(x) + else: + out = self.se_block(x) + return out + + +class RSEFPN(nn.Module): + def __init__(self, in_channels, out_channels, shortcut=True, **kwargs): + super(RSEFPN, self).__init__() + self.out_channels = out_channels + self.ins_conv = nn.ModuleList() + self.inp_conv = nn.ModuleList() + self.intracl = False + if "intracl" in kwargs.keys() and kwargs["intracl"] is True: + self.intracl = kwargs["intracl"] + self.incl1 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + self.incl2 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + self.incl3 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + self.incl4 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + + for i in range(len(in_channels)): + self.ins_conv.append( + RSELayer(in_channels[i], out_channels, kernel_size=1, shortcut=shortcut) + ) + self.inp_conv.append( + RSELayer( + out_channels, out_channels // 4, kernel_size=3, shortcut=shortcut + ) + ) + + def forward(self, x): + c2, c3, c4, c5 = x + + in5 = self.ins_conv[3](c5) + in4 = self.ins_conv[2](c4) + in3 = self.ins_conv[1](c3) + in2 = self.ins_conv[0](c2) + + out4 = in4 + F.interpolate(in5, scale_factor=2, mode="nearest") # 1/16 + out3 = in3 + F.interpolate(out4, scale_factor=2, mode="nearest") # 1/8 + out2 = in2 + F.interpolate(out3, scale_factor=2, mode="nearest") # 1/4 + + p5 = self.inp_conv[3](in5) + p4 = self.inp_conv[2](out4) + p3 = self.inp_conv[1](out3) + p2 = self.inp_conv[0](out2) + + if self.intracl is True: + p5 = self.incl4(p5) + p4 = self.incl3(p4) + p3 = self.incl2(p3) + p2 = self.incl1(p2) + + p5 = F.interpolate(p5, scale_factor=8, mode="nearest") + p4 = F.interpolate(p4, scale_factor=4, mode="nearest") + p3 = F.interpolate(p3, scale_factor=2, mode="nearest") + + fuse = torch.cat([p5, p4, p3, p2], dim=1) + return fuse + + +class LKPAN(nn.Module): + def __init__(self, in_channels, out_channels, mode="large", **kwargs): + super(LKPAN, self).__init__() + self.out_channels = out_channels + + self.ins_conv = nn.ModuleList() + self.inp_conv = nn.ModuleList() + # pan head + self.pan_head_conv = nn.ModuleList() + self.pan_lat_conv = nn.ModuleList() + + if mode.lower() == "lite": + p_layer = DSConv + elif mode.lower() == "large": + p_layer = nn.Conv2d + else: + raise ValueError( + "mode can only be one of ['lite', 'large'], but received {}".format( + mode + ) + ) + + for i in range(len(in_channels)): + self.ins_conv.append( + nn.Conv2d( + in_channels=in_channels[i], + out_channels=self.out_channels, + kernel_size=1, + bias=False, + ) + ) + + self.inp_conv.append( + p_layer( + in_channels=self.out_channels, + out_channels=self.out_channels // 4, + kernel_size=9, + padding=4, + bias=False, + ) + ) + + if i > 0: + self.pan_head_conv.append( + nn.Conv2d( + in_channels=self.out_channels // 4, + out_channels=self.out_channels // 4, + kernel_size=3, + padding=1, + stride=2, + bias=False, + ) + ) + self.pan_lat_conv.append( + p_layer( + in_channels=self.out_channels // 4, + out_channels=self.out_channels // 4, + kernel_size=9, + padding=4, + bias=False, + ) + ) + self.intracl = False + if "intracl" in kwargs.keys() and kwargs["intracl"] is True: + self.intracl = kwargs["intracl"] + self.incl1 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + self.incl2 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + self.incl3 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + self.incl4 = IntraCLBlock(self.out_channels // 4, reduce_factor=2) + + def forward(self, x): + c2, c3, c4, c5 = x + + in5 = self.ins_conv[3](c5) + in4 = self.ins_conv[2](c4) + in3 = self.ins_conv[1](c3) + in2 = self.ins_conv[0](c2) + + out4 = in4 + F.interpolate(in5, scale_factor=2, mode="nearest") # 1/16 + out3 = in3 + F.interpolate(out4, scale_factor=2, mode="nearest") # 1/8 + out2 = in2 + F.interpolate(out3, scale_factor=2, mode="nearest") # 1/4 + + f5 = self.inp_conv[3](in5) + f4 = self.inp_conv[2](out4) + f3 = self.inp_conv[1](out3) + f2 = self.inp_conv[0](out2) + + pan3 = f3 + self.pan_head_conv[0](f2) + pan4 = f4 + self.pan_head_conv[1](pan3) + pan5 = f5 + self.pan_head_conv[2](pan4) + + p2 = self.pan_lat_conv[0](f2) + p3 = self.pan_lat_conv[1](pan3) + p4 = self.pan_lat_conv[2](pan4) + p5 = self.pan_lat_conv[3](pan5) + + if self.intracl is True: + p5 = self.incl4(p5) + p4 = self.incl3(p4) + p3 = self.incl2(p3) + p2 = self.incl1(p2) + + p5 = F.interpolate(p5, scale_factor=8, mode="nearest") + p4 = F.interpolate(p4, scale_factor=4, mode="nearest") + p3 = F.interpolate(p3, scale_factor=2, mode="nearest") + + fuse = torch.cat([p5, p4, p3, p2], dim=1) + return fuse + + +class ASFBlock(nn.Module): + """ + This code is refered from: + https://github.com/MhLiao/DB/blob/master/decoders/feature_attention.py + """ + + def __init__(self, in_channels, inter_channels, out_features_num=4): + """ + Adaptive Scale Fusion (ASF) block of DBNet++ + Args: + in_channels: the number of channels in the input data + inter_channels: the number of middle channels + out_features_num: the number of fused stages + """ + super(ASFBlock, self).__init__() + self.in_channels = in_channels + self.inter_channels = inter_channels + self.out_features_num = out_features_num + self.conv = nn.Conv2d(in_channels, inter_channels, 3, padding=1) + + self.spatial_scale = nn.Sequential( + # Nx1xHxW + nn.Conv2d( + in_channels=1, + out_channels=1, + kernel_size=3, + bias=False, + padding=1, + ), + nn.ReLU(), + nn.Conv2d( + in_channels=1, + out_channels=1, + kernel_size=1, + bias=False, + ), + nn.Sigmoid(), + ) + + self.channel_scale = nn.Sequential( + nn.Conv2d( + in_channels=inter_channels, + out_channels=out_features_num, + kernel_size=1, + bias=False, + ), + nn.Sigmoid(), + ) + + def forward(self, fuse_features, features_list): + fuse_features = self.conv(fuse_features) + spatial_x = torch.mean(fuse_features, dim=1, keepdim=True) + attention_scores = self.spatial_scale(spatial_x) + fuse_features + attention_scores = self.channel_scale(attention_scores) + assert len(features_list) == self.out_features_num + + out_list = [] + for i in range(self.out_features_num): + out_list.append(attention_scores[:, i : i + 1] * features_list[i]) + return torch.cat(out_list, dim=1) diff --git a/python/rapidocr/networks/necks/intracl.py b/python/rapidocr/networks/necks/intracl.py new file mode 100644 index 0000000..0ba85fa --- /dev/null +++ b/python/rapidocr/networks/necks/intracl.py @@ -0,0 +1,117 @@ +from torch import nn + + +class IntraCLBlock(nn.Module): + def __init__(self, in_channels=96, reduce_factor=4): + super(IntraCLBlock, self).__init__() + self.channels = in_channels + self.rf = reduce_factor + self.conv1x1_reduce_channel = nn.Conv2d( + self.channels, self.channels // self.rf, kernel_size=1, stride=1, padding=0 + ) + self.conv1x1_return_channel = nn.Conv2d( + self.channels // self.rf, self.channels, kernel_size=1, stride=1, padding=0 + ) + + self.v_layer_7x1 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(7, 1), + stride=(1, 1), + padding=(3, 0), + ) + self.v_layer_5x1 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(5, 1), + stride=(1, 1), + padding=(2, 0), + ) + self.v_layer_3x1 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(3, 1), + stride=(1, 1), + padding=(1, 0), + ) + + self.q_layer_1x7 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(1, 7), + stride=(1, 1), + padding=(0, 3), + ) + self.q_layer_1x5 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(1, 5), + stride=(1, 1), + padding=(0, 2), + ) + self.q_layer_1x3 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(1, 3), + stride=(1, 1), + padding=(0, 1), + ) + + # base + self.c_layer_7x7 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(7, 7), + stride=(1, 1), + padding=(3, 3), + ) + self.c_layer_5x5 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(5, 5), + stride=(1, 1), + padding=(2, 2), + ) + self.c_layer_3x3 = nn.Conv2d( + self.channels // self.rf, + self.channels // self.rf, + kernel_size=(3, 3), + stride=(1, 1), + padding=(1, 1), + ) + + self.bn = nn.BatchNorm2d(self.channels) + self.relu = nn.ReLU() + + def forward(self, x): + x_new = self.conv1x1_reduce_channel(x) + + x_7_c = self.c_layer_7x7(x_new) + x_7_v = self.v_layer_7x1(x_new) + x_7_q = self.q_layer_1x7(x_new) + x_7 = x_7_c + x_7_v + x_7_q + + x_5_c = self.c_layer_5x5(x_7) + x_5_v = self.v_layer_5x1(x_7) + x_5_q = self.q_layer_1x5(x_7) + x_5 = x_5_c + x_5_v + x_5_q + + x_3_c = self.c_layer_3x3(x_5) + x_3_v = self.v_layer_3x1(x_5) + x_3_q = self.q_layer_1x3(x_5) + x_3 = x_3_c + x_3_v + x_3_q + + x_relation = self.conv1x1_return_channel(x_3) + + x_relation = self.bn(x_relation) + x_relation = self.relu(x_relation) + + return x + x_relation + + +def build_intraclblock_list(num_block): + IntraCLBlock_list = nn.ModuleList() + for i in range(num_block): + IntraCLBlock_list.append(IntraCLBlock()) + + return IntraCLBlock_list diff --git a/python/rapidocr/networks/necks/rnn.py b/python/rapidocr/networks/necks/rnn.py new file mode 100644 index 0000000..244ea44 --- /dev/null +++ b/python/rapidocr/networks/necks/rnn.py @@ -0,0 +1,228 @@ +import torch +from torch import nn + +from ..backbones.rec_svtrnet import Block, ConvBNLayer + + +class Im2Seq(nn.Module): + def __init__(self, in_channels, **kwargs): + super().__init__() + self.out_channels = in_channels + + def forward(self, x): + B, C, H, W = x.shape + # assert H == 1 + x = x.squeeze(dim=2) + # x = x.transpose([0, 2, 1]) # paddle (NTC)(batch, width, channels) + x = x.permute(0, 2, 1) + return x + + +class EncoderWithRNN_(nn.Module): + def __init__(self, in_channels, hidden_size): + super(EncoderWithRNN_, self).__init__() + self.out_channels = hidden_size * 2 + self.rnn1 = nn.LSTM( + in_channels, + hidden_size, + bidirectional=False, + batch_first=True, + num_layers=2, + ) + self.rnn2 = nn.LSTM( + in_channels, + hidden_size, + bidirectional=False, + batch_first=True, + num_layers=2, + ) + + def forward(self, x): + self.rnn1.flatten_parameters() + self.rnn2.flatten_parameters() + out1, h1 = self.rnn1(x) + out2, h2 = self.rnn2(torch.flip(x, [1])) + return torch.cat([out1, torch.flip(out2, [1])], 2) + + +class EncoderWithRNN(nn.Module): + def __init__(self, in_channels, hidden_size): + super(EncoderWithRNN, self).__init__() + self.out_channels = hidden_size * 2 + self.lstm = nn.LSTM( + in_channels, hidden_size, num_layers=2, batch_first=True, bidirectional=True + ) # batch_first:=True + + def forward(self, x): + x, _ = self.lstm(x) + return x + + +class EncoderWithFC(nn.Module): + def __init__(self, in_channels, hidden_size): + super(EncoderWithFC, self).__init__() + self.out_channels = hidden_size + self.fc = nn.Linear( + in_channels, + hidden_size, + bias=True, + ) + + def forward(self, x): + x = self.fc(x) + return x + + +class EncoderWithSVTR(nn.Module): + def __init__( + self, + in_channels, + dims=64, # XS + depth=2, + hidden_dims=120, + use_guide=False, + num_heads=8, + qkv_bias=True, + mlp_ratio=2.0, + drop_rate=0.1, + kernel_size=[3, 3], + attn_drop_rate=0.1, + drop_path=0.0, + qk_scale=None, + ): + super(EncoderWithSVTR, self).__init__() + self.depth = depth + self.use_guide = use_guide + self.conv1 = ConvBNLayer( + in_channels, + in_channels // 8, + kernel_size=kernel_size, + padding=[kernel_size[0] // 2, kernel_size[1] // 2], + act="swish", + ) + self.conv2 = ConvBNLayer( + in_channels // 8, hidden_dims, kernel_size=1, act="swish" + ) + + self.svtr_block = nn.ModuleList( + [ + Block( + dim=hidden_dims, + num_heads=num_heads, + mixer="Global", + HW=None, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + act_layer="swish", + attn_drop=attn_drop_rate, + drop_path=drop_path, + norm_layer="nn.LayerNorm", + epsilon=1e-05, + prenorm=False, + ) + for i in range(depth) + ] + ) + self.norm = nn.LayerNorm(hidden_dims, eps=1e-6) + self.conv3 = ConvBNLayer(hidden_dims, in_channels, kernel_size=1, act="swish") + # last conv-nxn, the input is concat of input tensor and conv3 output tensor + self.conv4 = ConvBNLayer( + 2 * in_channels, in_channels // 8, padding=1, act="swish" + ) + + self.conv1x1 = ConvBNLayer(in_channels // 8, dims, kernel_size=1, act="swish") + self.out_channels = dims + self.apply(self._init_weights) + + def _init_weights(self, m): + # weight initialization + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.BatchNorm2d): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.ConvTranspose2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out") + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + nn.init.ones_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, x): + # for use guide + if self.use_guide: + z = x.clone() + z.stop_gradient = True + else: + z = x + # for short cut + h = z + # reduce dim + z = self.conv1(z) + z = self.conv2(z) + # SVTR global block + B, C, H, W = z.shape + z = z.flatten(2).permute(0, 2, 1) + + for blk in self.svtr_block: + z = blk(z) + + z = self.norm(z) + # last stage + z = z.reshape([-1, H, W, C]).permute(0, 3, 1, 2) + z = self.conv3(z) + z = torch.cat((h, z), dim=1) + z = self.conv1x1(self.conv4(z)) + + return z + + +class SequenceEncoder(nn.Module): + def __init__(self, in_channels, encoder_type, hidden_size=48, **kwargs): + super(SequenceEncoder, self).__init__() + self.encoder_reshape = Im2Seq(in_channels) + self.out_channels = self.encoder_reshape.out_channels + self.encoder_type = encoder_type + if encoder_type == "reshape": + self.only_reshape = True + else: + support_encoder_dict = { + "reshape": Im2Seq, + "fc": EncoderWithFC, + "rnn": EncoderWithRNN, + "svtr": EncoderWithSVTR, + } + assert encoder_type in support_encoder_dict, "{} must in {}".format( + encoder_type, support_encoder_dict.keys() + ) + + if encoder_type == "svtr": + self.encoder = support_encoder_dict[encoder_type]( + self.encoder_reshape.out_channels, **kwargs + ) + else: + self.encoder = support_encoder_dict[encoder_type]( + self.encoder_reshape.out_channels, hidden_size + ) + self.out_channels = self.encoder.out_channels + self.only_reshape = False + + def forward(self, x): + if self.encoder_type != "svtr": + x = self.encoder_reshape(x) + if not self.only_reshape: + x = self.encoder(x) + return x + else: + x = self.encoder(x) + x = self.encoder_reshape(x) + return x diff --git a/python/rapidocr/utils/parse_parameters.py b/python/rapidocr/utils/parse_parameters.py index b946fed..f9c0064 100644 --- a/python/rapidocr/utils/parse_parameters.py +++ b/python/rapidocr/utils/parse_parameters.py @@ -100,6 +100,10 @@ def parse_lang(lang): lang = "cyrillic" elif lang in devanagari_lang: lang = "devanagari" + elif lang in ["ch", "en"]: + pass + else: + raise ValueError(f"lang: {lang} is not in the supported list.") if lang == "ch": det_lang = "ch" diff --git a/python/rapidocr_torch/modeling/common.py b/python/rapidocr_torch/modeling/common.py index 2d6e3d9..7378093 100644 --- a/python/rapidocr_torch/modeling/common.py +++ b/python/rapidocr_torch/modeling/common.py @@ -1,16 +1,16 @@ - - import torch import torch.nn as nn import torch.nn.functional as F + class Hswish(nn.Module): def __init__(self, inplace=True): super(Hswish, self).__init__() self.inplace = inplace def forward(self, x): - return x * F.relu6(x + 3., inplace=self.inplace) / 6. + return x * F.relu6(x + 3.0, inplace=self.inplace) / 6.0 + # out = max(0, min(1, slop*x+offset)) # paddle.fluid.layers.hard_sigmoid(x, slope=0.2, offset=0.5, name=None) @@ -22,7 +22,8 @@ class Hsigmoid(nn.Module): def forward(self, x): # torch: F.relu6(x + 3., inplace=self.inplace) / 6. # paddle: F.relu6(1.2 * x + 3., inplace=self.inplace) / 6. - return F.relu6(1.2 * x + 3., inplace=self.inplace) / 6. + return F.relu6(1.2 * x + 3.0, inplace=self.inplace) / 6.0 + class GELU(nn.Module): def __init__(self, inplace=True): @@ -43,31 +44,33 @@ class Swish(nn.Module): x.mul_(torch.sigmoid(x)) return x else: - return x*torch.sigmoid(x) + return x * torch.sigmoid(x) class Activation(nn.Module): def __init__(self, act_type, inplace=True): super(Activation, self).__init__() act_type = act_type.lower() - if act_type == 'relu': + if act_type == "relu": self.act = nn.ReLU(inplace=inplace) - elif act_type == 'relu6': + elif act_type == "relu6": self.act = nn.ReLU6(inplace=inplace) - elif act_type == 'sigmoid': + elif act_type == "sigmoid": raise NotImplementedError - elif act_type == 'hard_sigmoid': - self.act = Hsigmoid(inplace)#nn.Hardsigmoid(inplace=inplace)#Hsigmoid(inplace)# - elif act_type == 'hard_swish' or act_type == 'hswish': + elif act_type == "hard_sigmoid": + self.act = Hsigmoid( + inplace + ) # nn.Hardsigmoid(inplace=inplace)#Hsigmoid(inplace)# + elif act_type == "hard_swish" or act_type == "hswish": self.act = Hswish(inplace=inplace) - elif act_type == 'leakyrelu': + elif act_type == "leakyrelu": self.act = nn.LeakyReLU(inplace=inplace) - elif act_type == 'gelu': + elif act_type == "gelu": self.act = GELU(inplace=inplace) - elif act_type == 'swish': + elif act_type == "swish": self.act = Swish(inplace=inplace) else: raise NotImplementedError def forward(self, inputs): - return self.act(inputs) \ No newline at end of file + return self.act(inputs) diff --git a/python/rapidocr_torch/modeling/necks/intracl.py b/python/rapidocr_torch/modeling/necks/intracl.py index 414b723..0ba85fa 100644 --- a/python/rapidocr_torch/modeling/necks/intracl.py +++ b/python/rapidocr_torch/modeling/necks/intracl.py @@ -7,55 +7,55 @@ class IntraCLBlock(nn.Module): self.channels = in_channels self.rf = reduce_factor self.conv1x1_reduce_channel = nn.Conv2d( - self.channels, - self.channels // self.rf, - kernel_size=1, - stride=1, - padding=0) + self.channels, self.channels // self.rf, kernel_size=1, stride=1, padding=0 + ) self.conv1x1_return_channel = nn.Conv2d( - self.channels // self.rf, - self.channels, - kernel_size=1, - stride=1, - padding=0) + self.channels // self.rf, self.channels, kernel_size=1, stride=1, padding=0 + ) self.v_layer_7x1 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(7, 1), stride=(1, 1), - padding=(3, 0)) + padding=(3, 0), + ) self.v_layer_5x1 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(5, 1), stride=(1, 1), - padding=(2, 0)) + padding=(2, 0), + ) self.v_layer_3x1 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(3, 1), stride=(1, 1), - padding=(1, 0)) + padding=(1, 0), + ) self.q_layer_1x7 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(1, 7), stride=(1, 1), - padding=(0, 3)) + padding=(0, 3), + ) self.q_layer_1x5 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(1, 5), stride=(1, 1), - padding=(0, 2)) + padding=(0, 2), + ) self.q_layer_1x3 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(1, 3), stride=(1, 1), - padding=(0, 1)) + padding=(0, 1), + ) # base self.c_layer_7x7 = nn.Conv2d( @@ -63,19 +63,22 @@ class IntraCLBlock(nn.Module): self.channels // self.rf, kernel_size=(7, 7), stride=(1, 1), - padding=(3, 3)) + padding=(3, 3), + ) self.c_layer_5x5 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(5, 5), stride=(1, 1), - padding=(2, 2)) + padding=(2, 2), + ) self.c_layer_3x3 = nn.Conv2d( self.channels // self.rf, self.channels // self.rf, kernel_size=(3, 3), stride=(1, 1), - padding=(1, 1)) + padding=(1, 1), + ) self.bn = nn.BatchNorm2d(self.channels) self.relu = nn.ReLU() @@ -111,4 +114,4 @@ def build_intraclblock_list(num_block): for i in range(num_block): IntraCLBlock_list.append(IntraCLBlock()) - return IntraCLBlock_list \ No newline at end of file + return IntraCLBlock_list diff --git a/python/tests/test_files/config.yaml b/python/tests/test_files/config.yaml deleted file mode 100644 index b4f114b..0000000 --- a/python/tests/test_files/config.yaml +++ /dev/null @@ -1,62 +0,0 @@ -Global: - lang: "ch" - text_score: 0.5 - use_det: true - use_cls: true - use_rec: true - - print_verbose: false - min_height: 30 - width_height_ratio: 8 - max_side_len: 2000 - min_side_len: 30 - return_word_box: false - - with_onnx: false - with_openvino: false - with_paddle: false - with_torch: false - -EngineConfig: - onnxruntime: - intra_op_num_threads: -1 - inter_op_num_threads: -1 - use_cuda: false - use_dml: false - - paddlepaddle: - cpu_math_library_num_threads: -1 - gpu_id: 0 - gpu_mem: 500 - -Det: - model_path: models/ch_PP-OCRv4_det_infer.onnx - paddle_model_dir: models/paddle/ch_PP-OCRv4_det_infer - - limit_side_len: 736 - limit_type: min - std: [ 0.5, 0.5, 0.5 ] - mean: [ 0.5, 0.5, 0.5 ] - - thresh: 0.3 - box_thresh: 0.5 - max_candidates: 1000 - unclip_ratio: 1.6 - use_dilation: true - score_mode: fast - -Cls: - model_path: models/ch_ppocr_mobile_v2.0_cls_infer.onnx - paddle_model_dir: models/paddle/ch_ppocr_mobile_v2_cls_infer - - cls_image_shape: [3, 48, 192] - cls_batch_num: 6 - cls_thresh: 0.9 - label_list: ['0', '180'] - -Rec: - model_path: models/ch_PP-OCRv4_rec_infer.onnx - paddle_model_dir: models/paddle/ch_PP-OCRv4_rec_infer - - rec_img_shape: [3, 48, 320] - rec_batch_num: 6 diff --git a/python/tests/test_files/en.jpg b/python/tests/test_files/en.jpg new file mode 100644 index 0000000..bcfe003 Binary files /dev/null and b/python/tests/test_files/en.jpg differ diff --git a/python/tests/test_main.py b/python/tests/test_main.py index 6de89f8..2fe3f7e 100644 --- a/python/tests/test_main.py +++ b/python/tests/test_main.py @@ -34,11 +34,18 @@ def get_engine(params: Optional[Dict[str, Any]] = None): return engine +def test_error_lang(): + with pytest.raises(ValueError) as exc_info: + engine = get_engine(params={"Global.lang": "eh"}) + assert exc_info.type is ValueError + + def test_lang(): - engine = get_engine(params={"Global.lang": "ch", "Global.with_openvino": True}) + engine = get_engine(params={"Global.lang": "en"}) + img_path = tests_dir / "en.jpg" result = engine(img_path) assert result.txts is not None - assert result.txts[0] == "正品促销" + assert result.txts[0] == "MovieShots Dataset" def test_engine_openvino():