loss_functions.py

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import cv2
import math
from config import cfg, cfg_from_file, save_config_to_file


def MultiScale_UP(output,target,loss_type='L1',weight=[1.,0.5,0.25],valid_range=None,removezero=False,extra_mask=None):

    # Multi scale Loss, where flow upsampling is automaticly computed here
    loss = 0
    loss_list = []
    b, _, h, w = target.size()

    
    if (type(output) is not tuple) and (type(output) is not set):
        output = {output}
    
    for i, cur_output in enumerate(output):
        # Compute loss for each level
        realflow = F.interpolate(cur_output, (h,w), mode='bilinear', align_corners=True)
        realflow[:,0,:,:] = realflow[:,0,:,:]*(w/cur_output.shape[3])
        realflow[:,1,:,:] = realflow[:,1,:,:]*(h/cur_output.shape[2])

        with torch.no_grad():
            if i==0: epe = realEPE(realflow,target,extra_mask=extra_mask)

        if loss_type=='L2':
            lossvalue = torch.norm(realflow-target,p=2,dim=1)
        elif loss_type=='robust':
            lossvalue = ((realflow-target).abs().sum(dim=1)+1e-8)
            lossvalue = lossvalue**0.4
        elif loss_type=='L1':
            lossvalue = (realflow-target).abs().sum(dim=1)
        else:
            raise NotImplementedError

        if cfg.USE_VALID_RANGE and valid_range is not None:
            # Filter out the pixels whose gt flow is out of our search range
            with torch.no_grad():
                mask = (target[:,0,:,:].abs()<=valid_range[i][1]) & (target[:,1,:,:].abs()<=valid_range[i][0])
        else:
            with torch.no_grad():
                mask = torch.ones(target[:,0,:,:].shape).type_as(target)

        lossvalue = lossvalue*mask.float() 
        
        if extra_mask is not None:
            val = extra_mask > 0
            lossvalue = lossvalue[val]
            cur_loss = lossvalue.mean()*weight[i]
            assert lossvalue.shape[0] == extra_mask.sum()
        else:
            cur_loss = lossvalue.mean()*weight[i]

        loss+=cur_loss
        loss_list.append(cur_loss)

    loss = loss/len(output)
    return loss,loss_list,epe


#################################################

def random_select_points(x,y,x_,y_,samples=10):
    idx=torch.randperm(x.shape[0])
    x=x[idx[:samples],:]
    y=y[idx[:samples],:]
    x_=x_[idx[:samples],:]
    y_=y_[idx[:samples],:]
    return x,y,x_,y_

def subspace_loss_batch(flow):

    # https://openaccess.thecvf.com/content_CVPR_2019/papers/Zhong_Unsupervised_Deep_Epipolar_Flow_for_Stationary_or_Dynamic_Scenes_CVPR_2019_paper.pdf
    B, _, H, W = flow.size()
    xx = Variable(torch.arange(0, W).view(1,-1).repeat(H,1).cuda())
    yy = Variable(torch.arange(0, H).view(-1,1).repeat(1,W).cuda())
    grid_x = xx.view(1,1,H,W).repeat(B,1,1,1).float()
    grid_y = yy.view(1,1,H,W).repeat(B,1,1,1).float()
    
    flow_u = flow[:,0,:,:].unsqueeze(1)
    flow_v = flow[:,1,:,:].unsqueeze(1)
    
    pos_x = grid_x + flow_u
    pos_y = grid_y + flow_v

    inside_x = (pos_x <= (W-1)) & (pos_x >= 0.0)
    inside_y = (pos_y <= (H-1)) & (pos_y >= 0.0)
    
    inside = inside_x & inside_y 
    
    loss = 0

    least_num = 2000 

    list_X =[]
    list_X_ = []
    for i in range(B):
        grid_x_i = grid_x[i,:,:,:]
        grid_y_i = grid_y[i,:,:,:]
        pos_x_i = pos_x[i,:,:,:]
        pos_y_i = pos_y[i,:,:,:]
        inside_i= inside[i,:,:,:]
        
        if inside_i.sum()>least_num:
            x  = torch.masked_select(grid_x_i, inside_i).view(-1,1)
            y  = torch.masked_select(grid_y_i, inside_i).view(-1,1)
            x_ = torch.masked_select(pos_x_i, inside_i).view(-1,1)
            y_ = torch.masked_select(pos_y_i, inside_i).view(-1,1)
            x, y, x_, y_ = random_select_points(x,y,x_,y_,samples=least_num)
            o  = torch.ones_like(x)
            x, y, x_, y_ = x/W, y/W, x_/W, y_/W
            X  = torch.cat((x,x,x,y,y,y,o,o,o),1).permute(1,0)  
            X_ = torch.cat((x_,y_,o,x_,y_,o,x_,y_,o),1).permute(1,0)
            list_X.append(X.unsqueeze(0))
            list_X_.append(X_.unsqueeze(0))

    all_X = torch.cat(list_X)
    all_X_ = torch.cat(list_X_)
    M = all_X*all_X_
    lambda1 = 10
    MTM = lambda1 * torch.matmul(M.permute(0,2,1),M)
    I = torch.eye(MTM.size()[1]).type_as(MTM).unsqueeze(0).repeat(B,1,1)
    MTM_inverse = torch.inverse((I + MTM))
    C = torch.matmul(MTM_inverse,MTM)
    C2 = C**2
    loss1 = torch.sum(C2.view(-1,1),dim=0)

    loss2 = lambda1 * torch.sum(((torch.matmul(M,C)-M)**2).view(-1,1),dim=0)

    loss +=  (loss1 + loss2)

    return loss/B


