Active Contour Loss based on maxpooling & minpooling
| 54 | |
| 55 | |
| 56 | class ACLossV2(nn.Module): |
| 57 | """ |
| 58 | Active Contour Loss |
| 59 | based on maxpooling & minpooling |
| 60 | """ |
| 61 | |
| 62 | def __init__(self, miu=1.0, classes=3): |
| 63 | super(ACLossV2, self).__init__() |
| 64 | |
| 65 | self.miu = miu |
| 66 | self.classes = classes |
| 67 | |
| 68 | def forward(self, predication, label): |
| 69 | min_pool_x = nn.functional.max_pool2d( |
| 70 | predication * -1, (3, 3), 1, 1) * -1 |
| 71 | contour = torch.relu(nn.functional.max_pool2d( |
| 72 | min_pool_x, (3, 3), 1, 1) - min_pool_x) |
| 73 | |
| 74 | # length |
| 75 | length = torch.sum(torch.abs(contour)) |
| 76 | |
| 77 | # region |
| 78 | label = label.float() |
| 79 | c_in = torch.ones_like(predication) |
| 80 | c_out = torch.zeros_like(predication) |
| 81 | region_in = torch.abs(torch.sum(predication * ((label - c_in) ** 2))) |
| 82 | region_out = torch.abs( |
| 83 | torch.sum((1 - predication) * ((label - c_out) ** 2))) |
| 84 | region = self.miu * region_in + region_out |
| 85 | |
| 86 | return region + length |
| 87 | |
| 88 | |
| 89 | def ACELoss(y_pred, y_true, u=1, a=1, b=1): |