| 18 | |
| 19 | |
| 20 | def structure_loss(pred, mask, w=1): |
| 21 | weit = 1 + 5 * torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask) |
| 22 | wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='none') |
| 23 | wbce = (weit * wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3)) |
| 24 | |
| 25 | pred = torch.sigmoid(pred) |
| 26 | inter = ((pred * mask) * weit).sum(dim=(2, 3)) |
| 27 | union = ((pred + mask) * weit).sum(dim=(2, 3)) |
| 28 | wiou = 1 - (inter + 1) / (union - inter + 1) |
| 29 | |
| 30 | return (w * (wbce + wiou)).mean() |
| 31 | |
| 32 | def dice_coefficient(predicted, labels): |
| 33 | if predicted.device != labels.device: |