Spaces:
Runtime error
Runtime error
| # 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 | |