first commit

benchmark/PaddleOCR_DBNet/models/losses/DB_loss.py

@@ -0,0 +1,50 @@
+import paddle
+from models.losses.basic_loss import BalanceCrossEntropyLoss, MaskL1Loss, DiceLoss
+
+
+class DBLoss(paddle.nn.Layer):
+    def __init__(self, alpha=1.0, beta=10, ohem_ratio=3, reduction="mean", eps=1e-06):
+        """
+        Implement PSE Loss.
+        :param alpha: binary_map loss 前面的系数
+        :param beta: threshold_map loss 前面的系数
+        :param ohem_ratio: OHEM的比例
+        :param reduction: 'mean' or 'sum'对 batch里的loss 算均值或求和
+        """
+        super().__init__()
+        assert reduction in ["mean", "sum"], " reduction must in ['mean','sum']"
+        self.alpha = alpha
+        self.beta = beta
+        self.bce_loss = BalanceCrossEntropyLoss(negative_ratio=ohem_ratio)
+        self.dice_loss = DiceLoss(eps=eps)
+        self.l1_loss = MaskL1Loss(eps=eps)
+        self.ohem_ratio = ohem_ratio
+        self.reduction = reduction
+
+    def forward(self, pred, batch):
+        shrink_maps = pred[:, 0, :, :]
+        threshold_maps = pred[:, 1, :, :]
+        binary_maps = pred[:, 2, :, :]
+        loss_shrink_maps = self.bce_loss(
+            shrink_maps, batch["shrink_map"], batch["shrink_mask"]
+        )
+        loss_threshold_maps = self.l1_loss(
+            threshold_maps, batch["threshold_map"], batch["threshold_mask"]
+        )
+        metrics = dict(
+            loss_shrink_maps=loss_shrink_maps, loss_threshold_maps=loss_threshold_maps
+        )
+        if pred.shape[1] > 2:
+            loss_binary_maps = self.dice_loss(
+                binary_maps, batch["shrink_map"], batch["shrink_mask"]
+            )
+            metrics["loss_binary_maps"] = loss_binary_maps
+            loss_all = (
+                self.alpha * loss_shrink_maps
+                + self.beta * loss_threshold_maps
+                + loss_binary_maps
+            )
+            metrics["loss"] = loss_all
+        else:
+            metrics["loss"] = loss_shrink_maps
+        return metrics

benchmark/PaddleOCR_DBNet/models/losses/__init__.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2020/6/5 11:36
+# @Author  : zhoujun
+import copy
+from .DB_loss import DBLoss
+
+__all__ = ["build_loss"]
+support_loss = ["DBLoss"]
+
+
+def build_loss(config):
+    copy_config = copy.deepcopy(config)
+    loss_type = copy_config.pop("type")
+    assert loss_type in support_loss, f"all support loss is {support_loss}"
+    criterion = eval(loss_type)(**copy_config)
+    return criterion

benchmark/PaddleOCR_DBNet/models/losses/basic_loss.py

@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2019/12/4 14:39
+# @Author  : zhoujun
+import paddle
+import paddle.nn as nn
+
+
+class BalanceCrossEntropyLoss(nn.Layer):
+    """
+    Balanced cross entropy loss.
+    Shape:
+        - Input: :math:`(N, 1, H, W)`
+        - GT: :math:`(N, 1, H, W)`, same shape as the input
+        - Mask: :math:`(N, H, W)`, same spatial shape as the input
+        - Output: scalar.
+
+    """
+
+    def __init__(self, negative_ratio=3.0, eps=1e-6):
+        super(BalanceCrossEntropyLoss, self).__init__()
+        self.negative_ratio = negative_ratio
+        self.eps = eps
+
+    def forward(
+        self,
+        pred: paddle.Tensor,
+        gt: paddle.Tensor,
+        mask: paddle.Tensor,
+        return_origin=False,
+    ):
+        """
+        Args:
+            pred: shape :math:`(N, 1, H, W)`, the prediction of network
+            gt: shape :math:`(N, 1, H, W)`, the target
+            mask: shape :math:`(N, H, W)`, the mask indicates positive regions
+        """
+        positive = gt * mask
+        negative = (1 - gt) * mask
+        positive_count = int(positive.sum())
+        negative_count = min(
+            int(negative.sum()), int(positive_count * self.negative_ratio)
+        )
+        loss = nn.functional.binary_cross_entropy(pred, gt, reduction="none")
+        positive_loss = loss * positive
+        negative_loss = loss * negative
+        negative_loss, _ = negative_loss.reshape([-1]).topk(negative_count)
+
+        balance_loss = (positive_loss.sum() + negative_loss.sum()) / (
+            positive_count + negative_count + self.eps
+        )
+
+        if return_origin:
+            return balance_loss, loss
+        return balance_loss
+
+
+class DiceLoss(nn.Layer):
+    """
+    Loss function from https://arxiv.org/abs/1707.03237,
+    where iou computation is introduced heatmap manner to measure the
+    diversity between tow heatmaps.
+    """
+
+    def __init__(self, eps=1e-6):
+        super(DiceLoss, self).__init__()
+        self.eps = eps
+
+    def forward(self, pred: paddle.Tensor, gt, mask, weights=None):
+        """
+        pred: one or two heatmaps of shape (N, 1, H, W),
+            the losses of tow heatmaps are added together.
+        gt: (N, 1, H, W)
+        mask: (N, H, W)
+        """
+        return self._compute(pred, gt, mask, weights)
+
+    def _compute(self, pred, gt, mask, weights):
+        if len(pred.shape) == 4:
+            pred = pred[:, 0, :, :]
+            gt = gt[:, 0, :, :]
+        assert pred.shape == gt.shape
+        assert pred.shape == mask.shape
+        if weights is not None:
+            assert weights.shape == mask.shape
+            mask = weights * mask
+        intersection = (pred * gt * mask).sum()
+
+        union = (pred * mask).sum() + (gt * mask).sum() + self.eps
+        loss = 1 - 2.0 * intersection / union
+        assert loss <= 1
+        return loss
+
+
+class MaskL1Loss(nn.Layer):
+    def __init__(self, eps=1e-6):
+        super(MaskL1Loss, self).__init__()
+        self.eps = eps
+
+    def forward(self, pred: paddle.Tensor, gt, mask):
+        loss = (paddle.abs(pred - gt) * mask).sum() / (mask.sum() + self.eps)
+        return loss