| 69 | |
| 70 | |
| 71 | class IoULoss(torch.nn.Module): |
| 72 | def __init__(self): |
| 73 | super(IoULoss, self).__init__() |
| 74 | |
| 75 | def forward(self, pred, target): |
| 76 | b = pred.shape[0] |
| 77 | IoU = 0.0 |
| 78 | for i in range(0, b): |
| 79 | # compute the IoU of the foreground |
| 80 | Iand1 = torch.sum(target[i, :, :, :] * pred[i, :, :, :]) |
| 81 | Ior1 = torch.sum(target[i, :, :, :]) + torch.sum(pred[i, :, :, :]) - Iand1 |
| 82 | IoU1 = Iand1 / Ior1 |
| 83 | # IoU loss is (1-IoU1) |
| 84 | IoU = IoU + (1-IoU1) |
| 85 | # return IoU/b |
| 86 | return IoU |
| 87 | |
| 88 | |
| 89 | class StructureLoss(torch.nn.Module): |