first commit

ppocr/modeling/architectures/__init__.py

@@ -0,0 +1,145 @@
+# 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
+import importlib
+
+from paddle.jit import to_static
+from paddle.static import InputSpec
+
+from .base_model import BaseModel
+from .distillation_model import DistillationModel
+
+__all__ = ["build_model", "apply_to_static"]
+
+
+def build_model(config):
+    config = copy.deepcopy(config)
+    if not "name" in config:
+        arch = BaseModel(config)
+    else:
+        name = config.pop("name")
+        mod = importlib.import_module(__name__)
+        arch = getattr(mod, name)(config)
+    return arch
+
+
+def apply_to_static(model, config, logger):
+    if config["Global"].get("to_static", False) is not True:
+        return model
+    assert (
+        "d2s_train_image_shape" in config["Global"]
+    ), "d2s_train_image_shape must be assigned for static training mode..."
+    supported_list = [
+        "DB",
+        "SVTR_LCNet",
+        "TableMaster",
+        "LayoutXLM",
+        "SLANet",
+        "SVTR",
+        "SVTR_HGNet",
+        "LaTeXOCR",
+        "UniMERNet",
+        "PP-FormulaNet-S",
+        "PP-FormulaNet-L",
+    ]
+    if config["Architecture"]["algorithm"] in ["Distillation"]:
+        algo = list(config["Architecture"]["Models"].values())[0]["algorithm"]
+    else:
+        algo = config["Architecture"]["algorithm"]
+    assert (
+        algo in supported_list
+    ), f"algorithms that supports static training must in in {supported_list} but got {algo}"
+
+    specs = [
+        InputSpec([None] + config["Global"]["d2s_train_image_shape"], dtype="float32")
+    ]
+
+    if algo == "SVTR_LCNet":
+        specs.append(
+            [
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec([None], dtype="int64"),
+                InputSpec([None], dtype="float64"),
+            ]
+        )
+    elif algo == "TableMaster":
+        specs.append(
+            [
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"], 4], dtype="float32"
+                ),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"], 1], dtype="float32"
+                ),
+                InputSpec([None, 6], dtype="float32"),
+            ]
+        )
+    elif algo == "LayoutXLM":
+        specs = [
+            [
+                InputSpec(shape=[None, 512], dtype="int64"),  # input_ids
+                InputSpec(shape=[None, 512, 4], dtype="int64"),  # bbox
+                InputSpec(shape=[None, 512], dtype="int64"),  # attention_mask
+                InputSpec(shape=[None, 512], dtype="int64"),  # token_type_ids
+                InputSpec(shape=[None, 3, 224, 224], dtype="float32"),  # image
+                InputSpec(shape=[None, 512], dtype="int64"),  # label
+            ]
+        ]
+    elif algo == "SLANet":
+        specs.append(
+            [
+                InputSpec(
+                    [None, config["Global"]["max_text_length"] + 2], dtype="int64"
+                ),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"] + 2, 4], dtype="float32"
+                ),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"] + 2, 1], dtype="float32"
+                ),
+                InputSpec([None], dtype="int64"),
+                InputSpec([None, 6], dtype="float64"),
+            ]
+        )
+    elif algo == "SVTR":
+        specs.append(
+            [
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec([None], dtype="int64"),
+            ]
+        )
+    elif algo == "LaTeXOCR":
+        specs = [
+            [
+                InputSpec(shape=[None, 1, None, None], dtype="float32"),
+                InputSpec(shape=[None, None], dtype="float32"),
+                InputSpec(shape=[None, None], dtype="float32"),
+            ]
+        ]
+    elif algo in ["UniMERNet", "PP-FormulaNet-S", "PP-FormulaNet-L"]:
+        specs = [
+            [
+                InputSpec(
+                    [None] + config["Global"]["d2s_train_image_shape"], dtype="float32"
+                ),
+                InputSpec(shape=[None, None], dtype="float32"),
+                InputSpec(shape=[None, None], dtype="float32"),
+            ]
+        ]
+    model = to_static(model, input_spec=specs)
+    logger.info("Successfully to apply @to_static with specs: {}".format(specs))
+    return model

ppocr/modeling/architectures/base_model.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2021 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.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+from ppocr.modeling.transforms import build_transform
+from ppocr.modeling.backbones import build_backbone
+from ppocr.modeling.necks import build_neck
+from ppocr.modeling.heads import build_head
+
+__all__ = ["BaseModel"]
+
+
+class BaseModel(nn.Layer):
+    def __init__(self, config):
+        """
+        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 transform,
+        # for rec, transform can be TPS,None
+        # for det and cls, transform should to be None,
+        # if you make model differently, you can use transform in det and cls
+        if "Transform" not in config or config["Transform"] is None:
+            self.use_transform = False
+        else:
+            self.use_transform = True
+            config["Transform"]["in_channels"] = in_channels
+            self.transform = build_transform(config["Transform"])
+            in_channels = self.transform.out_channels
+
+        # 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"])
+
+        self.return_all_feats = config.get("return_all_feats", False)
+
+    def forward(self, x, data=None):
+        y = dict()
+        if self.use_transform:
+            x = self.transform(x)
+        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, targets=data)
+            # for multi head, save ctc neck out for udml
+            if isinstance(x, dict) and "ctc_neck" 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
+            final_name = "head_out"
+        if self.return_all_feats:
+            if self.training:
+                return y
+            elif isinstance(x, dict):
+                return x
+            else:
+                return {final_name: x}
+        else:
+            return x

ppocr/modeling/architectures/distillation_model.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2021 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.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+from ppocr.modeling.transforms import build_transform
+from ppocr.modeling.backbones import build_backbone
+from ppocr.modeling.necks import build_neck
+from ppocr.modeling.heads import build_head
+from .base_model import BaseModel
+from ppocr.utils.save_load import load_pretrained_params
+
+__all__ = ["DistillationModel"]
+
+
+class DistillationModel(nn.Layer):
+    def __init__(self, config):
+        """
+        the module for OCR distillation.
+        args:
+            config (dict): the super parameters for module.
+        """
+        super().__init__()
+        self.model_list = []
+        self.model_name_list = []
+        for key in config["Models"]:
+            model_config = config["Models"][key]
+            freeze_params = False
+            pretrained = None
+            if "freeze_params" in model_config:
+                freeze_params = model_config.pop("freeze_params")
+            if "pretrained" in model_config:
+                pretrained = model_config.pop("pretrained")
+            model = BaseModel(model_config)
+            if pretrained is not None:
+                load_pretrained_params(model, pretrained)
+            if freeze_params:
+                for param in model.parameters():
+                    param.trainable = False
+            self.model_list.append(self.add_sublayer(key, model))
+            self.model_name_list.append(key)
+
+    def forward(self, x, data=None):
+        result_dict = dict()
+        for idx, model_name in enumerate(self.model_name_list):
+            result_dict[model_name] = self.model_list[idx](x, data)
+        return result_dict