Spaces:
Running
Running
| # Copyright 2019-present NAVER Corp. | |
| # CC BY-NC-SA 3.0 | |
| # Available only for non-commercial use | |
| import pdb | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from nets.sampler import FullSampler | |
| class CosimLoss(nn.Module): | |
| """Try to make the repeatability repeatable from one image to the other.""" | |
| def __init__(self, N=16): | |
| nn.Module.__init__(self) | |
| self.name = f"cosim{N}" | |
| self.patches = nn.Unfold(N, padding=0, stride=N // 2) | |
| def extract_patches(self, sal): | |
| patches = self.patches(sal).transpose(1, 2) # flatten | |
| patches = F.normalize(patches, p=2, dim=2) # norm | |
| return patches | |
| def forward(self, repeatability, aflow, **kw): | |
| B, two, H, W = aflow.shape | |
| assert two == 2 | |
| # normalize | |
| sali1, sali2 = repeatability | |
| grid = FullSampler._aflow_to_grid(aflow) | |
| sali2 = F.grid_sample(sali2, grid, mode="bilinear", padding_mode="border") | |
| patches1 = self.extract_patches(sali1) | |
| patches2 = self.extract_patches(sali2) | |
| cosim = (patches1 * patches2).sum(dim=2) | |
| return 1 - cosim.mean() | |
| class PeakyLoss(nn.Module): | |
| """Try to make the repeatability locally peaky. | |
| Mechanism: we maximize, for each pixel, the difference between the local mean | |
| and the local max. | |
| """ | |
| def __init__(self, N=16): | |
| nn.Module.__init__(self) | |
| self.name = f"peaky{N}" | |
| assert N % 2 == 0, "N must be pair" | |
| self.preproc = nn.AvgPool2d(3, stride=1, padding=1) | |
| self.maxpool = nn.MaxPool2d(N + 1, stride=1, padding=N // 2) | |
| self.avgpool = nn.AvgPool2d(N + 1, stride=1, padding=N // 2) | |
| def forward_one(self, sali): | |
| sali = self.preproc(sali) # remove super high frequency | |
| return 1 - (self.maxpool(sali) - self.avgpool(sali)).mean() | |
| def forward(self, repeatability, **kw): | |
| sali1, sali2 = repeatability | |
| return (self.forward_one(sali1) + self.forward_one(sali2)) / 2 | |