Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from .geom import rnd_sample, interpolate | |
| class NghSampler2(nn.Module): | |
| """Similar to NghSampler, but doesnt warp the 2nd image. | |
| Distance to GT => 0 ... pos_d ... neg_d ... ngh | |
| Pixel label => + + + + + + 0 0 - - - - - - - | |
| Subsample on query side: if > 0, regular grid | |
| < 0, random points | |
| In both cases, the number of query points is = W*H/subq**2 | |
| """ | |
| def __init__( | |
| self, | |
| ngh, | |
| subq=1, | |
| subd=1, | |
| pos_d=0, | |
| neg_d=2, | |
| border=None, | |
| maxpool_pos=True, | |
| subd_neg=0, | |
| ): | |
| nn.Module.__init__(self) | |
| assert 0 <= pos_d < neg_d <= (ngh if ngh else 99) | |
| self.ngh = ngh | |
| self.pos_d = pos_d | |
| self.neg_d = neg_d | |
| assert subd <= ngh or ngh == 0 | |
| assert subq != 0 | |
| self.sub_q = subq | |
| self.sub_d = subd | |
| self.sub_d_neg = subd_neg | |
| if border is None: | |
| border = ngh | |
| assert border >= ngh, "border has to be larger than ngh" | |
| self.border = border | |
| self.maxpool_pos = maxpool_pos | |
| self.precompute_offsets() | |
| def precompute_offsets(self): | |
| pos_d2 = self.pos_d**2 | |
| neg_d2 = self.neg_d**2 | |
| rad2 = self.ngh**2 | |
| rad = (self.ngh // self.sub_d) * self.ngh # make an integer multiple | |
| pos = [] | |
| neg = [] | |
| for j in range(-rad, rad + 1, self.sub_d): | |
| for i in range(-rad, rad + 1, self.sub_d): | |
| d2 = i * i + j * j | |
| if d2 <= pos_d2: | |
| pos.append((i, j)) | |
| elif neg_d2 <= d2 <= rad2: | |
| neg.append((i, j)) | |
| self.register_buffer("pos_offsets", torch.LongTensor(pos).view(-1, 2).t()) | |
| self.register_buffer("neg_offsets", torch.LongTensor(neg).view(-1, 2).t()) | |
| def gen_grid(self, step, B, H, W, dev): | |
| b1 = torch.arange(B, device=dev) | |
| if step > 0: | |
| # regular grid | |
| x1 = torch.arange(self.border, W - self.border, step, device=dev) | |
| y1 = torch.arange(self.border, H - self.border, step, device=dev) | |
| H1, W1 = len(y1), len(x1) | |
| x1 = x1[None, None, :].expand(B, H1, W1).reshape(-1) | |
| y1 = y1[None, :, None].expand(B, H1, W1).reshape(-1) | |
| b1 = b1[:, None, None].expand(B, H1, W1).reshape(-1) | |
| shape = (B, H1, W1) | |
| else: | |
| # randomly spread | |
| n = (H - 2 * self.border) * (W - 2 * self.border) // step**2 | |
| x1 = torch.randint(self.border, W - self.border, (n,), device=dev) | |
| y1 = torch.randint(self.border, H - self.border, (n,), device=dev) | |
| x1 = x1[None, :].expand(B, n).reshape(-1) | |
| y1 = y1[None, :].expand(B, n).reshape(-1) | |
| b1 = b1[:, None].expand(B, n).reshape(-1) | |
| shape = (B, n) | |
| return b1, y1, x1, shape | |
| def forward(self, feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=2500): | |
| pscores_ls, nscores_ls, distractors_ls = [], [], [] | |
| valid_feat0_ls = [] | |
| valid_pos1_ls, valid_pos2_ls = [], [] | |
| qconf_ls = [] | |
| mask_ls = [] | |
| for i in range(B): | |
| # positions in the first image | |
| tmp_mask = ( | |
| (pos0[i][:, 1] >= self.border) | |
| * (pos0[i][:, 1] < W - self.border) | |
| * (pos0[i][:, 0] >= self.border) | |
| * (pos0[i][:, 0] < H - self.border) | |
| ) | |
| selected_pos0 = pos0[i][tmp_mask] | |
| selected_pos1 = pos1[i][tmp_mask] | |
| valid_pos0, valid_pos1 = rnd_sample([selected_pos0, selected_pos1], N) | |
| # sample features from first image | |
| valid_feat0 = interpolate(valid_pos0 / 4, feat0[i]) # [N, 128] | |
| valid_feat0 = F.normalize(valid_feat0, p=2, dim=-1) # [N, 128] | |
| qconf = interpolate(valid_pos0 / 4, conf0[i]) | |
| # sample GT from second image | |
| mask = ( | |
| (valid_pos1[:, 1] >= 0) | |
| * (valid_pos1[:, 1] < W) | |
| * (valid_pos1[:, 0] >= 0) | |
| * (valid_pos1[:, 0] < H) | |
| ) | |
| def clamp(xy): | |
| xy = xy | |
| torch.clamp(xy[0], 0, H - 1, out=xy[0]) | |
| torch.clamp(xy[1], 0, W - 1, out=xy[1]) | |
| return xy | |
| # compute positive scores | |
| valid_pos1p = clamp( | |
| valid_pos1.t()[:, None, :] | |
| + self.pos_offsets[:, :, None].to(valid_pos1.device) | |
| ) # [2, 29, N] | |
| valid_pos1p = valid_pos1p.permute(1, 2, 0).reshape( | |
| -1, 2 | |
| ) # [29, N, 2] -> [29*N, 2] | |
| valid_feat1p = interpolate(valid_pos1p / 4, feat1[i]).reshape( | |
| self.pos_offsets.shape[-1], -1, 128 | |
| ) # [29, N, 128] | |
| valid_feat1p = F.normalize(valid_feat1p, p=2, dim=-1) # [29, N, 128] | |
| pscores = ( | |
| (valid_feat0[None, :, :] * valid_feat1p).sum(dim=-1).t() | |
| ) # [N, 29] | |
| pscores, pos = pscores.max(dim=1, keepdim=True) | |
| sel = clamp( | |
| valid_pos1.t() + self.pos_offsets[:, pos.view(-1)].to(valid_pos1.device) | |
| ) | |
| qconf = (qconf + interpolate(sel.t() / 4, conf1[i])) / 2 | |
| # compute negative scores | |
| valid_pos1n = clamp( | |
| valid_pos1.t()[:, None, :] | |
| + self.neg_offsets[:, :, None].to(valid_pos1.device) | |
| ) # [2, 29, N] | |
| valid_pos1n = valid_pos1n.permute(1, 2, 0).reshape( | |
| -1, 2 | |
| ) # [29, N, 2] -> [29*N, 2] | |
| valid_feat1n = interpolate(valid_pos1n / 4, feat1[i]).reshape( | |
| self.neg_offsets.shape[-1], -1, 128 | |
| ) # [29, N, 128] | |
| valid_feat1n = F.normalize(valid_feat1n, p=2, dim=-1) # [29, N, 128] | |
| nscores = ( | |
| (valid_feat0[None, :, :] * valid_feat1n).sum(dim=-1).t() | |
| ) # [N, 29] | |
| if self.sub_d_neg: | |
| valid_pos2 = rnd_sample([selected_pos1], N)[0] | |
| distractors = interpolate(valid_pos2 / 4, feat1[i]) | |
| distractors = F.normalize(distractors, p=2, dim=-1) | |
| pscores_ls.append(pscores) | |
| nscores_ls.append(nscores) | |
| distractors_ls.append(distractors) | |
| valid_feat0_ls.append(valid_feat0) | |
| valid_pos1_ls.append(valid_pos1) | |
| valid_pos2_ls.append(valid_pos2) | |
| qconf_ls.append(qconf) | |
| mask_ls.append(mask) | |
| N = np.min([len(i) for i in qconf_ls]) | |
| # merge batches | |
| qconf = torch.stack([i[:N] for i in qconf_ls], dim=0).squeeze(-1) | |
| mask = torch.stack([i[:N] for i in mask_ls], dim=0) | |
| pscores = torch.cat([i[:N] for i in pscores_ls], dim=0) | |
| nscores = torch.cat([i[:N] for i in nscores_ls], dim=0) | |
| distractors = torch.cat([i[:N] for i in distractors_ls], dim=0) | |
| valid_feat0 = torch.cat([i[:N] for i in valid_feat0_ls], dim=0) | |
| valid_pos1 = torch.cat([i[:N] for i in valid_pos1_ls], dim=0) | |
| valid_pos2 = torch.cat([i[:N] for i in valid_pos2_ls], dim=0) | |
| dscores = torch.matmul(valid_feat0, distractors.t()) | |
| dis2 = (valid_pos2[:, 1] - valid_pos1[:, 1][:, None]) ** 2 + ( | |
| valid_pos2[:, 0] - valid_pos1[:, 0][:, None] | |
| ) ** 2 | |
| b = torch.arange(B, device=dscores.device)[:, None].expand(B, N).reshape(-1) | |
| dis2 += (b != b[:, None]).long() * self.neg_d**2 | |
| dscores[dis2 < self.neg_d**2] = 0 | |
| scores = torch.cat((pscores, nscores, dscores), dim=1) | |
| gt = scores.new_zeros(scores.shape, dtype=torch.uint8) | |
| gt[:, : pscores.shape[1]] = 1 | |
| return scores, gt, mask, qconf | |