AnimeRun/flow/core/augmentor.py

# %BANNER_BEGIN%
# ---------------------------------------------------------------------
# %COPYRIGHT_BEGIN%
#
#  Magic Leap, Inc. ("COMPANY") CONFIDENTIAL
#
#  Unpublished Copyright (c) 2020
#  Magic Leap, Inc., All Rights Reserved.
#
# NOTICE:  All information contained herein is, and remains the property
# of COMPANY. The intellectual and technical concepts contained herein
# are proprietary to COMPANY and may be covered by U.S. and Foreign
# Patents, patents in process, and are protected by trade secret or
# copyright law.  Dissemination of this information or reproduction of
# this material is strictly forbidden unless prior written permission is
# obtained from COMPANY.  Access to the source code contained herein is
# hereby forbidden to anyone except current COMPANY employees, managers
# or contractors who have executed Confidentiality and Non-disclosure
# agreements explicitly covering such access.
#
# The copyright notice above does not evidence any actual or intended
# publication or disclosure  of  this source code, which includes
# information that is confidential and/or proprietary, and is a trade
# secret, of  COMPANY.   ANY REPRODUCTION, MODIFICATION, DISTRIBUTION,
# PUBLIC  PERFORMANCE, OR PUBLIC DISPLAY OF OR THROUGH USE  OF THIS
# SOURCE CODE  WITHOUT THE EXPRESS WRITTEN CONSENT OF COMPANY IS
# STRICTLY PROHIBITED, AND IN VIOLATION OF APPLICABLE LAWS AND
# INTERNATIONAL TREATIES.  THE RECEIPT OR POSSESSION OF  THIS SOURCE
# CODE AND/OR RELATED INFORMATION DOES NOT CONVEY OR IMPLY ANY RIGHTS
# TO REPRODUCE, DISCLOSE OR DISTRIBUTE ITS CONTENTS, OR TO MANUFACTURE,
# USE, OR SELL ANYTHING THAT IT  MAY DESCRIBE, IN WHOLE OR IN PART.
#
# %COPYRIGHT_END%
# ----------------------------------------------------------------------
# %AUTHORS_BEGIN%
#
#  Originating Authors: Paul-Edouard Sarlin
#  Edited: Siyao Li 
# 
# %AUTHORS_END%
# --------------------------------------------------------------------*/
# %BANNER_END%

from copy import deepcopy
from pathlib import Path
import torch
from torch import nn
# from seg_desc import seg_descriptor
from torch_scatter import scatter as super_pixel_pooling
import argparse

def MLP(channels: list, do_bn=True):
    """ Multi-layer perceptron """
    n = len(channels)
    layers = []
    for i in range(1, n):
        layers.append(
            nn.Conv1d(channels[i - 1], channels[i], kernel_size=1, bias=True))
        if i < (n-1):
            if do_bn:
                # layers.append(nn.BatchNorm1d(channels[i]))
                layers.append(nn.InstanceNorm1d(channels[i]))
            layers.append(nn.ReLU())
    return nn.Sequential(*layers)


def normalize_keypoints(kpts, image_shape):
    """ Normalize keypoints locations based on image image_shape"""
    _, _, height, width = image_shape
    one = kpts.new_tensor(1)
    size = torch.stack([one*width, one*width, one*height, one*height])[None]
    center = size / 2
    scaling = size.max(1, keepdim=True).values * 0.7
    # print(kpts.size(), center[:, None, :].size(), scaling[:, None, :].size())
    return (kpts - center[:, None, :]) / scaling[:, None, :]

