image-matching-webui / third_party /DarkFeat /nets /noise_reliability_loss.py
Vincentqyw
fix: roma
c74a070
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,
):
# subsample things
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,
)
# compute pixel-wise AP
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