|
import torch |
|
import torch.nn as nn |
|
from .reliability_loss import APLoss |
|
|
|
|
|
class MultiPixelAPLoss(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.sampler = sampler |
|
self.base = 0.25 |
|
self.dec_base = 0.20 |
|
|
|
def loss_from_ap(self, ap, rel, noise_ap, noise_rel): |
|
dec_ap = torch.clamp(ap - noise_ap, min=0, max=1) |
|
return (1 - ap * noise_rel - (1 - noise_rel) * self.base), ( |
|
1.0 - dec_ap * (1 - noise_rel) - noise_rel * self.dec_base |
|
) |
|
|
|
def forward( |
|
self, |
|
feat0, |
|
feat1, |
|
noise_feat0, |
|
noise_feat1, |
|
conf0, |
|
conf1, |
|
noise_conf0, |
|
noise_conf1, |
|
pos0, |
|
pos1, |
|
B, |
|
H, |
|
W, |
|
N=1500, |
|
): |
|
|
|
scores, noise_scores, gt, msk, qconf, noise_qconf = self.sampler( |
|
feat0, |
|
feat1, |
|
noise_feat0, |
|
noise_feat1, |
|
conf0, |
|
conf1, |
|
noise_conf0, |
|
noise_conf1, |
|
pos0, |
|
pos1, |
|
B, |
|
H, |
|
W, |
|
N=1500, |
|
) |
|
|
|
|
|
n = qconf.numel() |
|
if n == 0: |
|
return 0, 0 |
|
scores, noise_scores, gt = scores.view(n, -1), noise_scores, gt.view(n, -1) |
|
ap = self.aploss(scores, gt).view(msk.shape) |
|
noise_ap = self.aploss(noise_scores, gt).view(msk.shape) |
|
|
|
pixel_loss = self.loss_from_ap(ap, qconf, noise_ap, noise_qconf) |
|
|
|
loss = pixel_loss[0][msk].mean(), pixel_loss[1][msk].mean() |
|
return loss |
|
|
