mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 13:03:58 +08:00
fix: align the performance of the torch version of the PP-OCRv6 Rec tiny model
This commit is contained in:
SWHL committed
1 parent
7fe716f8e3
commit
b64a46ffd4
7 files changed
+93
-64
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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(
|
||||
|
||||
Reference in new issue
Block a user