Spaces:
Runtime error
Runtime error
File size: 4,324 Bytes
01bb3bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
# -*- coding: utf-8 -*-
# @Time : 10/1/21
# @Author : GXYM
import torch
from torch import nn
import numpy as np
class SegmentLoss(nn.Module):
def __init__(self, Lambda, ratio=3, reduction='mean'):
"""Implement PSE Loss.
"""
super(SegmentLoss, self).__init__()
assert reduction in ['mean', 'sum'], " reduction must in ['mean','sum']"
self.Lambda = Lambda
self.ratio = ratio
self.reduction = reduction
def forward(self, outputs, labels, training_masks, th=0.5):
texts = outputs[:, -1, :, :]
kernels = outputs[:, :-1, :, :]
gt_texts = labels[:, -1, :, :]
gt_kernels = labels[:, :-1, :, :]
selected_masks = self.ohem_batch(texts, gt_texts, training_masks)
selected_masks = selected_masks.to(outputs.device)
loss_text = self.dice_loss(texts, gt_texts, selected_masks)
loss_kernels = []
# mask0 = torch.sigmoid(texts).data.cpu().numpy()
mask0 = texts.data.cpu().numpy()
mask1 = training_masks.data.cpu().numpy()
selected_masks = ((mask0 > th) & (mask1 > th)).astype('float32')
selected_masks = torch.from_numpy(selected_masks).float()
selected_masks = selected_masks.to(outputs.device)
kernels_num = gt_kernels.size()[1]
for i in range(kernels_num):
kernel_i = kernels[:, i, :, :]
gt_kernel_i = gt_kernels[:, i, :, :]
loss_kernel_i = self.dice_loss(kernel_i, gt_kernel_i, selected_masks)
loss_kernels.append(loss_kernel_i)
loss_kernels = torch.stack(loss_kernels).mean(0)
if self.reduction == 'mean':
loss_text = loss_text.mean()
loss_kernels = loss_kernels.mean()
elif self.reduction == 'sum':
loss_text = loss_text.sum()
loss_kernels = loss_kernels.sum()
loss = self.Lambda *loss_text + (1-self.Lambda)*loss_kernels
return loss_text, loss_kernels, loss
def dice_loss(self, input, target, mask):
# input = torch.sigmoid(input)
input = input.contiguous().view(input.size()[0], -1)
target = target.contiguous().view(target.size()[0], -1)
mask = mask.contiguous().view(mask.size()[0], -1)
input = input * mask
target = (target.float()) * mask
a = torch.sum(input * target, 1)
b = torch.sum(input * input, 1) + 0.001
c = torch.sum(target * target, 1) + 0.001
d = (2 * a) / (b + c)
return 1 - d
def ohem_single(self, score, gt_text, training_mask, th=0.5):
pos_num = (int)(np.sum(gt_text > th)) - (int)(np.sum((gt_text > th) & (training_mask <= th)))
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')
return selected_mask
neg_num = (int)(np.sum(gt_text <= th))
neg_num = (int)(min(pos_num * 3, neg_num))
if neg_num == 0:
selected_mask = training_mask
selected_mask = selected_mask.reshape(1, selected_mask.shape[0], selected_mask.shape[1]).astype('float32')
return selected_mask
neg_score = score[gt_text <= th]
# 将负样本得分从高到低排序
neg_score_sorted = np.sort(-neg_score)
threshold = -neg_score_sorted[neg_num - 1]
# 选出 得分高的 负样本 和正样本 的 mask
selected_mask = ((score >= threshold) | (gt_text > th)) & (training_mask > th)
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):
scores = scores.data.cpu().numpy()
gt_texts = gt_texts.data.cpu().numpy()
training_masks = training_masks.data.cpu().numpy()
selected_masks = []
for i in range(scores.shape[0]):
selected_masks.append(self.ohem_single(scores[i, :, :], gt_texts[i, :, :], training_masks[i, :, :]))
selected_masks = np.concatenate(selected_masks, 0)
selected_masks = torch.from_numpy(selected_masks).float()
return selected_masks
|