Spaces:
Running
Running
File size: 1,731 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
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 |
# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use
import pdb
import torch.nn as nn
import torch.nn.functional as F
from nets.ap_loss import APLoss
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, descriptors, aflow, **kw):
# subsample things
scores, gt, msk, qconf = self.sampler(descriptors, kw.get("reliability"), aflow)
# 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
self.name = "reliability"
def loss_from_ap(self, ap, rel):
return 1 - ap * rel - (1 - rel) * self.base
|