Active contour based elastic model loss based on minpooling & maxpooling and laplace filter
| 480 | |
| 481 | |
| 482 | class FastACELoss3DV2(nn.Module): |
| 483 | """ |
| 484 | Active contour based elastic model loss |
| 485 | based on minpooling & maxpooling and laplace filter |
| 486 | """ |
| 487 | |
| 488 | def __init__(self, miu=1, alpha=1e-3, beta=2.0, classes=4, types="other"): |
| 489 | super(FastACELoss3DV2, self).__init__() |
| 490 | self.miu = miu |
| 491 | self.alpha = alpha |
| 492 | self.beta = beta |
| 493 | self.classes = classes |
| 494 | self.types = types |
| 495 | laplace_kernel = np.ones((3, 3, 3)) |
| 496 | laplace_kernel[1, 1, 1] = -26 |
| 497 | |
| 498 | self.laplace = nn.Parameter( |
| 499 | torch.from_numpy(laplace_kernel).float().unsqueeze( |
| 500 | 0).unsqueeze(0).expand(self.classes, 1, 3, 3, 3), |
| 501 | requires_grad=False) |
| 502 | |
| 503 | self.laplace_operator = nn.Conv3d(self.classes, self.classes, groups=self.classes, kernel_size=3, stride=1, |
| 504 | padding=1, |
| 505 | bias=False) |
| 506 | self.laplace_operator.weight = self.laplace |
| 507 | |
| 508 | def forward(self, predication, label): |
| 509 | min_pool_x = nn.functional.max_pool3d(predication * -1, 3, 1, 1) * -1 |
| 510 | contour = torch.relu(nn.functional.max_pool3d( |
| 511 | min_pool_x, 3, 1, 1) - min_pool_x) |
| 512 | |
| 513 | diff = self.laplace_operator(predication) |
| 514 | |
| 515 | # length |
| 516 | length = torch.abs(contour) |
| 517 | |
| 518 | # curvature |
| 519 | if self.types: |
| 520 | curvature = torch.abs(diff) |
| 521 | curvature = (curvature - curvature.min()) / \ |
| 522 | (curvature.max() - curvature.min() + 1e-8) |
| 523 | else: |
| 524 | """ |
| 525 | maybe more powerful |
| 526 | """ |
| 527 | curvature = torch.abs(diff) / ((length ** 2 + 1) ** 0.5 + 1e-8) |
| 528 | curvature = (curvature - curvature.min()) / \ |
| 529 | (curvature.max() - curvature.min() + 1e-8) |
| 530 | # region |
| 531 | label = label.float() |
| 532 | c_in = torch.ones_like(predication) |
| 533 | c_out = torch.zeros_like(predication) |
| 534 | region_in = torch.abs(torch.sum(predication * ((label - c_in) ** 2))) |
| 535 | region_out = torch.abs( |
| 536 | torch.sum((1 - predication) * ((label - c_out) ** 2))) |
| 537 | region = self.miu * region_in + region_out |
| 538 | |
| 539 | # elastic |