e904b667c6
Build/Publish Develop Docs / deploy (push) Failing after 1s
PaddleOCR Code Style Check / check-code-style (push) Failing after 1s
PaddleOCR PR Tests GPU / detect-changes (push) Failing after 1s
PaddleOCR PR Tests / detect-changes (push) Failing after 1s
PaddleOCR PR Tests GPU / test-pr-gpu (push) Has been cancelled
PaddleOCR PR Tests / test-pr (push) Has been cancelled
PaddleOCR PR Tests GPU / test-pr-gpu-impl (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.13) (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.8) (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.9) (push) Has been cancelled
285 lines
11 KiB
Python
285 lines
11 KiB
Python
# copyright (c) 2019 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.
|
||
"""
|
||
This code is refer from:
|
||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/models/losses/basic_loss.py
|
||
"""
|
||
from __future__ import absolute_import
|
||
from __future__ import division
|
||
from __future__ import print_function
|
||
|
||
import numpy as np
|
||
|
||
import paddle
|
||
from paddle import nn
|
||
import paddle.nn.functional as F
|
||
|
||
|
||
class BalanceLoss(nn.Layer):
|
||
def __init__(
|
||
self,
|
||
balance_loss=True,
|
||
main_loss_type="DiceLoss",
|
||
negative_ratio=3,
|
||
return_origin=False,
|
||
eps=1e-6,
|
||
**kwargs,
|
||
):
|
||
"""
|
||
The BalanceLoss for Differentiable Binarization text detection
|
||
args:
|
||
balance_loss (bool): whether balance loss or not, default is True
|
||
main_loss_type (str): can only be one of ['CrossEntropy','DiceLoss',
|
||
'Euclidean','BCELoss', 'MaskL1Loss'], default is 'DiceLoss'.
|
||
negative_ratio (int|float): float, default is 3.
|
||
return_origin (bool): whether return unbalanced loss or not, default is False.
|
||
eps (float): default is 1e-6.
|
||
"""
|
||
super(BalanceLoss, self).__init__()
|
||
self.balance_loss = balance_loss
|
||
self.main_loss_type = main_loss_type
|
||
self.negative_ratio = negative_ratio
|
||
self.return_origin = return_origin
|
||
self.eps = eps
|
||
|
||
if self.main_loss_type == "CrossEntropy":
|
||
self.loss = nn.CrossEntropyLoss()
|
||
elif self.main_loss_type == "Euclidean":
|
||
self.loss = nn.MSELoss()
|
||
elif self.main_loss_type == "DiceLoss":
|
||
self.loss = DiceLoss(self.eps)
|
||
elif self.main_loss_type == "BCELoss":
|
||
self.loss = BCELoss(reduction="none")
|
||
elif self.main_loss_type == "MaskL1Loss":
|
||
self.loss = MaskL1Loss(self.eps)
|
||
else:
|
||
loss_type = [
|
||
"CrossEntropy",
|
||
"DiceLoss",
|
||
"Euclidean",
|
||
"BCELoss",
|
||
"MaskL1Loss",
|
||
]
|
||
raise Exception(
|
||
"main_loss_type in BalanceLoss() can only be one of {}".format(
|
||
loss_type
|
||
)
|
||
)
|
||
|
||
def forward(self, pred, gt, mask=None):
|
||
"""
|
||
The BalanceLoss for Differentiable Binarization text detection
|
||
args:
|
||
pred (variable): predicted feature maps.
|
||
gt (variable): ground truth feature maps.
|
||
mask (variable): masked maps.
|
||
return: (variable) balanced loss
|
||
"""
|
||
positive = gt * mask
|
||
negative = (1 - gt) * mask
|
||
|
||
positive_count = int(positive.sum())
|
||
negative_count = int(min(negative.sum(), positive_count * self.negative_ratio))
|
||
loss = self.loss(pred, gt, mask=mask)
|
||
|
||
if not self.balance_loss:
|
||
return loss
|
||
|
||
positive_loss = positive * loss
|
||
negative_loss = negative * loss
|
||
negative_loss = paddle.reshape(negative_loss, shape=[-1])
|
||
if negative_count > 0:
|
||
sort_loss = negative_loss.sort(descending=True)
|
||
negative_loss = sort_loss[:negative_count]
|
||
# negative_loss, _ = paddle.topk(negative_loss, k=negative_count_int)
|
||
balance_loss = (positive_loss.sum() + negative_loss.sum()) / (
|
||
positive_count + negative_count + self.eps
|
||
)
|
||
else:
|
||
balance_loss = positive_loss.sum() / (positive_count + self.eps)
|
||
if self.return_origin:
|
||
return balance_loss, loss
|
||
|
||
return balance_loss
|
||
|
||
|
||
class DiceLoss(nn.Layer):
|
||
def __init__(self, eps=1e-6):
|
||
super(DiceLoss, self).__init__()
|
||
self.eps = eps
|
||
|
||
def forward(self, pred, gt, mask, weights=None):
|
||
"""
|
||
DiceLoss function.
|
||
"""
|
||
|
||
assert pred.shape == gt.shape
|
||
assert pred.shape == mask.shape
|
||
if weights is not None:
|
||
assert weights.shape == mask.shape
|
||
mask = weights * mask
|
||
intersection = paddle.sum(pred * gt * mask)
|
||
|
||
union = paddle.sum(pred * mask) + paddle.sum(gt * mask) + 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, gt, mask):
|
||
"""
|
||
Mask L1 Loss
|
||
"""
|
||
loss = (paddle.abs(pred - gt) * mask).sum() / (mask.sum() + self.eps)
|
||
loss = paddle.mean(loss)
|
||
return loss
|
||
|
||
|
||
class BCELoss(nn.Layer):
|
||
def __init__(self, reduction="mean"):
|
||
super(BCELoss, self).__init__()
|
||
self.reduction = reduction
|
||
|
||
def forward(self, input, label, mask=None, weight=None, name=None):
|
||
loss = F.binary_cross_entropy(input, label, reduction=self.reduction)
|
||
return loss
|
||
|
||
|
||
class MaskedFocalLoss(nn.Layer):
|
||
"""
|
||
Binary Focal Loss with mask support, designed for text segmentation tasks.
|
||
|
||
Focal Loss addresses class imbalance by down-weighting easy examples and
|
||
focusing training on hard examples:
|
||
FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)
|
||
|
||
Compared to OHEM (which hard-selects a fixed ratio of negatives), Focal Loss
|
||
applies a continuous per-pixel weight that gracefully scales with difficulty,
|
||
making it a strictly superior drop-in for the OHEM + DiceLoss pattern when
|
||
DiceLoss returns a scalar and OHEM has no discriminating effect.
|
||
|
||
Args:
|
||
alpha (float): Balancing factor for the positive (text) class.
|
||
Since text pixels are a small minority, alpha > 0.5 gives them
|
||
higher weight. Default: 0.75.
|
||
gamma (float): Focusing parameter. gamma=0 reduces to masked BCE.
|
||
gamma=2 is the standard value from the original Focal Loss paper.
|
||
Default: 2.0.
|
||
eps (float): Small constant for numerical stability. Default: 1e-6.
|
||
"""
|
||
|
||
def __init__(self, alpha=0.25, gamma=2.0, eps=1e-6):
|
||
super(MaskedFocalLoss, self).__init__()
|
||
self.alpha = alpha
|
||
self.gamma = gamma
|
||
self.eps = eps
|
||
|
||
def forward(self, pred, gt, mask):
|
||
"""
|
||
Args:
|
||
pred (Tensor): Predicted probability map, shape (B, H, W), in [0, 1].
|
||
(i.e. after sigmoid — the direct output of DBHead.binarize)
|
||
gt (Tensor): Binary ground-truth map, shape (B, H, W), values 0 or 1.
|
||
mask (Tensor): Valid-pixel mask, shape (B, H, W), values 0 or 1.
|
||
Pixels with mask=0 are ignored regions (e.g. too-small text).
|
||
Returns:
|
||
Tensor: Scalar focal loss averaged over valid (mask=1) pixels.
|
||
"""
|
||
# F.sigmoid_focal_loss expects a logit (pre-sigmoid) input and applies
|
||
# sigmoid internally using the numerically stable log-sum-exp form:
|
||
# log(σ(x)) = -softplus(-x), log(1-σ(x)) = -softplus(x)
|
||
# This avoids the log(0) issue of the manual implementation.
|
||
# Since pred is already a probability (post-sigmoid from DBHead), we
|
||
# invert it: logit = log(p / (1-p)). The round-trip is numerically safe
|
||
# after clamping, and the stable path inside paddle takes over from there.
|
||
pred = paddle.clip(pred, self.eps, 1.0 - self.eps)
|
||
logit = paddle.log(pred / (1.0 - pred))
|
||
|
||
# Per-pixel focal loss, shape (B, H, W)
|
||
# reduction='none' so we can apply the mask ourselves
|
||
loss = F.sigmoid_focal_loss(
|
||
logit,
|
||
gt,
|
||
normalizer=None,
|
||
alpha=self.alpha,
|
||
gamma=self.gamma,
|
||
reduction="none",
|
||
)
|
||
|
||
# Average over valid (mask=1) pixels only
|
||
return (loss * mask).sum() / (mask.sum() + self.eps)
|
||
|
||
|
||
class DiceFocalLoss(nn.Layer):
|
||
"""
|
||
Combined DiceLoss + MaskedFocalLoss for binary text segmentation.
|
||
|
||
Rationale for the combination:
|
||
- DiceLoss optimizes the global F1 / region overlap between prediction and GT.
|
||
It is naturally robust to class imbalance (text vs background) because it
|
||
normalizes by the sum of both sets, not by pixel count.
|
||
- MaskedFocalLoss provides per-pixel supervision with adaptive hard-example
|
||
weighting. It compensates for DiceLoss being a global metric that cannot
|
||
distinguish which specific pixels are mispredicted.
|
||
Together they provide complementary supervision: DiceLoss for global shape
|
||
quality, FocalLoss for pixel-level precision on ambiguous boundaries.
|
||
|
||
This design follows the Dice + Focal combination used in mmsegmentation and
|
||
segmentation_models_pytorch for binary segmentation with class imbalance.
|
||
|
||
This class is a drop-in replacement for BalanceLoss when main_loss_type is
|
||
'DiceLoss' — both share the same forward(pred, gt, mask) signature and
|
||
return a scalar.
|
||
|
||
Args:
|
||
dice_weight (float): Weight for the DiceLoss term. Default: 1.0.
|
||
focal_weight (float): Weight for the MaskedFocalLoss term. Default: 1.0.
|
||
focal_alpha (float): Positive-class balancing factor for FocalLoss.
|
||
Default: 0.75.
|
||
focal_gamma (float): Focusing exponent for FocalLoss. Default: 2.0.
|
||
eps (float): Numerical stability constant. Default: 1e-6.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
dice_weight=1.0,
|
||
focal_weight=1.0,
|
||
focal_alpha=0.75,
|
||
focal_gamma=2.0,
|
||
eps=1e-6,
|
||
):
|
||
super(DiceFocalLoss, self).__init__()
|
||
self.dice_weight = dice_weight
|
||
self.focal_weight = focal_weight
|
||
self.dice_loss = DiceLoss(eps=eps)
|
||
self.focal_loss = MaskedFocalLoss(alpha=focal_alpha, gamma=focal_gamma, eps=eps)
|
||
|
||
def forward(self, pred, gt, mask=None):
|
||
"""
|
||
Args:
|
||
pred (Tensor): Predicted probability map, shape (B, H, W), in [0, 1].
|
||
gt (Tensor): Binary ground-truth shrink map, shape (B, H, W).
|
||
mask (Tensor): Valid-pixel mask, shape (B, H, W).
|
||
Returns:
|
||
Tensor: Scalar combined loss.
|
||
"""
|
||
loss_dice = self.dice_loss(pred, gt, mask)
|
||
loss_focal = self.focal_loss(pred, gt, mask)
|
||
return self.dice_weight * loss_dice + self.focal_weight * loss_focal
|