69 lines
2.5 KiB
Python
69 lines
2.5 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.
|
||
|
|
|
||
|
|
from __future__ import absolute_import
|
||
|
|
from __future__ import division
|
||
|
|
from __future__ import print_function
|
||
|
|
|
||
|
|
from .det_basic_loss import BalanceLoss, MaskL1Loss, DiceLoss
|
||
|
|
|
||
|
|
|
||
|
|
class DBLoss(object):
|
||
|
|
"""
|
||
|
|
Differentiable Binarization (DB) Loss Function
|
||
|
|
args:
|
||
|
|
param (dict): the super paramter for DB Loss
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, params):
|
||
|
|
super(DBLoss, self).__init__()
|
||
|
|
self.balance_loss = params['balance_loss']
|
||
|
|
self.main_loss_type = params['main_loss_type']
|
||
|
|
|
||
|
|
self.alpha = params['alpha']
|
||
|
|
self.beta = params['beta']
|
||
|
|
self.ohem_ratio = params['ohem_ratio']
|
||
|
|
|
||
|
|
def __call__(self, predicts, labels):
|
||
|
|
label_shrink_map = labels['shrink_map']
|
||
|
|
label_shrink_mask = labels['shrink_mask']
|
||
|
|
label_threshold_map = labels['threshold_map']
|
||
|
|
label_threshold_mask = labels['threshold_mask']
|
||
|
|
pred = predicts['maps']
|
||
|
|
shrink_maps = pred[:, 0, :, :]
|
||
|
|
threshold_maps = pred[:, 1, :, :]
|
||
|
|
binary_maps = pred[:, 2, :, :]
|
||
|
|
|
||
|
|
loss_shrink_maps = BalanceLoss(
|
||
|
|
shrink_maps,
|
||
|
|
label_shrink_map,
|
||
|
|
label_shrink_mask,
|
||
|
|
balance_loss=self.balance_loss,
|
||
|
|
main_loss_type=self.main_loss_type,
|
||
|
|
negative_ratio=self.ohem_ratio)
|
||
|
|
loss_threshold_maps = MaskL1Loss(threshold_maps, label_threshold_map,
|
||
|
|
label_threshold_mask)
|
||
|
|
loss_binary_maps = DiceLoss(binary_maps, label_shrink_map,
|
||
|
|
label_shrink_mask)
|
||
|
|
loss_shrink_maps = self.alpha * loss_shrink_maps
|
||
|
|
loss_threshold_maps = self.beta * loss_threshold_maps
|
||
|
|
|
||
|
|
loss_all = loss_shrink_maps + loss_threshold_maps\
|
||
|
|
+ loss_binary_maps
|
||
|
|
losses = {'total_loss':loss_all,\
|
||
|
|
"loss_shrink_maps":loss_shrink_maps,\
|
||
|
|
"loss_threshold_maps":loss_threshold_maps,\
|
||
|
|
"loss_binary_maps":loss_binary_maps}
|
||
|
|
return losses
|