class ThreeLayerDecoder(nn.Module):
    """ Joint encoding of visual appearance and location using MLPs"""
    def __init__(self, enc_dim):
        super().__init__()
        # input must be 3 channel (r, g, b)
        self.layer1 = nn.Conv2d(enc_dim, enc_dim, 3, padding=1)
        self.non_linear1 = nn.ReLU()
        self.layer2 = nn.Conv2d(enc_dim, enc_dim, 3, padding=1)
        self.non_linear2 = nn.ReLU()
        self.layer3 = nn.Conv2d(enc_dim, enc_dim, 1)

        self.norm1 = nn.InstanceNorm2d(enc_dim)
        self.norm2 = nn.InstanceNorm2d(enc_dim)
        self.norm3 = nn.InstanceNorm2d(enc_dim)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                nn.init.constant_(m.bias, 0.0)

    def forward(self, img):
        x = self.non_linear1(self.norm1(self.layer1(img)))
        x = self.non_linear2(self.norm2(self.layer2(x)))
        x = self.norm3(self.layer3(x))
        # x = self.non_linear1(self.layer1(img))
        # x = self.non_linear2(self.layer2(x))
        # x = self.layer3(x)
        return x


class SegmentDescriptor(nn.Module):
    """ Joint encoding of visual appearance and location using MLPs"""
    def __init__(self, enc_dim):
        super().__init__()
        # self.encoder = ThreeLayerEncoder(enc_dim)
        # self.super_pixel_pooling = 
        # use scatter
        # nn.init.constant_(self.encoder[-1].bias, 0.0)

    def forward(self, img, seg):
        # x = self.encoder(img)
        n, c, h, w = x.size()
        assert((h, w) == img.size()[2:4])
        return super_pixel_pooling(x.view(n, c, -1), seg.view(-1).long(), reduce='mean')
        # here return size is [1]xCx|Seg|

class SegmentAdaptor(nn.Module):
    """ Joint encoding of visual appearance and location using MLPs"""
    def __init__(self, enc_dim):
        super().__init__()
        self.encoder = ThreeLayerDecoder(enc_dim)
        # self.super_pixel_pooling = 
        # use scatter
        # nn.init.constant_(self.encoder[-1].bias, 0.0)

    def forward(self, desc, seg):
        x = desc[seg]
        x = self.decoder(x)
        return x
        # here return size is [1]xCxHxW

class KeypointEncoder(nn.Module):
    """ Joint encoding of visual appearance and location using MLPs"""
    def __init__(self, feature_dim, layers):
        super().__init__()
        self.encoder = MLP([4] + layers + [feature_dim])
        # for m in self.encoder.modules():
        #     if isinstance(m, nn.Conv2d):
        #         nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        #         nn.init.constant_(m.bias, 0.0)
        nn.init.constant_(self.encoder[-1].bias, 0.0)

    def forward(self, kpts):
        inputs = kpts.transpose(1, 2)
        # print(inputs.size(), 'wula!')
        x = self.encoder(inputs)
        # print(x.size())
        return x


def attention(query, key, value):
    dim = query.shape[1]
    scores = torch.einsum('bdhn,bdhm->bhnm', query, key) / dim**.5
    prob = torch.nn.functional.softmax(scores, dim=-1)
    return torch.einsum('bhnm,bdhm->bdhn', prob, value), prob


