text: from enum import EnumMeta
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  import numpy as np
  import copy
  
  class NetVLAD(nn.Module):
      """NetVLAD layer implementation"""
  
      def __init__(self, num_clusters=64, dim=512, alpha=100.0, normalize_input=True):
          """
          Args:
              num_clusters : int
                  The number of clusters
              dim : int
                  Dimension of descriptors
              alpha : float
                  Parameter of initialization. Larger value is harder assignment.
              normalize_input : bool
                  If true, descriptor-wise L2 normalization is applied to input.
          """
          super(NetVLAD, self).__init__()
          self.num_clusters = num_clusters
          self.dim = dim
          self.alpha = alpha
          self.normalize_input = normalize_input
          self.conv = nn.Conv2d(dim, num_clusters, kernel_size=(1, 1), bias=False)
          self.centroids = nn.Parameter(torch.rand(num_clusters, dim), requires_grad=True)
  
          self.clsts = None
          self.traindescs = None
  
      def _init_params(self):
          clstsAssign = self.clsts / np.linalg.norm(self.clsts, axis=1, keepdims=True)
          dots = np.dot(clstsAssign, self.traindescs.T)
          dots.sort(0)
          dots = dots[::-1, :] # sort, descending
  
          self.alpha = (-np.log(0.01) / np.mean(dots[0,:] - dots[1,:])).item()
          self.centroids.data.copy_(torch.from_numpy(self.clsts))
          self.conv.weight.data.copy_(torch.from_numpy(self.alpha*clstsAssign).unsqueeze(2).unsqueeze(3))
  
      def forward(self, x):
          N, C = x.shape[:2]
          if self.normalize_input:
              x = F.normalize(x, p=2, dim=1)  # across descriptor dim
  
          # soft-assignment
          soft_assign = self.conv(x).view(N, self.num_clusters, -1)
          soft_assign = F.softmax(soft_assign, dim=1)
  
          x_flatten = x.view(N, C, -1)
  
          # calculate residuals to each clusters in one loop
          residual = x_flatten.expand(self.num_clusters, -1, -1, -1).permute(1, 0, 2, 3) - \
              self.centroids.expand(x_flatten.size(-1), -1, -1).permute(1, 2, 0).unsqueeze(0)
          residual *= soft_assign.unsqueeze(2)
          vlad = residual.sum(dim=-1)
  
          return vlad
  
      def forward_patch(self,x):
          N, C, H, W = x.shape
          x = self.forward(x)
          # [IMPORTANT] normalize
          x = F.normalize(x, p=2, dim=2)  # intra-normalization
          x = x.view(N, -1)  # flatten
          x = F.normalize(x, p=2, dim=1)  # L2 normalize
          return x
  
  class EmbedNet(nn.Module):
      def __init__(self, base_model, net_vlad, region_sim_strategy='netvlad'):
          super(EmbedNet, self).__init__()
          self.base_model = base_model
          self.net_vlad = net_vlad
          self.region_sim_strategy = region_sim_strategy
  
      def _init_params(self):
          self.base_model._init_params()
          self.net_vlad._init_params()
  
      def forward(self, x):
          pool_x, x = self.base_model(x)
          vlad_x = self.net_vlad(x)
  
          # [IMPORTANT] normalize
          if self.region_sim_strategy == 'netvlad':
              # normalize
              vlad_x = F.normalize(vlad_x, p=2, dim=2)  # intra-normalization
              vlad_x = vlad_x.view(x.size(0), -1)  # flatten
              vlad_x = F.normalize(vlad_x, p=2, dim=1)  # L2 normalize
  
          return pool_x, vlad_x
  
  class EmbedNetPCA(nn.Module):
      def __init__(self, base_model, net_vlad, dim=4096, region_sim_strategy='netvlad'):
          super(EmbedNetPCA, self).__init__()
          self.base_model = base_model
          self.net_vlad = net_vlad
          self.pca_layer = nn.Conv2d(net_vlad.num_clusters*net_vlad.dim, dim, 1, stride=1, padding=0)
          self.region_sim_strategy = region_sim_strategy
  
      def _init_params(self):
          self.base_model._init_params()
          self.net_vlad._init_params()
  
      def forward(self, x):
          _, x = self.base_model(x)
          vlad_x = self.net_vlad(x)
  
          if self.region_sim_strategy == 'netvlad':
              # normalize
              vlad_x = F.normalize(vlad_x, p=2, dim=2)  # intra-normalization
              vlad_x = vlad_x.view(x.size(0), -1)  # flatten
              vlad_x = F.normalize(vlad_x, p=2, dim=1)  # L2 normalize
  
          # reduction
          N, D = vlad_x.size()
          vlad_x = vlad_x.view(N, D, 1, 1)
          vlad_x = self.pca_layer(vlad_x).view(N, -1)
          vlad_x = F.normalize(vlad_x, p=2, dim=-1)  # L2 normalize
  
          return vlad_x
  
  class EmbedRegionNet(nn.Module):
      def __init__(self, base_model, net_vlad, tuple_size=1, region_sim_strategy='netvlad'):
          super(EmbedRegionNet, self).__init__()
          self.base_model = base_model
          self.net_vlad = net_vlad
          self.tuple_size = tuple_size
          self.region_sim_strategy = region_sim_strategy
  
      def _init_params(self):
          self.base_model._init_params()
          self.net_vlad._init_params()
          
      def compute_region_sim(self, feature_A, feature_B):
          # if self.region_sim_strategy == 'netvlad':
          #     score, vlad_A, vlad_B = self._compute_region_sim(feature_A, feature_B)
          # else:
          score, vlad_A, vlad_B = self._compute_region_sim_alter(feature_A, feature_B)
          return score, vlad_A, vlad_B
  
      def _compute_region_sim(self, feature_A, feature_B):
          # feature_A: B*C*H*W
          # feature_B: (B*(1+neg_num))*C*H*W
  
          def reshape(x):
              # re-arrange local features for aggregating quarter regions
              N, C, H, W = x.size()
              x = x.view(N, C, 2, int(H/2), 2, int(W/2))
              x = x.permute(0,1,2,4,3,5).contiguous() # N, C, 2, 2, H/2, W/2
              x = x.view(N, C, -1, int(H/2), int(W/2)) # N, C, 4, H/2, W/2
              return x
  
          feature_A = reshape(feature_A) # [1, 512, 4, 15, 20]
          feature_B = reshape(feature_B) # [11, 512, 4, 15, 20]
  
          # compute quarter-region features
          def aggregate_quarter(x):
              N, C, B, H, W = x.size()
              x = x.permute(0,2,1,3,4).contiguous()
              x = x.view(-1,C,H,W)
              vlad_x = self.net_vlad(x) # (N*B)*64*512
              _, cluster_num, feat_dim = vlad_x.size()
              vlad_x = vlad_x.view(N,B,cluster_num,feat_dim)
              return vlad_x
  
          vlad_A_quarter = aggregate_quarter(feature_A) # [1,4,64,512]
          vlad_B_quarter = aggregate_quarter(feature_B) # [11,4,64,512]
  
          # compute half-region features
          def quarter_to_half(vlad_x):
              return torch.stack((vlad_x[:,0]+vlad_x[:,1], vlad_x[:,2]+vlad_x[:,3], \
                                  vlad_x[:,0]+vlad_x[:,2], vlad_x[:,1]+vlad_x[:,3]), dim=1).contiguous()
  
          vlad_A_half = quarter_to_half(vlad_A_quarter) # [1,4,64,512]
          vlad_B_half = quarter_to_half(vlad_B_quarter) # [11,4,64,512]
  
          # compute global-image features
          def quarter_to_global(vlad_x):
              return vlad_x.sum(1).unsqueeze(1).contiguous()
  
          vlad_A_global = quarter_to_global(vlad_A_quarter) # [1,1,64,512]
          vlad_B_global = quarter_to_global(vlad_B_quarter) # [11,1,64,512]
  
          def norm(vlad_x):
              N, B, C, _ = vlad_x.size()
              vlad_x = F.normalize(vlad_x, p=2, dim=3)  # intra-normalization
              vlad_x = vlad_x.view(N, B, -1)  # flatten
              vlad_x = F.normalize(vlad_x, p=2, dim=2)  # L2 normalize
              return vlad_x
  
          vlad_A = torch.cat((vlad_A_global, vlad_A_half, vlad_A_quarter), dim=1) 
          vlad_B = torch.cat((vlad_B_global, vlad_B_half, vlad_B_quarter), dim=1) 
          vlad_A = norm(vlad_A) # [1, 9, 32768]
          vlad_B = norm(vlad_B) # [11, 9, 32768]
  
          _, B, L = vlad_B.size() # 9,32768 ##9=1+4+4
          vlad_A = vlad_A.view(self.tuple_size,-1,B,L)
          vlad_B = vlad_B.view(self.tuple_size,-1,B,L)
  
          score = torch.bmm(vlad_A.expand_as(vlad_B).view(-1,B,L), vlad_B.view(-1,B,L).transpose(1,2)) # [11,9,9]
          score = score.view(self.tuple_size,-1,B,B)
  
          return score, vlad_A, vlad_B
  
  
      def _compute_region_sim_alter(self, feature_A, feature_B):
          # feature_A: B*C*H*W
          # feature_B: (B*(1+neg_num))*C*H*W
  
          def scale_encode(x,scaleh,scalew):
              # re-arrange local features for aggregating quarter regions
              N, C, H, W = x.size()
              x = x.view(N, C, scaleh, int(H/scaleh), scalew, int(W/scalew))
              x = x.permute(0,1,2,4,3,5).contiguous() # N, C, scaleh, scalew, H/scaleh, W/scalew
              x = x.view(N, C, -1, int(H/scaleh), int(W/scalew)) # N, C, scaleh*scalew, H/scaleh, W/scalew
              N, C, B, H, W = x.size()
              x = x.permute(0,2,1,3,4).contiguous() # N, B, C, H', W'
              ## comment for replacement
              x = x.view(-1,C,H,W) # N*B, C, H', W'
              vlad_x = self.net_vlad.forward_patch(x) # (N*B)*64*512
              _, des_dim = vlad_x.size()
              vlad_x = vlad_x.view(N,B,des_dim) # N, B, 64*512
              ## 2022/03/22 split forward to save memory cost
              # for i in range(B):
              #     vlad_x_i = self.net_vlad.forward_patch(x[:,i,:,:,:]).unsqueeze(1)
              #     if i == 0:
              #         vlad_x = vlad_x_i
              #     else:
              #         vlad_x = torch.cat((vlad_x,vlad_x_i),1)
              return vlad_x
  
          # def scale_encode_from_feat(feature_A):
          #     vlad_11 = scale_encode(feature_A,1,1)
          #     vlad_21 = scale_encode(feature_A,2,1)
          #     vlad_12 = scale_encode(feature_A,1,2)
          #     vlad_22 = scale_encode(feature_A,2,2)
          #     vlad = torch.cat((vlad_11,vlad_21,vlad_12,vlad_22),dim=1)
          #     return vlad
  
          # vlad_A = scale_encode_from_feat(feature_A)
          # vlad_B = scale_encode_from_feat(feature_B)
  
          vlad_A = torch.cat((scale_encode(feature_A,1,1), scale_encode(feature_A,2,1), scale_encode(feature_A,1,2), scale_encode(feature_A,2,2)), dim=1)
          vlad_B = torch.cat((scale_encode(feature_B,1,1), scale_encode(feature_B,2,1), scale_encode(feature_B,1,2), scale_encode(feature_B,2,2)), dim=1)
  
          _, B, L = vlad_B.size() # 9,32768 ##9=1+4+4
          vlad_A = vlad_A.view(self.tuple_size,-1,B,L)
          vlad_B = vlad_B.view(self.tuple_size,-1,B,L)
  
          score = torch.bmm(vlad_A.expand_as(vlad_B).view(-1,B,L), vlad_B.view(-1,B,L).transpose(1,2)) # [11,9,9]
          score = score.view(self.tuple_size,-1,B,B)
  
          return score, vlad_A, vlad_B
  
  
      def _forward_train(self, x):
          B, C, H, W = x.size()
          x = x.view(self.tuple_size, -1, C, H, W)
  
          anchors = x[:, 0].unsqueeze(1).contiguous().view(-1,C,H,W) # B*C*H*W
          pairs = x[:, 1:].view(-1,C,H,W) # (B*(1+neg_num))*C*H*W
  
          # return self._compute_region_sim(anchors, pairs)
          return self.compute_region_sim(anchors, pairs)
  
      def forward(self, x):
          pool_x, x = self.base_model(x)
          if (not self.training):
              vlad_x = self.net_vlad(x)
              if self.region_sim_strategy == 'netvlad':
                  # normalize
                  vlad_x = F.normalize(vlad_x, p=2, dim=2)  # intra-normalization
                  vlad_x = vlad_x.view(x.size(0), -1)  # flatten
                  vlad_x = F.normalize(vlad_x, p=2, dim=1)  # L2 normalize
              return pool_x, vlad_x
  
          return self._forward_train(x)
  