imcui/third_party/LiftFeat/loss/loss.py

import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import time


def dual_softmax_loss(X, Y, temp = 0.2):
    if X.size() != Y.size() or X.dim() != 2 or Y.dim() != 2:
        raise RuntimeError('Error: X and Y shapes must match and be 2D matrices')

    dist_mat = (X @ Y.t()) * temp
    conf_matrix12 = F.log_softmax(dist_mat, dim=1)
    conf_matrix21 = F.log_softmax(dist_mat.t(), dim=1)

    with torch.no_grad():
        conf12 = torch.exp( conf_matrix12 ).max(dim=-1)[0]
        conf21 = torch.exp( conf_matrix21 ).max(dim=-1)[0]
        conf = conf12 * conf21

    target = torch.arange(len(X), device = X.device)

    loss = F.nll_loss(conf_matrix12, target) + \
           F.nll_loss(conf_matrix21, target)

    return loss, conf


class LiftFeatLoss(nn.Module):
    def __init__(self,device,lam_descs=1,lam_fb_descs=1,lam_kpts=1,lam_heatmap=1,lam_normals=1,lam_coordinates=1,lam_fb_coordinates=1,depth_spvs=False):
        super().__init__()
        
        # loss parameters
        self.lam_descs=lam_descs
        self.lam_fb_descs=lam_fb_descs
        self.lam_kpts=lam_kpts
        self.lam_heatmap=lam_heatmap
        self.lam_normals=lam_normals
        self.lam_coordinates=lam_coordinates
        self.lam_fb_coordinates=lam_fb_coordinates
        self.depth_spvs=depth_spvs
        self.running_descs_loss=0
        self.running_kpts_loss=0
        self.running_heatmaps_loss=0
        self.loss_descs=0
        self.loss_fb_descs=0
        self.loss_kpts=0
        self.loss_heatmaps=0
        self.loss_normals=0
        self.loss_coordinates=0
        self.loss_fb_coordinates=0
        self.acc_coarse=0
        self.acc_fb_coarse=0
        self.acc_kpt=0
        self.acc_coordinates=0
        self.acc_fb_coordinates=0
        
        # device
        self.dev=device
        
    
    def check_accuracy(self,m1,m2,pts1=None,pts2=None,plot=False):
        with torch.no_grad():
            #dist_mat = torch.cdist(X,Y)
            dist_mat = m1 @ m2.t()
            nn = torch.argmax(dist_mat, dim=1)
            #nn = torch.argmin(dist_mat, dim=1)
            correct = nn == torch.arange(len(m1), device = m1.device)

            if pts1 is not None and plot:
                import matplotlib.pyplot as plt
                canvas = torch.zeros((60, 80),device=m1.device)
                pts1 = pts1[~correct]
                canvas[pts1[:,1].long(), pts1[:,0].long()] = 1
                canvas = canvas.cpu().numpy()
                plt.imshow(canvas), plt.show()

            acc = correct.sum().item() / len(m1)
            return acc    
    
    def compute_descriptors_loss(self,descs1,descs2,pts):
        loss=[]
        acc=0
        B,_,H,W=descs1.shape
        conf_list=[]
        
        for b in range(B):
            pts1,pts2=pts[b][:,:2],pts[b][:,2:]
            m1=descs1[b,:,pts1[:,1].long(),pts1[:,0].long()].permute(1,0)
            m2=descs2[b,:,pts2[:,1].long(),pts2[:,0].long()].permute(1,0)
            
            loss_per,conf_per=dual_softmax_loss(m1,m2)
            loss.append(loss_per.unsqueeze(0))
            conf_list.append(conf_per)
            
            acc_coarse_per=self.check_accuracy(m1,m2)
            acc += acc_coarse_per
            
        loss=torch.cat(loss,dim=-1).mean()
        acc /= B
        return loss,acc,conf_list
    
    
    def alike_distill_loss(self,kpts,alike_kpts):
        C, H, W = kpts.shape
        kpts = kpts.permute(1,2,0)
        # get ALike keypoints
        with torch.no_grad():
            labels = torch.ones((H, W), dtype = torch.long, device = kpts.device) * 64 # -> Default is non-keypoint (bin 64)
            offsets = (((alike_kpts/8) - (alike_kpts/8).long())*8).long()
            offsets =  offsets[:, 0] + 8*offsets[:, 1]  # Linear IDX
            labels[(alike_kpts[:,1]/8).long(), (alike_kpts[:,0]/8).long()] = offsets

        kpts = kpts.view(-1,C)
        labels = labels.view(-1)

        mask = labels < 64
        idxs_pos = mask.nonzero().flatten()
        idxs_neg = (~mask).nonzero().flatten()
        perm = torch.randperm(idxs_neg.size(0))[:len(idxs_pos)//32]
        idxs_neg = idxs_neg[perm]
        idxs = torch.cat([idxs_pos, idxs_neg])

        kpts = kpts[idxs]
        labels = labels[idxs]

        with torch.no_grad():
            predicted = kpts.max(dim=-1)[1]
            acc =  (labels == predicted)
            acc = acc.sum() / len(acc)

        kpts = F.log_softmax(kpts,dim=-1)
        loss = F.nll_loss(kpts, labels, reduction = 'mean')

        return loss, acc
    
    
    def compute_keypoints_loss(self,kpts1,kpts2,alike_kpts1,alike_kpts2):
        loss=[]
        acc=0
        B,_,H,W=kpts1.shape
        
        for b in range(B):
            loss_per1,acc_per1=self.alike_distill_loss(kpts1[b],alike_kpts1[b])
            loss_per2,acc_per2=self.alike_distill_loss(kpts2[b],alike_kpts2[b])
            loss_per=(loss_per1+loss_per2)
            acc_per=(acc_per1+acc_per2)/2
            loss.append(loss_per.unsqueeze(0))
            acc += acc_per
            
        loss=torch.cat(loss,dim=-1).mean()
        acc /= B
        return loss,acc
    
    
    def compute_heatmaps_loss(self,heatmaps1,heatmaps2,pts,conf_list):
        loss=[]
        B,_,H,W=heatmaps1.shape
        
        for b in range(B):
            pts1,pts2=pts[b][:,:2],pts[b][:,2:]
            h1=heatmaps1[b,0,pts1[:,1].long(),pts1[:,0].long()]
            h2=heatmaps2[b,0,pts2[:,1].long(),pts2[:,0].long()]
            
            conf=conf_list[b]
            loss_per1=F.l1_loss(h1,conf)
            loss_per2=F.l1_loss(h2,conf)
            loss_per=(loss_per1+loss_per2)
            loss.append(loss_per.unsqueeze(0))
            
        loss=torch.cat(loss,dim=-1).mean()
        return loss
    
    
    def normal_loss(self,normal,target_normal):
        # import pdb;pdb.set_trace()
        normal = normal.permute(1, 2, 0)
        target_normal = target_normal.permute(1,2,0)
        # loss = F.l1_loss(d_feat, depth_anything_normal_feat)
        dot = torch.cosine_similarity(normal, target_normal, dim=2)
        valid_mask = target_normal[:, :, 0].float() \
                    * (dot.detach() < 0.999).float() \
                    * (dot.detach() > -0.999).float()
        valid_mask = valid_mask > 0.0
        al = torch.acos(dot[valid_mask])
        loss = torch.mean(al)
        return loss
    
    
    def compute_normals_loss(self,normals1,normals2,DA_normals1,DA_normals2,megadepth_batch_size,coco_batch_size):
        loss=[]
        
        # import pdb;pdb.set_trace()
        
        # only MegaDepth image need depth-normal
        normals1=normals1[coco_batch_size:,...]
        normals2=normals2[coco_batch_size:,...]
        for b in range(len(DA_normals1)):
            normal1,normal2=normals1[b],normals2[b]
            loss_per1=self.normal_loss(normal1,DA_normals1[b].permute(2,0,1))
            loss_per2=self.normal_loss(normal2,DA_normals2[b].permute(2,0,1))
            loss_per=(loss_per1+loss_per2)
            loss.append(loss_per.unsqueeze(0))
        
        loss=torch.cat(loss,dim=-1).mean()
        return loss
    
    
    def coordinate_loss(self,coordinate,conf,pts1):
        with torch.no_grad():
            coordinate_detached = pts1 * 8
            offset_detached = (coordinate_detached/8) - (coordinate_detached/8).long()
            offset_detached = (offset_detached * 8).long()
            label = offset_detached[:, 0] + 8*offset_detached[:, 1]

        #pdb.set_trace()
        coordinate_log = F.log_softmax(coordinate, dim=-1)

        predicted = coordinate.max(dim=-1)[1]
        acc =  (label == predicted)
        acc = acc[conf > 0.1]
        acc = acc.sum() / len(acc)

        loss = F.nll_loss(coordinate_log, label, reduction = 'none')
        
        #Weight loss by confidence, giving more emphasis on reliable matches
        conf = conf / conf.sum()
        loss = (loss * conf).sum()

        return loss*2., acc
    
    def compute_coordinates_loss(self,coordinates,pts,conf_list):
        loss=[]
        acc=0
        B,_,H,W=coordinates.shape
        
        for b in range(B):
            pts1,pts2=pts[b][:,:2],pts[b][:,2:]
            coordinate=coordinates[b,:,pts1[:,1].long(),pts1[:,0].long()].permute(1,0)
            conf=conf_list[b]
            
            loss_per,acc_per=self.coordinate_loss(coordinate,conf,pts1)
            loss.append(loss_per.unsqueeze(0))
            acc += acc_per
            
        loss=torch.cat(loss,dim=-1).mean()
        acc /= B
        
        return loss,acc
        
        
    def forward(self,
                descs1,fb_descs1,kpts1,normals1,
                descs2,fb_descs2,kpts2,normals2,
                pts,coordinates,fb_coordinates,
                alike_kpts1,alike_kpts2,
                DA_normals1,DA_normals2,
                megadepth_batch_size,coco_batch_size
                ):
        # import pdb;pdb.set_trace()
        self.loss_descs,self.acc_coarse,conf_list=self.compute_descriptors_loss(descs1,descs2,pts)
        self.loss_fb_descs,self.acc_fb_coarse,fb_conf_list=self.compute_descriptors_loss(fb_descs1,fb_descs2,pts)
        
        # start=time.perf_counter()
        self.loss_kpts,self.acc_kpt=self.compute_keypoints_loss(kpts1,kpts2,alike_kpts1,alike_kpts2)
        # end=time.perf_counter()
        # print(f"kpts loss cost {end-start} seconds")
        
        # start=time.perf_counter()
        self.loss_normals=self.compute_normals_loss(normals1,normals2,DA_normals1,DA_normals2,megadepth_batch_size,coco_batch_size)
        # end=time.perf_counter()
        # print(f"normal loss cost {end-start} seconds")
        
        self.loss_coordinates,self.acc_coordinates=self.compute_coordinates_loss(coordinates,pts,conf_list)
        self.loss_fb_coordinates,self.acc_fb_coordinates=self.compute_coordinates_loss(fb_coordinates,pts,fb_conf_list)
        
        return {
            'loss_descs':self.lam_descs*self.loss_descs,
            'acc_coarse':self.acc_coarse,
            'loss_coordinates':self.lam_coordinates*self.loss_coordinates,
            'acc_coordinates':self.acc_coordinates,
            'loss_fb_descs':self.lam_fb_descs*self.loss_fb_descs,
            'acc_fb_coarse':self.acc_fb_coarse,
            'loss_fb_coordinates':self.lam_fb_coordinates*self.loss_fb_coordinates,
            'acc_fb_coordinates':self.acc_fb_coordinates,
            'loss_kpts':self.lam_kpts*self.loss_kpts,
            'acc_kpt':self.acc_kpt,
            'loss_normals':self.lam_normals*self.loss_normals,
        }