class MultiHeadedAttention(nn.Module):
    """ Multi-head attention to increase model expressivitiy """
    def __init__(self, num_heads: int, d_model: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.dim = d_model // num_heads
        self.num_heads = num_heads
        self.merge = nn.Conv1d(d_model, d_model, kernel_size=1)
        self.proj = nn.ModuleList([deepcopy(self.merge) for _ in range(3)])

    def forward(self, query, key, value):
        batch_dim = query.size(0)
        query, key, value = [l(x).view(batch_dim, self.dim, self.num_heads, -1)
                             for l, x in zip(self.proj, (query, key, value))]
        x, prob = attention(query, key, value)
        self.prob.append(prob)
        return self.merge(x.contiguous().view(batch_dim, self.dim*self.num_heads, -1))


class AttentionalPropagation(nn.Module):
    def __init__(self, feature_dim: int, num_heads: int):
        super().__init__()
        self.attn = MultiHeadedAttention(num_heads, feature_dim)
        self.mlp = MLP([feature_dim*2, feature_dim*2, feature_dim])
        nn.init.constant_(self.mlp[-1].bias, 0.0)

    def forward(self, x, source):
        message = self.attn(x, source, source)
        return self.mlp(torch.cat([x, message], dim=1))


class AttentionalGNN(nn.Module):
    def __init__(self, feature_dim: int, layer_names: list):
        super().__init__()
        self.layers = nn.ModuleList([
            AttentionalPropagation(feature_dim, 4)
            for _ in range(len(layer_names))])
        self.names = layer_names

    def forward(self, desc0, desc1):
        for layer, name in zip(self.layers, self.names):
            layer.attn.prob = []
            if name == 'cross':
                src0, src1 = desc1, desc0
            else:  # if name == 'self':
                src0, src1 = desc0, desc1
            delta0, delta1 = layer(desc0, src0), layer(desc1, src1)
            desc0, desc1 = (desc0 + delta0), (desc1 + delta1)
        return desc0, desc1


def log_sinkhorn_iterations(Z, log_mu, log_nu, iters: int):
    """ Perform Sinkhorn Normalization in Log-space for stability"""
    u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
    for _ in range(iters):
        u = log_mu - torch.logsumexp(Z + v.unsqueeze(1), dim=2)
        v = log_nu - torch.logsumexp(Z + u.unsqueeze(2), dim=1)
    return Z + u.unsqueeze(2) + v.unsqueeze(1)


def log_optimal_transport(scores, alpha, iters: int):
    """ Perform Differentiable Optimal Transport in Log-space for stability"""
    b, m, n = scores.shape
    one = scores.new_tensor(1)
    ms, ns = (m*one).to(scores), (n*one).to(scores)

    bins0 = alpha.expand(b, m, 1)
    bins1 = alpha.expand(b, 1, n)
    alpha = alpha.expand(b, 1, 1)

    couplings = torch.cat([torch.cat([scores, bins0], -1),
                           torch.cat([bins1, alpha], -1)], 1)

    norm = - (ms + ns).log()
    log_mu = torch.cat([norm.expand(m), ns.log()[None] + norm])
    log_nu = torch.cat([norm.expand(n), ms.log()[None] + norm])
    log_mu, log_nu = log_mu[None].expand(b, -1), log_nu[None].expand(b, -1)

    Z = log_sinkhorn_iterations(couplings, log_mu, log_nu, iters)
    Z = Z - norm  # multiply probabilities by M+N
    return Z


def arange_like(x, dim: int):
    return x.new_ones(x.shape[dim]).cumsum(0) - 1  # traceable in 1.1


class Augmentor(nn.Module):
    """SuperGlue feature matching middle-end

    Given two sets of keypoints and locations, we determine the
    correspondences by:
      1. Keypoint Encoding (normalization + visual feature and location fusion)
      2. Graph Neural Network with multiple self and cross-attention layers
      3. Final projection layer
      4. Optimal Transport Layer (a differentiable Hungarian matching algorithm)
      5. Thresholding matrix based on mutual exclusivity and a match_threshold

    The correspondence ids use -1 to indicate non-matching points.

    Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew
    Rabinovich. SuperGlue: Learning Feature Matching with Graph Neural
    Networks. In CVPR, 2020. https://arxiv.org/abs/1911.11763

    """
    # default_config = {
    #     'descriptor_dim': 128,
    #     'weights': 'indoor',
    #     'keypoint_encoder': [32, 64, 128],
    #     'GNN_layers': ['self', 'cross'] * 9,
    #     'sinkhorn_iterations': 100,
    #     'match_threshold': 0.2,
    # }

    def __init__(self, config=None):
        super().__init__()

        default_config = argparse.Namespace()
        default_config.descriptor_dim = 128
        # default_config.weights = 
        # default_config.keypoint_encoder = [32, 64, 128]
        default_config.GNN_layers = ['self', 'cross'] * 9
        default_config.sinkhorn_iterations = 100
        default_config.match_threshold = 0.2
        # self.config = {**self.default_config, **config}

        if config is None:
            self.config = default_config
        else:
            self.config = config   
            self.config.GNN_layers = ['self', 'cross'] * self.config.GNN_layer_num
            # print('WULA!', self.config.GNN_layer_num)

        self.kenc = KeypointEncoder(
            self.config.descriptor_dim, self.config.keypoint_encoder)

        self.gnn = AttentionalGNN(
            self.config.descriptor_dim, self.config.GNN_layers)

        self.final_proj = nn.Conv1d(
            self.config.descriptor_dim, self.config.descriptor_dim,
            kernel_size=1, bias=True)

        bin_score = torch.nn.Parameter(torch.tensor(1.))
        self.register_parameter('bin_score', bin_score)
        self.segment_adaptor = SegmentAdaptor(self.config.descriptor_dim)

        # assert self.config.weights in ['indoor', 'outdoor']
        # path = Path(__file__).parent
        # path = path / 'weights/superglue_{}.pth'.format(self.config.weights)
        # self.load_state_dict(torch.load(path))
        # print('Loaded SuperGlue model (\"{}\" weights)'.format(
        #     self.config.weights))

    def forward(self, data):
        """Run SuperGlue on a pair of keypoints and descriptors"""
        # print(data['segment0'].size())
        # desc0, desc1 = data['descriptors0'].float()(), data['descriptors1'].float()()
        desc0, desc1 = self.segment_adaptor(data['feat0'], data['segment0']), self.segment_adaptor(data['feat1'], data['segment1'])
        # print(desc0.size())
        kpts0, kpts1 = data['keypoints0'].float(), data['keypoints1'].float()

        # desc0 = desc0.transpose(0,1)
        # desc1 = desc1.transpose(0,1)
        # kpts0 = torch.reshape(kpts0, (1, -1, 2))
        # kpts1 = torch.reshape(kpts1, (1, -1, 2))

        # if kpts0.shape[1] < 2 or kpts1.shape[1] < 2:  # no keypoints
        #     shape0, shape1 = kpts0.shape[:-1], kpts1.shape[:-1]
        #     # print(data['file_name'])
        #     return {
        #         'matches0': kpts0.new_full(shape0, -1, dtype=torch.int)[0],
        #         # 'matches1': kpts1.new_full(shape1, -1, dtype=torch.int)[0],
        #         'matching_scores0': kpts0.new_zeros(shape0)[0],
        #         # 'matching_scores1': kpts1.new_zeros(shape1)[0],
        #         'skip_train': True
        #     }

        # file_name = data['file_name']
        # all_matches = data['all_matches'] if 'all_matches' in data else None# shape = (1, K1)
        # .permute(1,2,0) # shape=torch.Size([1, 87,])
        
        # positional embedding
        # Keypoint normalization.
        kpts0 = normalize_keypoints(kpts0, data['image0'].shape)
        kpts1 = normalize_keypoints(kpts1, data['image1'].shape)

        # Keypoint MLP encoder.
        # print(data['file_name'])
        # print(kpts0.size())
    
        pos0 = self.kenc(kpts0)
        pos1 = self.kenc(kpts1)
        # print(desc0.size(), pos0.size())
        desc0 = desc0 + pos0
        desc1 = desc1 + pos1

        # Multi-layer Transformer network.
        desc0, desc1 = self.gnn(desc0, desc1)

        # Final MLP projection.
        mdesc0, mdesc1 = self.final_proj(desc0), self.final_proj(desc1)
        return mdesc0, mdesc1, None
        # TODO + loss


        # # Compute matching descriptor distance.
        # # print(mdesc0.size(), mdesc1.size())
        # scores = torch.einsum('bdn,bdm->bnm', mdesc0, mdesc1)
        # # #print('here1!!', scores.size())

        # # b k1 k2
        # scores = scores / self.config.descriptor_dim**.5
        # # print(scores.size())
        # # Run the optimal transport.
        # scores = log_optimal_transport(
        #     scores, self.bin_score,
        #     iters=self.config.sinkhorn_iterations)

        # # print(scores)
        # # print(scores.sum())
        # # print(scores.sum(1))
        # # print(scores.sum(0))

        # # Get the matches with score above "match_threshold".
        # max0, max1 = scores[:, :-1, :-1].max(2), scores[:, :-1, :-1].max(1)
        # indices0, indices1 = max0.indices, max1.indices
        # mutual0 = arange_like(indices0, 1)[None] == indices1.gather(1, indices0)
        # mutual1 = arange_like(indices1, 1)[None] == indices0.gather(1, indices1)
        # zero = scores.new_tensor(0)
        # mscores0 = torch.where(mutual0, max0.values.exp(), zero)
        # mscores1 = torch.where(mutual1, mscores0.gather(1, indices1), zero)
        # valid0 = mutual0 & (mscores0 > self.config.match_threshold)
        # valid1 = mutual1 & valid0.gather(1, indices1)
        # indices0 = torch.where(valid0, indices0, indices0.new_tensor(-1))
        # indices1 = torch.where(valid1, indices1, indices1.new_tensor(-1))

        # # check if indexed correctly
        # # #print(scores.size())
        # loss = []
        # weights = data['num0'].float().cuda()
        # avg_w = weights.mean()
        # # print(weights)
        
        # #print(scores.size())
        # if all_matches is not None:
        #     for i in range(len(all_matches[0])):
        #         # x = all_matches[0][i][0]
        #         x = i
        #         y = all_matches[0][i].long()
        #         # #print(y.data.cpu().numpy())
                
        #         loss.append(-scores[0][x][y] * weights[0][0][x]/avg_w) # check batch size == 1 ?
        # # print(len(indices1[0]))
        
        # # for j in range(len(indices1[0])):
        # #     if j not in all_matches[0]:
        # #         loss.append(-scores[0][-1][j])
        # # loss.append()
        # # for p0 in unmatched0:
        # #     loss += -torch.log(scores[0][p0][-1])
        # # for p1 in unmatched1:
        # #     loss += -torch.log(scores[0][-1][p1])
        #     loss_mean = torch.mean(torch.stack(loss))
        #     loss_mean = torch.reshape(loss_mean, (1, -1))
        #     #print((all_matches[0] == indices0[0]).sum())
        #     return {
        #         'matches0': indices0[0], # use -1 for invalid match
        #         # 'matches1': indices1[0], # use -1 for invalid match
        #         'matching_scores0': mscores0[0],
        #         # 'matching_scores1': mscores1[0],
        #         'loss': loss_mean[0],
        #         'skip_train': False,
        #         'accuracy': ((all_matches[0] == indices0[0]).sum() / len(all_matches[0])).item(),
        #         'area_accuracy': (torch.tensor([ (data['segment0'] == ii).sum() for ii in torch.arange(0, all_matches[0].shape[0])[all_matches[0] == indices0[0]]]).sum() / (weights.sum() * 1.0)).item(),
        #         'valid_accuracy': (((all_matches[0] == indices0[0]) & (all_matches[0] != -1)).sum() / (all_matches[0] != -1).sum()).item(),
        #     }
        # else:
        #     return {
        #         'matches0': indices0[0], # use -1 for invalid match
        #         'matching_scores0': mscores0[0],
        #         'loss': -1,
        #         'skip_train': True,
        #         'accuracy': -1,
        #         'area_accuracy': -1,
        #         'valid_accuracy': -1,
        #     }


if __name__ == '__main__':
    from anime_seg_mat_dataset import fetch_dataloader
    args = argparse.Namespace()
    ss = SuperGlue()
    args.subset = 'trytry'
    args.batch_size = 1
    args.stage = 'anime'
    args.image_size = (368, 368)
    loader = fetch_dataloader(args)
    # #print(len(loader))
    for data in loader:
        # p1, p2, s1, s2, mi = data
        dict1 = data

        kp1 = dict1['keypoints0']
        kp2 = dict1['keypoints1']
        p1 = dict1['image0']
        p2 = dict1['image1']  
        s1 = dict1['segment0']
        s2 = dict1['segment1']
        # #print(s1)
        # #print(s1.type)
        mi = dict1['all_matches']
        fname = dict1['file_name']   
        # #print(mi.size())  
        # #print(mi)

        a = ss(data)
        #print(dict1['file_name'])
        # print(a['loss'])
        # a['loss'].backward()
        # print(a['matches0'].size())
        # print(a['accuracy'], a['valid_accuracy'])