text: #!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  """
  _____________________________________________________________________________
  Created By  : Nguyen Minh Trang - Trangnm5
  Created Date: Mon November 03 10:00:00 VNT 2020
  Project : AkaOCR core
  _____________________________________________________________________________
  
  This file contain transformation layer of model recognition
  _____________________________________________________________________________
  """
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  class TPSSpatialTransformerNetwork(nn.Module):
      """
      Rectification Network of RARE, namely TPS based STN
      """
      def __init__(self, F, I_size, I_r_size, I_channel_num=1, device=torch.device('cpu')):
          """Based on RARE TPS
  
          Parameters
          ----------
          F : int
              number of fiducial (in opt)
          I_size : tuple
              (height, width) of the input image I
          I_r_size : tuple
              (height, width) of the rectified image I_r
          I_channel_num : int
               the number of channels of the input image I
          """
          super(TPSSpatialTransformerNetwork, self).__init__()
          self.F = F
          self.I_size = I_size
          self.I_r_size = I_r_size  # = (I_r_height, I_r_width)
          self.I_channel_num = I_channel_num
          self.LocalizationNetwork = LocalizationNetwork(self.F, self.I_channel_num)
          self.GridGenerator = GridGenerator(self.F, self.I_r_size, device=device)
  
      def forward(self, batch_I):
          """
          Parameters
          ----------
          batch_I : Tensor
              Batch Input Image [batch_size x I_channel_num x I_height x I_width]
  
          Returns
          ------
          Tensor
              rectified image [batch_size x I_channel_num x I_r_height x I_r_width]
          """
          batch_C_prime = self.LocalizationNetwork(batch_I)  # batch_size x K x 2
          build_P_prime = self.GridGenerator.build_P_prime(batch_C_prime)  # batch_size x n (= I_r_width x I_r_height) x 2
          build_P_prime_reshape = build_P_prime.reshape([build_P_prime.size(0), self.I_r_size[0], self.I_r_size[1], 2])
          
          # Fix error "Type mismatch error with torch.nn.functional.grid_sample() under AMP" (https://github.com/pytorch/pytorch/issues/42218)
          # Pytorch fixed in 1.9.0, remove below line (.float()) when run on torch 1.9.0
          build_P_prime_reshape = build_P_prime_reshape.float()  
      
          if torch.__version__ > "1.2.0":
              batch_I_r = F.grid_sample(batch_I, build_P_prime_reshape, padding_mode='border', align_corners=True)
          else:
              batch_I_r = F.grid_sample(batch_I, build_P_prime_reshape, padding_mode='border')
  
          return batch_I_r
  
  
  class LocalizationNetwork(nn.Module):
      """
      Localization Network of RARE, which predicts C' (K x 2) from I (I_width x I_height)
      """
      def __init__(self, F, I_channel_num):
          super(LocalizationNetwork, self).__init__()
          self.F = F
          self.I_channel_num = I_channel_num
          self.conv = nn.Sequential(
              nn.Conv2d(in_channels=self.I_channel_num, out_channels=64, kernel_size=3, stride=1, padding=1,
                        bias=False), nn.BatchNorm2d(64), nn.ReLU(True),
              nn.MaxPool2d(2, 2),  # batch_size x 64 x I_height/2 x I_width/2
              nn.Conv2d(64, 128, 3, 1, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(True),
              nn.MaxPool2d(2, 2),  # batch_size x 128 x I_height/4 x I_width/4
              nn.Conv2d(128, 256, 3, 1, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(True),
              nn.MaxPool2d(2, 2),  # batch_size x 256 x I_height/8 x I_width/8
              nn.Conv2d(256, 512, 3, 1, 1, bias=False), nn.BatchNorm2d(512), nn.ReLU(True),
              nn.AdaptiveAvgPool2d(1)  # batch_size x 512
          )
  
          self.localization_fc1 = nn.Sequential(nn.Linear(512, 256), nn.ReLU(True))
          self.localization_fc2 = nn.Linear(256, self.F * 2)
  
          # Init fc2 in LocalizationNetwork
          self.localization_fc2.weight.data.fill_(0)
          """ see RARE paper Fig. 6 (a) """
          ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
          ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2))
          ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2))
          ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
          ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
          initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
          self.localization_fc2.bias.data = torch.from_numpy(initial_bias).float().view(-1)
  
      def forward(self, batch_I):
          """
          Parameters
          ----------
          batch_I : Tensor
              Batch Input Image [batch_size x I_channel_num x I_height x I_width]
  
          Returns
          ------
          Tensor
              Predicted coordinates of fiducial points for input batch [batch_size x F x 2]
          """
          batch_size = batch_I.size(0)
          features = self.conv(batch_I).view(batch_size, -1)
          batch_C_prime = self.localization_fc2(self.localization_fc1(features)).view(batch_size, self.F, 2)
          return batch_C_prime
  
  
  class GridGenerator(nn.Module):
      """
      Grid Generator of RARE, which produces P_prime by multipling T with P
      """
      def __init__(self, F, I_r_size, device):
          """Generate P_hat and inv_delta_C for later"""
          super(GridGenerator, self).__init__()
          self.device = device
          self.eps = 1e-6
          self.I_r_height, self.I_r_width = I_r_size
          self.F = F
          self.C = self._build_C(self.F)  # F x 2
          self.P = self._build_P(self.I_r_width, self.I_r_height)
          ## for multi-gpu, you need register buffer
          self.register_buffer("inv_delta_C", torch.tensor(self._build_inv_delta_C(self.F, self.C)).float())  # F+3 x F+3
          self.register_buffer("P_hat", torch.tensor(self._build_P_hat(self.F, self.C, self.P)).float())  # n x F+3
          ## for fine-tuning with different image width, you may use below instead of self.register_buffer
          #self.inv_delta_C = torch.tensor(self._build_inv_delta_C(self.F, self.C)).float().cuda()  # F+3 x F+3
          #self.P_hat = torch.tensor(self._build_P_hat(self.F, self.C, self.P)).float().cuda()  # n x F+3
  
      @staticmethod
      def _build_C(F):
          """
          Returns
          ------
          Tensor
              coordinates of fiducial points in I_r
          """
          ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
          ctrl_pts_y_top = -1 * np.ones(int(F / 2))
          ctrl_pts_y_bottom = np.ones(int(F / 2))
          ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
          ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
          C = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
          return C  # F x 2
  
      @staticmethod
      def _build_inv_delta_C(F, C):
          """
          Returns
          ------
          Tensor
              inv_delta_C which is needed to calculate T
          """
          hat_C = np.zeros((F, F), dtype=float)  # F x F
          for i in range(0, F):
              for j in range(i, F):
                  r = np.linalg.norm(C[i] - C[j])
                  hat_C[i, j] = r
                  hat_C[j, i] = r
          np.fill_diagonal(hat_C, 1)
          hat_C = (hat_C ** 2) * np.log(hat_C)
          # print(C.shape, hat_C.shape)
          delta_C = np.concatenate(  # F+3 x F+3
              [
                  np.concatenate([np.ones((F, 1)), C, hat_C], axis=1),  # F x F+3
                  np.concatenate([np.zeros((2, 3)), np.transpose(C)], axis=1),  # 2 x F+3
                  np.concatenate([np.zeros((1, 3)), np.ones((1, F))], axis=1)  # 1 x F+3
              ],
              axis=0
          )
          inv_delta_C = np.linalg.inv(delta_C)
          return inv_delta_C  # F+3 x F+3
  
      @staticmethod
      def _build_P(I_r_width, I_r_height):
          I_r_grid_x = (np.arange(-I_r_width, I_r_width, 2) + 1.0) / I_r_width  # self.I_r_width
          I_r_grid_y = (np.arange(-I_r_height, I_r_height, 2) + 1.0) / I_r_height  # self.I_r_height
          P = np.stack(  # self.I_r_width x self.I_r_height x 2
              np.meshgrid(I_r_grid_x, I_r_grid_y),
              axis=2
          )
          return P.reshape([-1, 2])  # n (= self.I_r_width x self.I_r_height) x 2
  
      def _build_P_hat(self, F, C, P):
          n = P.shape[0]  # n (= self.I_r_width x self.I_r_height)
          P_tile = np.tile(np.expand_dims(P, axis=1), (1, F, 1))  # n x 2 -> n x 1 x 2 -> n x F x 2
          C_tile = np.expand_dims(C, axis=0)  # 1 x F x 2
          P_diff = P_tile - C_tile  # n x F x 2
          rbf_norm = np.linalg.norm(P_diff, ord=2, axis=2, keepdims=False)  # n x F
          rbf = np.multiply(np.square(rbf_norm), np.log(rbf_norm + self.eps))  # n x F
          P_hat = np.concatenate([np.ones((n, 1)), P, rbf], axis=1)
          return P_hat  # n x F+3
  
      def build_P_prime(self, batch_C_prime):
          """
          Returns
          ------
          Tensor
              Generated Grid from batch_C_prime [batch_size x F x 2]
          """
          batch_size = batch_C_prime.size(0)
          batch_inv_delta_C = self.inv_delta_C.repeat(batch_size, 1, 1)
          batch_P_hat = self.P_hat.repeat(batch_size, 1, 1)
          batch_C_prime_with_zeros = torch.cat((batch_C_prime, torch.zeros(
              batch_size, 3, 2).float().to(self.device)), dim=1)  # batch_size x F+3 x 2
          batch_T = torch.bmm(batch_inv_delta_C, batch_C_prime_with_zeros)  # batch_size x F+3 x 2
          batch_P_prime = torch.bmm(batch_P_hat, batch_T)  # batch_size x n x 2
          return batch_P_prime  # batch_size x n x 2
  