<reponame>SamGalanakis/FlowCompare
import os
import sys
import copy
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.nets import MLP
# Original code from : https://github.com/WangYueFt/dgcnn/blob/master/pytorch/model.py


def knn(x, k):
    inner = -2*torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    # (batch_size, num_points, k)
    idx = pairwise_distance.topk(k=k, dim=-1)[1]
    return idx


def get_graph_feature(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(2)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        idx = knn(x, k=k)   # (batch_size, num_points, k)

    idx_base = torch.arange(
        0, batch_size, device=x.device).view(-1, 1, 1)*num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()

    # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    x = x.transpose(2, 1).contiguous()
    feature = x.view(batch_size*num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims)
    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)

    feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2).contiguous()

    return feature


class DGCNNembedder(nn.Module):
    """DGCNN embedder with n_points x embedding dim as output"""
    def __init__(self, out_mlp_dims, emb_dim=22, dropout=0, n_neighbors=20):
        super().__init__()

        self.n_neighbors = n_neighbors

        self.bn1 = nn.BatchNorm2d(64)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(256)
        self.bn5 = nn.BatchNorm1d(512)

        self.conv1 = nn.Sequential(nn.Conv2d(12, 64, kernel_size=1, bias=False),
                                   self.bn1,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
                                   self.bn2,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False),
                                   self.bn3,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False),
                                   self.bn4,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv5 = nn.Sequential(nn.Conv1d(512, 512, kernel_size=1, bias=False),
                                   self.bn5,
                                   nn.LeakyReLU(negative_slope=0.2))

        self.out_mlp = MLP(512, out_mlp_dims, emb_dim, torch.nn.GELU())

    def forward(self, x):
        batch_size = x.size(0)
        x = x.permute((0, 2, 1))
        
    
        x = get_graph_feature(x, k=self.n_neighbors)
        x = self.conv1(x)
        x1 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x1, self.n_neighbors)
        x = self.conv2(x)
        x2 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x2, k=self.n_neighbors)
        x = self.conv3(x)
        x3 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x3, k=self.n_neighbors)
        x = self.conv4(x)
        x4 = x.max(dim=-1, keepdim=False)[0]

        x = torch.cat((x1, x2, x3, x4),dim=1) 

        x = self.conv5(x).permute((0, 2, 1))

        x = self.out_mlp(x)
        return x


class DGCNNembedderGlobal(nn.Module):
    """DGCNN embedder with single embedding vector as output"""
    def __init__(self, input_dim, out_mlp_dims, emb_dim=22, n_neighbors=20):
        super().__init__()

        self.n_neighbors = n_neighbors

        self.input_dim = input_dim

        self.bn1 = nn.BatchNorm2d(64)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(256)
        self.bn5 = nn.BatchNorm1d(512)

        self.conv1 = nn.Sequential(nn.Conv2d(self.input_dim*2, 64, kernel_size=1, bias=False),
                                   self.bn1,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
                                   self.bn2,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False),
                                   self.bn3,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False),
                                   self.bn4,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv5 = nn.Sequential(nn.Conv1d(512, 512, kernel_size=1, bias=False),
                                   self.bn5,
                                   nn.LeakyReLU(negative_slope=0.2))

        self.out_mlp = MLP(512*2, out_mlp_dims, emb_dim, torch.nn.GELU())

    def forward(self, x):
        batch_size = x.size(0)
        x = x.permute((0, 2, 1))
        # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k)
        x = get_graph_feature(x, k=self.n_neighbors)
        # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k)
        x = self.conv1(x)
        # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
        x1 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 64, num_points) -> (batch_size, 64*2, num_points, k)
        x = get_graph_feature(x1, k=self.n_neighbors)
        # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k)
        x = self.conv2(x)
        # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
        x2 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 64, num_points) -> (batch_size, 64*2, num_points, k)
        x = get_graph_feature(x2, k=self.n_neighbors)
        # (batch_size, 64*2, num_points, k) -> (batch_size, 128, num_points, k)
        x = self.conv3(x)
        # (batch_size, 128, num_points, k) -> (batch_size, 128, num_points)
        x3 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 128, num_points) -> (batch_size, 128*2, num_points, k)
        x = get_graph_feature(x3, k=self.n_neighbors)
        # (batch_size, 128*2, num_points, k) -> (batch_size, 256, num_points, k)
        x = self.conv4(x)
        # (batch_size, 256, num_points, k) -> (batch_size, 256, num_points)
        x4 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 64+64+128+256, num_points)
        x = torch.cat((x1, x2, x3, x4), dim=1)

        # (batch_size, 64+64+128+256, num_points) -> (batch_size, emb_dims, num_points)
        x = self.conv5(x)
        # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims)
        x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1)
        # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims)
        x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1)
        x = torch.cat((x1, x2), 1)              # (batch_size, emb_dims*2)

        # (batch_size, 256) -> (batch_size, output_channels)
        x = self.out_mlp(x)

        return x




import os
import sys
import copy
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.nets import MLP
# Original code from : https://github.com/WangYueFt/dgcnn/blob/master/pytorch/model.py


def knn(x, k):
    inner = -2*torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    # (batch_size, num_points, k)
    idx = pairwise_distance.topk(k=k, dim=-1)[1]
    return idx


def get_graph_feature(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(2)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        idx = knn(x, k=k)   # (batch_size, num_points, k)

    idx_base = torch.arange(
        0, batch_size, device=x.device).view(-1, 1, 1)*num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()

    # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    x = x.transpose(2, 1).contiguous()
    feature = x.view(batch_size*num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims)
    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)

    feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2).contiguous()

    return feature