def EPE_flow(input_flow, target_flow):
    return torch.norm(target_flow-input_flow,p=2,dim=1).mean()

class L1(nn.Module):
    def __init__(self):
        super(L1, self).__init__()
    def forward(self, output, target):
        lossvalue = torch.abs(output - target).mean()
        return lossvalue

class L2(nn.Module):
    def __init__(self):
        super(L2, self).__init__()
    def forward(self, output, target):
        lossvalue = torch.norm(output-target,p=2,dim=1).mean()
        return lossvalue

class L1Loss(nn.Module):
    def __init__(self):
        super(L1Loss, self).__init__()
        self.loss = L1()
        self.loss_labels = ['L1', 'EPE']

    def forward(self, output, target):
        lossvalue = self.loss(output, target)
        epevalue = EPE_flow(output, target)
        return [lossvalue, epevalue]

class L2Loss(nn.Module):
    def __init__(self):
        super(L2Loss, self).__init__()
        self.loss = L2()
        self.loss_labels = ['L2', 'EPE']

    def forward(self, output, target):
        lossvalue = self.loss(output, target)
        epevalue = EPE_flow(output, target)
        return [lossvalue, epevalue]




########################################################


def warp(x, flo):
    """
    warp an image/tensor (im2) back to im1, according to the optical flow

    x: [B, C, H, W] (im2)
    flo: [B, 2, H, W] flow

    """
    B, C, H, W = x.size()
        # mesh grid 
    xx = torch.arange(0, W).view(1,-1).repeat(H,1)
    yy = torch.arange(0, H).view(-1,1).repeat(1,W)
    xx = xx.view(1,1,H,W).repeat(B,1,1,1)
    yy = yy.view(1,1,H,W).repeat(B,1,1,1)
    grid = torch.cat((xx,yy),1).float()

    if x.is_cuda:
        grid = grid.cuda()
    vgrid = Variable(grid) + flo

        # scale grid to [-1,1] 
    vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:]/max(W-1,1)-1.0
    vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:]/max(H-1,1)-1.0

    vgrid = vgrid.permute(0,2,3,1)        
    output = nn.functional.grid_sample(x, vgrid)
    mask = torch.autograd.Variable(torch.ones(x.size())).cuda()
    mask = nn.functional.grid_sample(mask, vgrid)

    mask[mask<0.9999] = 0
    mask[mask>0] = 1
        
    return output*mask
      
def warp_(self, tensorInput, tensorFlow):
    if hasattr(self, 'tensorPartial') == False or self.tensorPartial.size(0) != tensorFlow.size(0) or self.tensorPartial.size(2) != tensorFlow.size(2) or self.tensorPartial.size(3) != tensorFlow.size(3):
        self.tensorPartial = tensorFlow.new_ones(tensorFlow.size(0), 1, tensorFlow.size(2), tensorFlow.size(3))
        # end

    if hasattr(self, 'tensorGrid') == False or self.tensorGrid.size(0) != tensorFlow.size(0) or self.tensorGrid.size(2) != tensorFlow.size(2) or self.tensorGrid.size(3) != tensorFlow.size(3):
        tensorHorizontal = torch.linspace(-1.0, 1.0, tensorFlow.size(3)).view(1, 1, 1, tensorFlow.size(3)).expand(tensorFlow.size(0), -1, tensorFlow.size(2), -1)
        tensorVertical = torch.linspace(-1.0, 1.0, tensorFlow.size(2)).view(1, 1, tensorFlow.size(2), 1).expand(tensorFlow.size(0), -1, -1, tensorFlow.size(3))

        self.tensorGrid = torch.cat([ tensorHorizontal, tensorVertical ], 1).cuda()
        # end

    tensorInput = torch.cat([ tensorInput, self.tensorPartial ], 1)
    tensorFlow = torch.cat([ tensorFlow[:, 0:1, :, :] / ((tensorInput.size(3) - 1.0) / 2.0), tensorFlow[:, 1:2, :, :] / ((tensorInput.size(2) - 1.0) / 2.0) ], 1)

    tensorOutput = torch.nn.functional.grid_sample(input=tensorInput, grid=(self.tensorGrid + tensorFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros')

    tensorMask = tensorOutput[:, -1:, :, :]; tensorMask[tensorMask > 0.999] = 1.0; tensorMask[tensorMask < 1.0] = 0.0

    return tensorOutput[:, :-1, :, :] * tensorMask


def EPE(input_flow, target_flow, sparse=False, mean=True,extra_mask=None):
    EPE_map = torch.norm(target_flow-input_flow,2,1)
    batch_size = EPE_map.size(0)

    if sparse:
        # invalid flow is defined with both flow coordinates to be exactly 0
        mask = (target_flow[:,0] == 0) & (target_flow[:,1] == 0)
        EPE_map = EPE_map[~mask]

    if extra_mask is not None:
        EPE_map = EPE_map[extra_mask.byte()]

    if mean:
        return EPE_map.mean()
    else:
        return EPE_map.sum()/batch_size

def realEPE(output, target, sparse=False, valid_range=None,extra_mask=None):
    b, _, h, w = target.size()
    upsampled_output = output
    if cfg.USE_VALID_RANGE and valid_range is not None:
        mask = (target[:,0,:,:].abs()<=valid_range[1]) & (target[:,1,:,:].abs()<=valid_range[0])
        mask = mask.unsqueeze(1).expand(-1,2,-1,-1).float()
        upsampled_output = upsampled_output*mask
        target = target*mask
    return EPE(upsampled_output, target, sparse, mean=True,extra_mask=extra_mask)