first commit

ppocr/losses/table_master_loss.py

@@ -0,0 +1,74 @@
+# 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.
+"""
+This code is refer from:
+https://github.com/JiaquanYe/TableMASTER-mmocr/tree/master/mmocr/models/textrecog/losses
+"""
+
+import paddle
+from paddle import nn
+
+
+class TableMasterLoss(nn.Layer):
+    def __init__(self, ignore_index=-1):
+        super(TableMasterLoss, self).__init__()
+        self.structure_loss = nn.CrossEntropyLoss(
+            ignore_index=ignore_index, reduction="mean"
+        )
+        self.box_loss = nn.L1Loss(reduction="sum")
+        self.eps = 1e-12
+
+    def forward(self, predicts, batch):
+        # structure_loss
+        structure_probs = predicts["structure_probs"]
+        structure_targets = batch[1]
+        structure_targets = structure_targets[:, 1:]
+        structure_probs = structure_probs.reshape([-1, structure_probs.shape[-1]])
+        structure_targets = structure_targets.reshape([-1])
+
+        structure_loss = self.structure_loss(structure_probs, structure_targets)
+        structure_loss = structure_loss.mean()
+        losses = dict(structure_loss=structure_loss)
+
+        # box loss
+        bboxes_preds = predicts["loc_preds"]
+        bboxes_targets = batch[2][:, 1:, :]
+        bbox_masks = batch[3][:, 1:]
+        # mask empty-bbox or non-bbox structure token's bbox.
+
+        masked_bboxes_preds = bboxes_preds * bbox_masks
+        masked_bboxes_targets = bboxes_targets * bbox_masks
+
+        # horizon loss (x and width)
+        horizon_sum_loss = self.box_loss(
+            masked_bboxes_preds[:, :, 0::2], masked_bboxes_targets[:, :, 0::2]
+        )
+        horizon_loss = horizon_sum_loss / (bbox_masks.sum() + self.eps)
+        # vertical loss (y and height)
+        vertical_sum_loss = self.box_loss(
+            masked_bboxes_preds[:, :, 1::2], masked_bboxes_targets[:, :, 1::2]
+        )
+        vertical_loss = vertical_sum_loss / (bbox_masks.sum() + self.eps)
+
+        horizon_loss = horizon_loss.mean()
+        vertical_loss = vertical_loss.mean()
+        all_loss = structure_loss + horizon_loss + vertical_loss
+        losses.update(
+            {
+                "loss": all_loss,
+                "horizon_bbox_loss": horizon_loss,
+                "vertical_bbox_loss": vertical_loss,
+            }
+        )
+        return losses