class DGCNNembedder(nn.Module):
    """DGCNN embedder with n_points x embedding dim as output"""
    def __init__(self, out_mlp_dims, emb_dim=22, dropout=0, n_neighbors=20):
        super().__init__()

        self.n_neighbors = n_neighbors

        self.bn1 = nn.BatchNorm2d(64)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(256)
        self.bn5 = nn.BatchNorm1d(512)

        self.conv1 = nn.Sequential(nn.Conv2d(12, 64, kernel_size=1, bias=False),
                                   self.bn1,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
                                   self.bn2,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False),
                                   self.bn3,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False),
                                   self.bn4,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv5 = nn.Sequential(nn.Conv1d(512, 512, kernel_size=1, bias=False),
                                   self.bn5,
                                   nn.LeakyReLU(negative_slope=0.2))

        self.out_mlp = MLP(512, out_mlp_dims, emb_dim, torch.nn.GELU())

    def forward(self, x):
        batch_size = x.size(0)
        x = x.permute((0, 2, 1))
        
    
        x = get_graph_feature(x, k=self.n_neighbors)
        x = self.conv1(x)
        x1 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x1, self.n_neighbors)
        x = self.conv2(x)
        x2 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x2, k=self.n_neighbors)
        x = self.conv3(x)
        x3 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x3, k=self.n_neighbors)
        x = self.conv4(x)
        x4 = x.max(dim=-1, keepdim=False)[0]

        x = torch.cat((x1, x2, x3, x4),dim=1) 

        x = self.conv5(x).permute((0, 2, 1))

        x = self.out_mlp(x)
        return x


class DGCNNembedderGlobal(nn.Module):
    """DGCNN embedder with single embedding vector as output"""
    def __init__(self, input_dim, out_mlp_dims, emb_dim=22, n_neighbors=20):
        super().__init__()

        self.n_neighbors = n_neighbors

        self.input_dim = input_dim

        self.bn1 = nn.BatchNorm2d(64)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(256)
        self.bn5 = nn.BatchNorm1d(512)

        self.conv1 = nn.Sequential(nn.Conv2d(self.input_dim*2, 64, kernel_size=1, bias=False),
                                   self.bn1,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
                                   self.bn2,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False),
                                   self.bn3,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False),
                                   self.bn4,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv5 = nn.Sequential(nn.Conv1d(512, 512, kernel_size=1, bias=False),
                                   self.bn5,
                                   nn.LeakyReLU(negative_slope=0.2))

        self.out_mlp = MLP(512*2, out_mlp_dims, emb_dim, torch.nn.GELU())

    def forward(self, x):
        batch_size = x.size(0)
        x = x.permute((0, 2, 1))
        # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k)
        x = get_graph_feature(x, k=self.n_neighbors)
        # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k)
        x = self.conv1(x)
        # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
        x1 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 64, num_points) -> (batch_size, 64*2, num_points, k)
        x = get_graph_feature(x1, k=self.n_neighbors)
        # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k)
        x = self.conv2(x)
        # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
        x2 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 64, num_points) -> (batch_size, 64*2, num_points, k)
        x = get_graph_feature(x2, k=self.n_neighbors)
        # (batch_size, 64*2, num_points, k) -> (batch_size, 128, num_points, k)
        x = self.conv3(x)
        # (batch_size, 128, num_points, k) -> (batch_size, 128, num_points)
        x3 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 128, num_points) -> (batch_size, 128*2, num_points, k)
        x = get_graph_feature(x3, k=self.n_neighbors)
        # (batch_size, 128*2, num_points, k) -> (batch_size, 256, num_points, k)
        x = self.conv4(x)
        # (batch_size, 256, num_points, k) -> (batch_size, 256, num_points)
        x4 = x.max(dim=-1, keepdim=False)[0]

        # (batch_size, 64+64+128+256, num_points)
        x = torch.cat((x1, x2, x3, x4), dim=1)

        # (batch_size, 64+64+128+256, num_points) -> (batch_size, emb_dims, num_points)
        x = self.conv5(x)
        # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims)
        x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1)
        # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims)
        x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1)
        x = torch.cat((x1, x2), 1)              # (batch_size, emb_dims*2)

        # (batch_size, 256) -> (batch_size, output_channels)
        x = self.out_mlp(x)

        return x
# Original code from : https://github.com/WangYueFt/dgcnn/blob/master/pytorch/model.py # (batch_size, num_points, k) # (batch_size, num_points, k) # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points) DGCNN embedder with n_points x embedding dim as output DGCNN embedder with single embedding vector as output # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k) # (batch_size, 3*2, num_points, k) -> (batch_size, 64, num_points, k) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) # (batch_size, 64, num_points) -> (batch_size, 64*2, num_points, k) # (batch_size, 64*2, num_points, k) -> (batch_size, 64, num_points, k) # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) # (batch_size, 64, num_points) -> (batch_size, 64*2, num_points, k) # (batch_size, 64*2, num_points, k) -> (batch_size, 128, num_points, k) # (batch_size, 128, num_points, k) -> (batch_size, 128, num_points) # (batch_size, 128, num_points) -> (batch_size, 128*2, num_points, k) # (batch_size, 128*2, num_points, k) -> (batch_size, 256, num_points, k) # (batch_size, 256, num_points, k) -> (batch_size, 256, num_points) # (batch_size, 64+64+128+256, num_points) # (batch_size, 64+64+128+256, num_points) -> (batch_size, emb_dims, num_points) # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims) # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims) # (batch_size, emb_dims*2) # (batch_size, 256) -> (batch_size, output_channels)