| 87 | |
| 88 | |
| 89 | class StructureLoss(torch.nn.Module): |
| 90 | def __init__(self): |
| 91 | super(StructureLoss, self).__init__() |
| 92 | |
| 93 | def forward(self, pred, target): |
| 94 | weit = 1+5*torch.abs(F.avg_pool2d(target, kernel_size=31, stride=1, padding=15)-target) |
| 95 | wbce = F.binary_cross_entropy_with_logits(pred, target, reduction='none') |
| 96 | wbce = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3)) |
| 97 | |
| 98 | pred = torch.sigmoid(pred) |
| 99 | inter = ((pred * target) * weit).sum(dim=(2, 3)) |
| 100 | union = ((pred + target) * weit).sum(dim=(2, 3)) |
| 101 | wiou = 1-(inter+1)/(union-inter+1) |
| 102 | |
| 103 | return (wbce+wiou).mean() |
| 104 | |
| 105 | |
| 106 | class PatchIoULoss(torch.nn.Module): |