From b64a46ffd4fc370e4e2c5a10c485cb26e40b8b0f Mon Sep 17 00:00:00 2001 From: SWHL Date: Sat, 4 Jul 2026 22:42:55 +0800 Subject: [PATCH] fix: align the performance of the torch version of the PP-OCRv6 Rec tiny model --- python/rapidocr/default_models.yaml | 4 +- .../pytorch/networks/arch_config.yaml | 3 ++ .../pytorch/networks/backbones/__init__.py | 46 +++++++---------- .../pytorch/networks/heads/rec_ctc_head.py | 49 +++++++++++++------ .../pytorch/networks/heads/rec_multi_head.py | 28 +++++++++-- .../pytorch/networks/necks/__init__.py | 24 ++++----- .../pytorch/networks/necks/db_fpn.py | 3 +- 7 files changed, 93 insertions(+), 64 deletions(-) diff --git a/python/rapidocr/default_models.yaml b/python/rapidocr/default_models.yaml index e1b031f..a5474dc 100644 --- a/python/rapidocr/default_models.yaml +++ b/python/rapidocr/default_models.yaml @@ -802,8 +802,8 @@ torch: SHA256: 40c40ff2d432888f53697d73b37d8b280c65ad66696d6412834a90b40cebc81d rec: multi_PP-OCRv6_rec_tiny: - model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.1/torch/PP-OCRv6/rec/PP-OCRv6_rec_tiny.pth - SHA256: 82a9b90ec19bcdbc6f2cd54f7b08394aeb01d947b9dfc43e483b90fb30238b85 + model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/torch/PP-OCRv6/rec/PP-OCRv6_rec_tiny.pth + SHA256: dceca3cd44855565e9c10391f22767aa2937c2195df3ce7dd0d9d3d5b6bfa011 dict_url: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.1/paddle/PP-OCRv6/rec/PP-OCRv6_rec_tiny/ppocrv6_tiny_dict.txt multi_PP-OCRv6_rec_small: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.1/torch/PP-OCRv6/rec/PP-OCRv6_rec_small.pth diff --git a/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml b/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml index f92ce39..30c51cc 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml +++ b/python/rapidocr/inference_engine/pytorch/networks/arch_config.yaml @@ -9,6 +9,8 @@ PP-OCRv6_rec_tiny: name: MultiHead out_channels_list: CTCLabelDecode: 6906 #'blank' + ...(6904) + ' ' + NRTRLabelDecode: 6909 + build_gtc_head: false head_list: - CTCHead: Neck: @@ -16,6 +18,7 @@ PP-OCRv6_rec_tiny: Head: mid_channels: 80 use_guide: true + guide_layer_type: ppocrv6_tiny fc_decay: 0.00001 - NRTRHead: nrtr_dim: 384 diff --git a/python/rapidocr/inference_engine/pytorch/networks/backbones/__init__.py b/python/rapidocr/inference_engine/pytorch/networks/backbones/__init__.py index ea18e42..eb4cfcc 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/backbones/__init__.py +++ b/python/rapidocr/inference_engine/pytorch/networks/backbones/__init__.py @@ -23,16 +23,14 @@ def build_backbone(config, model_type): from .rec_lcnetv4 import PPLCNetV4 from .rec_pphgnetv2 import PPHGNetV2_B4 - support_dict = [ - "MobileNetV3", - "ResNet", - "ResNet_vd", - "ResNet_SAST", - "PPLCNetV3", - "PPHGNet_small", - "PPHGNetV2_B4", - "PPLCNetV4", - ] + support_dict = { + "MobileNetV3": MobileNetV3, + "PPLCNetV3": PPLCNetV3, + "PPLCNetV4": PPLCNetV4, + "PPHGNet_small": PPHGNet_small, + "PPHGNetV2_B4": PPHGNetV2_B4, + } + elif model_type == "rec" or model_type == "cls": from .rec_hgnet import PPHGNet_small from .rec_lcnetv3 import PPLCNetV3 @@ -45,22 +43,16 @@ def build_backbone(config, model_type): from .rec_pplcnet import PPLCNet from .rec_svtrnet import SVTRNet - support_dict = [ - "MobileNetV1Enhance", - "MobileNetV3", - "ResNet", - "ResNetFPN", - "MTB", - "ResNet31", - "SVTRNet", - "ViTSTR", - "DenseNet", - "PPLCNetV3", - "PPHGNet_small", - "PPHGNetV2_B4", - "PPLCNetV4", - "PPLCNet", - ] + support_dict = { + "MobileNetV1Enhance": MobileNetV1Enhance, + "MobileNetV3": MobileNetV3, + "SVTRNet": SVTRNet, + "PPLCNetV3": PPLCNetV3, + "PPLCNetV4": PPLCNetV4, + "PPHGNet_small": PPHGNet_small, + "PPHGNetV2_B4": PPHGNetV2_B4, + "PPLCNet": PPLCNet, + } else: raise NotImplementedError @@ -70,5 +62,5 @@ def build_backbone(config, model_type): model_type, support_dict ) ) - module_class = eval(module_name)(**config) + module_class = support_dict[module_name](**config) return module_class diff --git a/python/rapidocr/inference_engine/pytorch/networks/heads/rec_ctc_head.py b/python/rapidocr/inference_engine/pytorch/networks/heads/rec_ctc_head.py index d11c7eb..dda4502 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/heads/rec_ctc_head.py +++ b/python/rapidocr/inference_engine/pytorch/networks/heads/rec_ctc_head.py @@ -11,27 +11,44 @@ class CTCHead(nn.Module): mid_channels=None, return_feats=False, use_guide=False, + guide_layer_type="default", **kwargs, ): super(CTCHead, self).__init__() self.use_guide = use_guide if use_guide: - # Depthwise-separable Conv1d block: - # DW Conv1d(k=5, groups=in_channels) → BN → ReLU → PW Conv1d(k=1) → BN - self.guide_layer = nn.Sequential( - nn.Conv1d( - in_channels, - in_channels, - 5, - padding=2, - groups=in_channels, - bias=True, - ), - nn.BatchNorm1d(in_channels), - nn.ReLU(), - nn.Conv1d(in_channels, in_channels, 1, bias=True), - nn.BatchNorm1d(in_channels), - ) + if guide_layer_type == "ppocrv6_tiny": + self.guide_layer = nn.Sequential( + nn.Conv1d( + in_channels, + in_channels, + 5, + padding=2, + groups=in_channels, + bias=False, + ), + nn.BatchNorm1d(in_channels), + nn.Hardswish(), + nn.Conv1d(in_channels, in_channels, 1, bias=False), + nn.BatchNorm1d(in_channels), + nn.Hardswish(), + ) + else: + self.guide_layer = nn.Sequential( + nn.Conv1d( + in_channels, + in_channels, + 5, + padding=2, + groups=in_channels, + bias=True, + ), + nn.BatchNorm1d(in_channels), + nn.ReLU(), + nn.Conv1d(in_channels, in_channels, 1, bias=True), + nn.BatchNorm1d(in_channels), + ) + if mid_channels is None: self.fc = nn.Linear( in_channels, diff --git a/python/rapidocr/inference_engine/pytorch/networks/heads/rec_multi_head.py b/python/rapidocr/inference_engine/pytorch/networks/heads/rec_multi_head.py index cd14349..a104f73 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/heads/rec_multi_head.py +++ b/python/rapidocr/inference_engine/pytorch/networks/heads/rec_multi_head.py @@ -1,7 +1,10 @@ +import copy + from torch import nn from ..necks.rnn import Im2Seq, SequenceEncoder from .rec_ctc_head import CTCHead +from .rec_nrtr_head import Transformer class FCTranspose(nn.Module): @@ -20,7 +23,8 @@ class FCTranspose(nn.Module): class MultiHead(nn.Module): def __init__(self, in_channels, out_channels_list, **kwargs): super().__init__() - self.head_list = kwargs.pop("head_list") + self.head_list = copy.deepcopy(kwargs.pop("head_list")) + self.build_gtc_head = kwargs.pop("build_gtc_head", False) self.gtc_head = "sar" assert len(self.head_list) >= 2 @@ -29,7 +33,25 @@ class MultiHead(nn.Module): if name == "SARHead": pass elif name == "NRTRHead": - pass + if not self.build_gtc_head: + continue + gtc_args = self.head_list[idx][name] + max_text_length = gtc_args.get("max_text_length", 25) + nrtr_dim = gtc_args.get("nrtr_dim", 256) + num_decoder_layers = gtc_args.get("num_decoder_layers", 4) + self.before_gtc = nn.Sequential( + nn.Flatten(2), FCTranspose(in_channels, nrtr_dim) + ) + self.gtc_head = Transformer( + d_model=nrtr_dim, + nhead=nrtr_dim // 32, + num_encoder_layers=-1, + beam_size=-1, + num_decoder_layers=num_decoder_layers, + max_len=max_text_length, + dim_feedforward=nrtr_dim * 4, + out_channels=out_channels_list["NRTRLabelDecode"], + ) elif name == "CTCHead": # ctc neck self.encoder_reshape = Im2Seq(in_channels) @@ -65,6 +87,6 @@ class MultiHead(nn.Module): sar_out = self.sar_head(x, data[1:])["res"] head_out["sar"] = sar_out else: - gtc_out = self.gtc_head(self.before_gtc(x), data[1:])["res"] + gtc_out = self.gtc_head(self.before_gtc(x), data[1:]) head_out["nrtr"] = gtc_out return head_out diff --git a/python/rapidocr/inference_engine/pytorch/networks/necks/__init__.py b/python/rapidocr/inference_engine/pytorch/networks/necks/__init__.py index bc5c479..9252952 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/necks/__init__.py +++ b/python/rapidocr/inference_engine/pytorch/networks/necks/__init__.py @@ -19,24 +19,18 @@ def build_neck(config): from .db_fpn import DBFPN, LKPAN, RSEFPN, RepLKFPN, RepLKPAN from .rnn import SequenceEncoder - support_dict = [ - "FPN", - "DBFPN", - "EASTFPN", - "SASTFPN", - "SequenceEncoder", - "PGFPN", - "TableFPN", - "RSEFPN", - "LKPAN", - "FCEFPN", - "RepLKFPN", - "RepLKPAN", - ] + support_dict = { + "DBFPN": DBFPN, + "SequenceEncoder": SequenceEncoder, + "RSEFPN": RSEFPN, + "LKPAN": LKPAN, + "RepLKFPN": RepLKFPN, + "RepLKPAN": RepLKPAN, + } module_name = config.pop("name") assert module_name in support_dict, Exception( "neck only support {}".format(support_dict) ) - module_class = eval(module_name)(**config) + module_class = support_dict[module_name](**config) return module_class diff --git a/python/rapidocr/inference_engine/pytorch/networks/necks/db_fpn.py b/python/rapidocr/inference_engine/pytorch/networks/necks/db_fpn.py index 9fcf01c..512ea13 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/necks/db_fpn.py +++ b/python/rapidocr/inference_engine/pytorch/networks/necks/db_fpn.py @@ -24,8 +24,9 @@ class DSConv(nn.Module): **kwargs, ): super(DSConv, self).__init__() - if groups == None: + if groups is None: groups = in_channels + self.if_act = if_act self.act = act self.conv1 = nn.Conv2d(