Active Contour Loss based on sobel filter
| 165 | |
| 166 | |
| 167 | class ACLoss3D(nn.Module): |
| 168 | """ |
| 169 | Active Contour Loss |
| 170 | based on sobel filter |
| 171 | """ |
| 172 | |
| 173 | def __init__(self, classes=4, alpha=1): |
| 174 | super(ACLoss3D, self).__init__() |
| 175 | self.alpha = alpha |
| 176 | self.classes = classes |
| 177 | sobel = np.array([[[1., 2., 1.], |
| 178 | [2., 4., 2.], |
| 179 | [1., 2., 1.]], |
| 180 | |
| 181 | [[0., 0., 0.], |
| 182 | [0., 0., 0.], |
| 183 | [0., 0., 0.]], |
| 184 | |
| 185 | [[-1., -2., -1.], |
| 186 | [-2., -4., -2.], |
| 187 | [-1., -2., -1.]]]) |
| 188 | |
| 189 | self.sobel_x = nn.Parameter( |
| 190 | torch.from_numpy(sobel.transpose(0, 1, 2)).float().unsqueeze(0).unsqueeze(0).expand(self.classes, 1, 3, 3, |
| 191 | 3), requires_grad=False) |
| 192 | self.sobel_y = nn.Parameter( |
| 193 | torch.from_numpy(sobel.transpose(1, 0, 2)).float().unsqueeze(0).unsqueeze(0).expand(self.classes, 1, 3, 3, |
| 194 | 3), requires_grad=False) |
| 195 | self.sobel_z = nn.Parameter( |
| 196 | torch.from_numpy(sobel.transpose(1, 2, 0)).float().unsqueeze(0).unsqueeze(0).expand(self.classes, 1, 3, 3, |
| 197 | 3), requires_grad=False) |
| 198 | |
| 199 | self.diff_x = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, padding=1, |
| 200 | bias=False) |
| 201 | self.diff_x.weight = self.sobel_x |
| 202 | self.diff_y = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, padding=1, |
| 203 | bias=False) |
| 204 | self.diff_y.weight = self.sobel_y |
| 205 | self.diff_z = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, padding=1, |
| 206 | bias=False) |
| 207 | self.diff_z.weight = self.sobel_z |
| 208 | |
| 209 | def forward(self, predication, label): |
| 210 | grd_x = self.diff_x(predication) |
| 211 | grd_y = self.diff_y(predication) |
| 212 | grd_z = self.diff_z(predication) |
| 213 | |
| 214 | # length |
| 215 | length = torch.sqrt(grd_x ** 2 + grd_y ** 2 + grd_z ** 2 + 1e-8) |
| 216 | length = (length - length.min()) / (length.max() - length.min() + 1e-8) |
| 217 | length = torch.sum(length) |
| 218 | |
| 219 | # region |
| 220 | label = label.float() |
| 221 | c_in = torch.ones_like(predication) |
| 222 | c_out = torch.zeros_like(predication) |
| 223 | region_in = torch.abs(torch.sum(predication * ((label - c_in) ** 2))) |
| 224 | region_out = torch.abs( |