| import torch
|
| from metrics.registry import LOSSFUNC
|
| from .abstract_loss_func import AbstractLossClass
|
|
|
|
|
| def mahalanobis_distance(values: torch.Tensor, mean: torch.Tensor, inv_covariance: torch.Tensor) -> torch.Tensor:
|
| """Compute the batched mahalanobis distance.
|
|
|
| values is a batch of feature vectors.
|
| mean is either the mean of the distribution to compare, or a second
|
| batch of feature vectors.
|
| inv_covariance is the inverse covariance of the target distribution.
|
| """
|
| assert values.dim() == 2
|
| assert 1 <= mean.dim() <= 2
|
| assert inv_covariance.dim() == 2
|
| assert values.shape[1] == mean.shape[-1]
|
| assert mean.shape[-1] == inv_covariance.shape[0]
|
| assert inv_covariance.shape[0] == inv_covariance.shape[1]
|
|
|
| if mean.dim() == 1:
|
| mean = mean.unsqueeze(0)
|
| x_mu = values - mean
|
|
|
| dist = torch.einsum("im,mn,in->i", x_mu, inv_covariance, x_mu)
|
|
|
| return dist.sqrt()
|
|
|
|
|
| @LOSSFUNC.register_module(module_name="patch_consistency_loss")
|
| class PatchConsistencyLoss(AbstractLossClass):
|
| def __init__(self, c_real, c_fake, c_cross):
|
| super().__init__()
|
| self.c_real = c_real
|
| self.c_fake = c_fake
|
| self.c_cross = c_cross
|
|
|
| def forward(self, attention_map_real, attention_map_fake, feature_patch, real_feature_mean, real_inv_covariance,
|
| fake_feature_mean, fake_inv_covariance, labels):
|
|
|
| B, H, W, C = feature_patch.size()
|
| dist_real = mahalanobis_distance(feature_patch.reshape(B * H * W, C), real_feature_mean.cuda(),
|
| real_inv_covariance.cuda())
|
| dist_fake = mahalanobis_distance(feature_patch.reshape(B * H * W, C), fake_feature_mean.cuda(),
|
| fake_inv_covariance.cuda())
|
| fake_indices = torch.where(labels == 1.0)[0]
|
| index_map = torch.relu(dist_real - dist_fake).reshape((B, H, W))[fake_indices, :]
|
|
|
|
|
| if attention_map_real.shape[0] == 0:
|
| loss_real = 0
|
| else:
|
| B, PP, PP = attention_map_real.shape
|
| c_matrix = (1 - self.c_real) * torch.eye(PP).cuda() + self.c_real * torch.ones(PP).cuda()
|
| c_matrix = c_matrix.expand(B, -1, -1)
|
| loss_real = torch.sum(torch.abs(attention_map_real - c_matrix)) / (B * (PP * PP - PP))
|
|
|
| if attention_map_fake.shape[0] == 0:
|
| loss_fake = 0
|
| else:
|
| B, PP, PP = attention_map_fake.shape
|
| c_matrix = []
|
| for b in range(B):
|
| fake_indices = torch.where(index_map[b].reshape(-1) > 0)[0]
|
| real_indices = torch.where(index_map[b].reshape(-1) <= 0)[0]
|
| tmp = torch.zeros((PP, PP)).cuda() + self.c_cross
|
| for i in fake_indices:
|
| tmp[i, fake_indices] = self.c_fake
|
| for i in real_indices:
|
| tmp[i, real_indices] = self.c_real
|
| c_matrix.append(tmp)
|
|
|
| c_matrix = torch.stack(c_matrix).cuda()
|
| loss_fake = torch.sum(torch.abs(attention_map_fake - c_matrix)) / (B * (PP * PP - PP))
|
|
|
| return loss_real + loss_fake
|
|
|
