use paddlepaddle license

ppocr/losses/det_pse_loss.py

-# -*- coding: utf-8 -*-
-# @Time    : 3/29/19 11:03 AM
-# @Author  : zhoujun
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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 paddle
 from paddle import nn
 from paddle.nn import functional as F
@@ -9,7 +20,12 @@ from ppocr.utils.iou import iou
 
 
 class PSELoss(nn.Layer):
-    def __init__(self, alpha, ohem_ratio=3, kernel_sample_mask='pred', reduction='sum', **kwargs):
+    def __init__(self,
+                 alpha,
+                 ohem_ratio=3,
+                 kernel_sample_mask='pred',
+                 reduction='sum',
+                 **kwargs):
         """Implement PSE Loss.
         """
         super(PSELoss, self).__init__()
@@ -31,32 +47,32 @@ class PSELoss(nn.Layer):
         selected_masks = self.ohem_batch(texts, gt_texts, training_masks)
 
         loss_text = self.dice_loss(texts, gt_texts, selected_masks)
-        iou_text = iou((texts > 0).astype('int64'), gt_texts, training_masks, reduce=False)
-        losses = dict(
-            loss_text=loss_text,
-            iou_text=iou_text
-        )
+        iou_text = iou((texts > 0).astype('int64'),
+                       gt_texts,
+                       training_masks,
+                       reduce=False)
+        losses = dict(loss_text=loss_text, iou_text=iou_text)
 
         # kernel loss
         loss_kernels = []
         if self.kernel_sample_mask == 'gt':
             selected_masks = gt_texts * training_masks
         elif self.kernel_sample_mask == 'pred':
-            selected_masks = (F.sigmoid(texts) > 0.5).astype('float32') * training_masks
+            selected_masks = (
+                F.sigmoid(texts) > 0.5).astype('float32') * training_masks
 
         for i in range(kernels.shape[1]):
             kernel_i = kernels[:, i, :, :]
             gt_kernel_i = gt_kernels[:, i, :, :]
-            loss_kernel_i = self.dice_loss(kernel_i, gt_kernel_i, selected_masks)
+            loss_kernel_i = self.dice_loss(kernel_i, gt_kernel_i,
+                                           selected_masks)
             loss_kernels.append(loss_kernel_i)
         loss_kernels = paddle.mean(paddle.stack(loss_kernels, axis=1), axis=1)
-        iou_kernel = iou(
-            (kernels[:, -1, :, :] > 0).astype('int64'), gt_kernels[:, -1, :, :], training_masks * gt_texts,
-            reduce=False)
-        losses.update(dict(
-            loss_kernels=loss_kernels,
-            iou_kernel=iou_kernel
-        ))
+        iou_kernel = iou((kernels[:, -1, :, :] > 0).astype('int64'),
+                         gt_kernels[:, -1, :, :],
+                         training_masks * gt_texts,
+                         reduce=False)
+        losses.update(dict(loss_kernels=loss_kernels, iou_kernel=iou_kernel))
         loss = self.alpha * loss_text + (1 - self.alpha) * loss_kernels
         losses['loss'] = loss
         if self.reduction == 'sum':
@@ -83,13 +99,16 @@ class PSELoss(nn.Layer):
 
     def ohem_single(self, score, gt_text, training_mask, ohem_ratio=3):
         pos_num = int(paddle.sum((gt_text > 0.5).astype('float32'))) - int(
-            paddle.sum(paddle.logical_and((gt_text > 0.5), (training_mask <= 0.5)).astype('float32')))
+            paddle.sum(
+                paddle.logical_and((gt_text > 0.5), (training_mask <= 0.5))
+                .astype('float32')))
 
         if pos_num == 0:
             # selected_mask = gt_text.copy() * 0 # may be not good
             selected_mask = training_mask
-            selected_mask = selected_mask.reshape([1, selected_mask.shape[0], selected_mask.shape[1]]).astype(
-                'float32')
+            selected_mask = selected_mask.reshape(
+                [1, selected_mask.shape[0], selected_mask.shape[1]]).astype(
+                    'float32')
             return selected_mask
 
         neg_num = int(paddle.sum((gt_text <= 0.5).astype('float32')))
@@ -97,23 +116,29 @@ class PSELoss(nn.Layer):
 
         if neg_num == 0:
             selected_mask = training_mask
-            selected_mask = selected_mask.view(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
+            selected_mask = selected_mask.view(
+                1, selected_mask.shape[0],
+                selected_mask.shape[1]).astype('float32')
             return selected_mask
 
         neg_score = paddle.masked_select(score, gt_text <= 0.5)
         neg_score_sorted = paddle.sort(-neg_score)
         threshold = -neg_score_sorted[neg_num - 1]
 
-        selected_mask = paddle.logical_and(paddle.logical_or((score >= threshold), (gt_text > 0.5)),
-                                           (training_mask > 0.5))
-        selected_mask = selected_mask.reshape([1, selected_mask.shape[0], selected_mask.shape[1]]).astype('float32')
+        selected_mask = paddle.logical_and(
+            paddle.logical_or((score >= threshold), (gt_text > 0.5)),
+            (training_mask > 0.5))
+        selected_mask = selected_mask.reshape(
+            [1, selected_mask.shape[0], selected_mask.shape[1]]).astype(
+                'float32')
         return selected_mask
 
     def ohem_batch(self, scores, gt_texts, training_masks, ohem_ratio=3):
         selected_masks = []
         for i in range(scores.shape[0]):
             selected_masks.append(
-                self.ohem_single(scores[i, :, :], gt_texts[i, :, :], training_masks[i, :, :], ohem_ratio))
+                self.ohem_single(scores[i, :, :], gt_texts[i, :, :],
+                                 training_masks[i, :, :], ohem_ratio))
 
         selected_masks = paddle.concat(selected_masks, 0).astype('float32')
-        return selected_masks
\ No newline at end of file
+        return selected_masks