Spaces:
Running
Running
File size: 3,949 Bytes
404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
import torch
import torch.nn as nn
import numpy as np
class APLoss(nn.Module):
"""differentiable AP loss, through quantization.
Input: (N, M) values in [min, max]
label: (N, M) values in {0, 1}
Returns: list of query AP (for each n in {1..N})
Note: typically, you want to minimize 1 - mean(AP)
"""
def __init__(self, nq=25, min=0, max=1, euc=False):
nn.Module.__init__(self)
assert isinstance(nq, int) and 2 <= nq <= 100
self.nq = nq
self.min = min
self.max = max
self.euc = euc
gap = max - min
assert gap > 0
# init quantizer = non-learnable (fixed) convolution
self.quantizer = q = nn.Conv1d(1, 2 * nq, kernel_size=1, bias=True)
a = (nq - 1) / gap
# 1st half = lines passing to (min+x,1) and (min+x+1/a,0) with x = {nq-1..0}*gap/(nq-1)
q.weight.data[:nq] = -a
q.bias.data[:nq] = torch.from_numpy(
a * min + np.arange(nq, 0, -1)
) # b = 1 + a*(min+x)
# 2nd half = lines passing to (min+x,1) and (min+x-1/a,0) with x = {nq-1..0}*gap/(nq-1)
q.weight.data[nq:] = a
q.bias.data[nq:] = torch.from_numpy(
np.arange(2 - nq, 2, 1) - a * min
) # b = 1 - a*(min+x)
# first and last one are special: just horizontal straight line
q.weight.data[0] = q.weight.data[-1] = 0
q.bias.data[0] = q.bias.data[-1] = 1
def compute_AP(self, x, label):
N, M = x.shape
# print(x.shape, label.shape)
if self.euc: # euclidean distance in same range than similarities
x = 1 - torch.sqrt(2.001 - 2 * x)
# quantize all predictions
q = self.quantizer(x.unsqueeze(1))
q = torch.min(q[:, : self.nq], q[:, self.nq :]).clamp(
min=0
) # N x Q x M [1600, 20, 1681]
nbs = q.sum(dim=-1) # number of samples N x Q = c
rec = (q * label.view(N, 1, M).float()).sum(
dim=-1
) # nb of correct samples = c+ N x Q
prec = rec.cumsum(dim=-1) / (1e-16 + nbs.cumsum(dim=-1)) # precision
rec /= rec.sum(dim=-1).unsqueeze(1) # norm in [0,1]
ap = (prec * rec).sum(dim=-1) # per-image AP
return ap
def forward(self, x, label):
assert x.shape == label.shape # N x M
return self.compute_AP(x, label)
class PixelAPLoss(nn.Module):
"""Computes the pixel-wise AP loss:
Given two images and ground-truth optical flow, computes the AP per pixel.
feat1: (B, C, H, W) pixel-wise features extracted from img1
feat2: (B, C, H, W) pixel-wise features extracted from img2
aflow: (B, 2, H, W) absolute flow: aflow[...,y1,x1] = x2,y2
"""
def __init__(self, sampler, nq=20):
nn.Module.__init__(self)
self.aploss = APLoss(nq, min=0, max=1, euc=False)
self.name = "pixAP"
self.sampler = sampler
def loss_from_ap(self, ap, rel):
return 1 - ap
def forward(self, feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=1200):
# subsample things
scores, gt, msk, qconf = self.sampler(
feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=1200
)
# compute pixel-wise AP
n = qconf.numel()
if n == 0:
return 0
scores, gt = scores.view(n, -1), gt.view(n, -1)
ap = self.aploss(scores, gt).view(msk.shape)
pixel_loss = self.loss_from_ap(ap, qconf)
loss = pixel_loss[msk].mean()
return loss
class ReliabilityLoss(PixelAPLoss):
"""same than PixelAPLoss, but also train a pixel-wise confidence
that this pixel is going to have a good AP.
"""
def __init__(self, sampler, base=0.5, **kw):
PixelAPLoss.__init__(self, sampler, **kw)
assert 0 <= base < 1
self.base = base
def loss_from_ap(self, ap, rel):
return 1 - ap * rel - (1 - rel) * self.base
|