Active contour based elastic model loss based on sobel and laplace filter
| 387 | |
| 388 | |
| 389 | class FastACELoss3D(nn.Module): |
| 390 | """ |
| 391 | Active contour based elastic model loss |
| 392 | based on sobel and laplace filter |
| 393 | """ |
| 394 | |
| 395 | def __init__(self, miu=1, alpha=1e-3, beta=2.0, classes=4, types="laplace"): |
| 396 | super(FastACELoss3D, self).__init__() |
| 397 | self.miu = miu |
| 398 | self.alpha = alpha |
| 399 | self.beta = beta |
| 400 | self.classes = classes |
| 401 | self.types = types |
| 402 | sobel = np.array([[[1., 2., 1.], |
| 403 | [2., 4., 2.], |
| 404 | [1., 2., 1.]], |
| 405 | |
| 406 | [[0., 0., 0.], |
| 407 | [0., 0., 0.], |
| 408 | [0., 0., 0.]], |
| 409 | |
| 410 | [[-1., -2., -1.], |
| 411 | [-2., -4., -2.], |
| 412 | [-1., -2., -1.]]]) |
| 413 | laplace_kernel = np.ones((3, 3, 3)) |
| 414 | laplace_kernel[1, 1, 1] = -26 |
| 415 | |
| 416 | self.sobel_x = nn.Parameter( |
| 417 | torch.from_numpy(sobel.transpose(0, 1, 2)).float().unsqueeze(0).unsqueeze(0).expand(self.classes, 1, 3, 3, |
| 418 | 3), requires_grad=False) |
| 419 | self.sobel_y = nn.Parameter( |
| 420 | torch.from_numpy(sobel.transpose(1, 0, 2)).float().unsqueeze(0).unsqueeze(0).expand(self.classes, 1, 3, 3, |
| 421 | 3), requires_grad=False) |
| 422 | self.sobel_z = nn.Parameter( |
| 423 | torch.from_numpy(sobel.transpose(1, 2, 0)).float().unsqueeze(0).unsqueeze(0).expand(self.classes, 1, 3, 3, |
| 424 | 3), requires_grad=False) |
| 425 | self.laplace = nn.Parameter( |
| 426 | torch.from_numpy(laplace_kernel).float().unsqueeze( |
| 427 | 0).unsqueeze(0).expand(self.classes, 1, 3, 3, 3), |
| 428 | requires_grad=False) |
| 429 | |
| 430 | self.diff_x = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, padding=1, |
| 431 | bias=False) |
| 432 | self.diff_x.weight = self.sobel_x |
| 433 | self.diff_y = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, padding=1, |
| 434 | bias=False) |
| 435 | self.diff_y.weight = self.sobel_y |
| 436 | self.diff_z = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, padding=1, |
| 437 | bias=False) |
| 438 | self.diff_z.weight = self.sobel_z |
| 439 | |
| 440 | self.laplace_operator = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, |
| 441 | padding=1, |
| 442 | bias=False) |
| 443 | self.laplace_operator.weight = self.laplace |
| 444 | |
| 445 | def forward(self, predication, label): |
| 446 | grd_x = self.diff_x(predication) |