openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/modeling/transforms/aster_tps.py

@@ -0,0 +1,262 @@
+import itertools
+import math
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def conv3x3_block(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding."""
+    conv_layer = nn.Conv2d(in_planes,
+                           out_planes,
+                           kernel_size=3,
+                           stride=1,
+                           padding=1)
+
+    block = nn.Sequential(
+        conv_layer,
+        nn.BatchNorm2d(out_planes),
+        nn.ReLU(inplace=True),
+    )
+    return block
+
+
+class STNHead(nn.Module):
+
+    def __init__(self, in_planes, num_ctrlpoints, activation='none'):
+        super(STNHead, self).__init__()
+
+        self.in_planes = in_planes
+        self.num_ctrlpoints = num_ctrlpoints
+        self.activation = activation
+        self.stn_convnet = nn.Sequential(
+            conv3x3_block(in_planes, 32),  # 32*64
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(32, 64),  # 16*32
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(64, 128),  # 8*16
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(128, 256),  # 4*8
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(256, 256),  # 2*4,
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(256, 256))  # 1*2
+
+        self.stn_fc1 = nn.Sequential(nn.Linear(2 * 256, 512),
+                                     nn.BatchNorm1d(512),
+                                     nn.ReLU(inplace=True))
+        self.stn_fc2 = nn.Linear(512, num_ctrlpoints * 2)
+
+        self.init_weights(self.stn_convnet)
+        self.init_weights(self.stn_fc1)
+        self.init_stn(self.stn_fc2)
+
+    def init_weights(self, module):
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, 0.001)
+                m.bias.data.zero_()
+
+    def init_stn(self, stn_fc2):
+        margin = 0.01
+        sampling_num_per_side = int(self.num_ctrlpoints / 2)
+        ctrl_pts_x = np.linspace(margin, 1. - margin, sampling_num_per_side)
+        ctrl_pts_y_top = np.ones(sampling_num_per_side) * margin
+        ctrl_pts_y_bottom = np.ones(sampling_num_per_side) * (1 - margin)
+        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)
+        ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom],
+                                     axis=0).astype(np.float32)
+        if self.activation == 'none':
+            pass
+        elif self.activation == 'sigmoid':
+            ctrl_points = -np.log(1. / ctrl_points - 1.)
+        stn_fc2.weight.data.zero_()
+        stn_fc2.bias.data = torch.Tensor(ctrl_points).view(-1)
+
+    def forward(self, x):
+        x = self.stn_convnet(x)
+        batch_size, _, h, w = x.size()
+        x = x.view(batch_size, -1)
+        img_feat = self.stn_fc1(x)
+        x = self.stn_fc2(0.1 * img_feat)
+        if self.activation == 'sigmoid':
+            x = F.sigmoid(x)
+        x = x.view(-1, self.num_ctrlpoints, 2)
+        return x
+
+
+def grid_sample(input, grid, canvas=None):
+    output = F.grid_sample(input, grid)
+    if canvas is None:
+        return output
+    else:
+        input_mask = input.data.new(input.size()).fill_(1)
+        output_mask = F.grid_sample(input_mask, grid)
+        padded_output = output * output_mask + canvas * (1 - output_mask)
+        return padded_output
+
+
+# phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
+def compute_partial_repr(input_points, control_points):
+    N = input_points.size(0)
+    M = control_points.size(0)
+    pairwise_diff = input_points.view(N, 1, 2) - control_points.view(1, M, 2)
+    # original implementation, very slow
+    # pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
+    pairwise_diff_square = pairwise_diff * pairwise_diff
+    pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :,
+                                                                         1]
+    repr_matrix = 0.5 * pairwise_dist * torch.log(pairwise_dist)
+    # fix numerical error for 0 * log(0), substitute all nan with 0
+    mask = repr_matrix != repr_matrix
+    repr_matrix.masked_fill_(mask, 0)
+    return repr_matrix
+
+
+# output_ctrl_pts are specified, according to our task.
+def build_output_control_points(num_control_points, margins):
+    margin_x, margin_y = margins
+    num_ctrl_pts_per_side = num_control_points // 2
+    ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+    ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+    ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+    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)
+    # ctrl_pts_top = ctrl_pts_top[1:-1,:]
+    # ctrl_pts_bottom = ctrl_pts_bottom[1:-1,:]
+    output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom],
+                                         axis=0)
+    output_ctrl_pts = torch.Tensor(output_ctrl_pts_arr)
+    return output_ctrl_pts
+
+
+class TPSSpatialTransformer(nn.Module):
+
+    def __init__(
+        self,
+        output_image_size,
+        num_control_points,
+        margins,
+    ):
+        super(TPSSpatialTransformer, self).__init__()
+        self.output_image_size = output_image_size
+        self.num_control_points = num_control_points
+        self.margins = margins
+
+        self.target_height, self.target_width = output_image_size
+        target_control_points = build_output_control_points(
+            num_control_points, margins)
+        N = num_control_points
+        # N = N - 4
+
+        # create padded kernel matrix
+        forward_kernel = torch.zeros(N + 3, N + 3)
+        target_control_partial_repr = compute_partial_repr(
+            target_control_points, target_control_points)
+        forward_kernel[:N, :N].copy_(target_control_partial_repr)
+        forward_kernel[:N, -3].fill_(1)
+        forward_kernel[-3, :N].fill_(1)
+        forward_kernel[:N, -2:].copy_(target_control_points)
+        forward_kernel[-2:, :N].copy_(target_control_points.transpose(0, 1))
+        # compute inverse matrix
+        inverse_kernel = torch.inverse(forward_kernel)
+
+        # create target cordinate matrix
+        HW = self.target_height * self.target_width
+        target_coordinate = list(
+            itertools.product(range(self.target_height),
+                              range(self.target_width)))
+        target_coordinate = torch.Tensor(target_coordinate)  # HW x 2
+        Y, X = target_coordinate.split(1, dim=1)
+        Y = Y / (self.target_height - 1)
+        X = X / (self.target_width - 1)
+        target_coordinate = torch.cat([X, Y],
+                                      dim=1)  # convert from (y, x) to (x, y)
+        target_coordinate_partial_repr = compute_partial_repr(
+            target_coordinate, target_control_points)
+        target_coordinate_repr = torch.cat([
+            target_coordinate_partial_repr,
+            torch.ones(HW, 1), target_coordinate
+        ],
+                                           dim=1)
+
+        # register precomputed matrices
+        self.register_buffer('inverse_kernel', inverse_kernel)
+        self.register_buffer('padding_matrix', torch.zeros(3, 2))
+        self.register_buffer('target_coordinate_repr', target_coordinate_repr)
+        self.register_buffer('target_control_points', target_control_points)
+
+    def forward(self, input, source_control_points):
+        assert source_control_points.ndimension() == 3
+        assert source_control_points.size(1) == self.num_control_points
+        assert source_control_points.size(2) == 2
+        batch_size = source_control_points.size(0)
+
+        Y = torch.cat([
+            source_control_points,
+            self.padding_matrix.expand(batch_size, 3, 2)
+        ], 1)
+        mapping_matrix = torch.matmul(self.inverse_kernel, Y)
+        source_coordinate = torch.matmul(self.target_coordinate_repr,
+                                         mapping_matrix)
+
+        grid = source_coordinate.view(-1, self.target_height,
+                                      self.target_width, 2)
+        grid = torch.clamp(
+            grid, 0, 1)  # the source_control_points may be out of [0, 1].
+        # the input to grid_sample is normalized [-1, 1], but what we get is [0, 1]
+        grid = 2.0 * grid - 1.0
+        output_maps = grid_sample(input, grid, canvas=None)
+        return output_maps
+
+
+class Aster_TPS(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        tps_inputsize=[32, 64],
+        tps_outputsize=[32, 100],
+        num_control_points=20,
+        tps_margins=[0.05, 0.05],
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        #TODO
+        self.out_channels = in_channels
+        self.tps_inputsize = tps_inputsize
+        self.num_control_points = num_control_points
+
+        self.stn_head = STNHead(
+            in_planes=3,
+            num_ctrlpoints=num_control_points,
+        )
+
+        self.tps = TPSSpatialTransformer(
+            output_image_size=tps_outputsize,
+            num_control_points=num_control_points,
+            margins=tps_margins,
+        )
+
+    def forward(self, img):
+        stn_input = F.interpolate(img,
+                                  self.tps_inputsize,
+                                  mode='bilinear',
+                                  align_corners=True)
+
+        ctrl_points = self.stn_head(stn_input)
+
+        img = self.tps(img, ctrl_points)
+
+        return img

openocr/openrec/modeling/transforms/moran.py

@@ -0,0 +1,136 @@
+"""This code is refer from:
+https://github.com/Canjie-Luo/MORAN_v2
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class MORN(nn.Module):
+
+    def __init__(self, in_channels, target_shape=[32, 100], enhance=1):
+        super(MORN, self).__init__()
+        self.targetH = target_shape[0]
+        self.targetW = target_shape[1]
+        self.enhance = enhance
+        self.out_channels = in_channels
+        self.cnn = nn.Sequential(nn.MaxPool2d(2, 2),
+                                 nn.Conv2d(in_channels, 64, 3, 1, 1),
+                                 nn.BatchNorm2d(64), nn.ReLU(True),
+                                 nn.MaxPool2d(2,
+                                              2), nn.Conv2d(64, 128, 3, 1, 1),
+                                 nn.BatchNorm2d(128), nn.ReLU(True),
+                                 nn.MaxPool2d(2,
+                                              2), nn.Conv2d(128, 64, 3, 1, 1),
+                                 nn.BatchNorm2d(64), nn.ReLU(True),
+                                 nn.Conv2d(64, 16, 3, 1, 1),
+                                 nn.BatchNorm2d(16), nn.ReLU(True),
+                                 nn.Conv2d(16, 1, 3, 1, 1), nn.BatchNorm2d(1))
+        self.pool = nn.MaxPool2d(2, 1)
+        h_list = np.arange(self.targetH) * 2. / (self.targetH - 1) - 1
+        w_list = np.arange(self.targetW) * 2. / (self.targetW - 1) - 1
+        grid = np.meshgrid(w_list, h_list, indexing='ij')
+        grid = np.stack(grid, axis=-1)
+        grid = np.transpose(grid, (1, 0, 2))
+        grid = np.expand_dims(grid, 0)
+        self.grid = nn.Parameter(
+            torch.from_numpy(grid).float(),
+            requires_grad=False,
+        )
+
+    def forward(self, x):
+
+        bs = x.shape[0]
+        grid = self.grid.tile([bs, 1, 1, 1])
+        grid_x = self.grid[:, :, :, 0].unsqueeze(3).tile([bs, 1, 1, 1])
+        grid_y = self.grid[:, :, :, 1].unsqueeze(3).tile([bs, 1, 1, 1])
+        x_small = F.upsample(x,
+                             size=(self.targetH, self.targetW),
+                             mode='bilinear')
+
+        offsets = self.cnn(x_small)
+        offsets_posi = F.relu(offsets, inplace=False)
+        offsets_nega = F.relu(-offsets, inplace=False)
+        offsets_pool = self.pool(offsets_posi) - self.pool(offsets_nega)
+
+        offsets_grid = F.grid_sample(offsets_pool, grid)
+        offsets_grid = offsets_grid.permute(0, 2, 3, 1).contiguous()
+        offsets_x = torch.cat([grid_x, grid_y + offsets_grid], 3)
+        x_rectified = F.grid_sample(x, offsets_x)
+
+        for iteration in range(self.enhance):
+            offsets = self.cnn(x_rectified)
+
+            offsets_posi = F.relu(offsets, inplace=False)
+            offsets_nega = F.relu(-offsets, inplace=False)
+            offsets_pool = self.pool(offsets_posi) - self.pool(offsets_nega)
+
+            offsets_grid += F.grid_sample(offsets_pool,
+                                          grid).permute(0, 2, 3,
+                                                        1).contiguous()
+            offsets_x = torch.cat([grid_x, grid_y + offsets_grid], 3)
+            x_rectified = F.grid_sample(x, offsets_x)
+
+        # if debug:
+
+        #     offsets_mean = torch.mean(offsets_grid.view(x.size(0), -1), 1)
+        #     offsets_max, _ = torch.max(offsets_grid.view(x.size(0), -1), 1)
+        #     offsets_min, _ = torch.min(offsets_grid.view(x.size(0), -1), 1)
+
+        #     import matplotlib.pyplot as plt
+        #     from colour import Color
+        #     from torchvision import transforms
+        #     import cv2
+
+        #     alpha = 0.7
+        #     density_range = 256
+        #     color_map = np.empty([self.targetH, self.targetW, 3], dtype=int)
+        #     cmap = plt.get_cmap("rainbow")
+        #     blue = Color("blue")
+        #     hex_colors = list(blue.range_to(Color("red"), density_range))
+        #     rgb_colors = [[rgb * 255 for rgb in color.rgb] for color in hex_colors][::-1]
+        #     to_pil_image = transforms.ToPILImage()
+
+        #     for i in range(x.size(0)):
+
+        #         img_small = x_small[i].data.cpu().mul_(0.5).add_(0.5)
+        #         img = to_pil_image(img_small)
+        #         img = np.array(img)
+        #         if len(img.shape) == 2:
+        #             img = cv2.merge([img.copy()]*3)
+        #         img_copy = img.copy()
+
+        #         v_max = offsets_max.data[i]
+        #         v_min = offsets_min.data[i]
+        #         if self.cuda:
+        #             img_offsets = (offsets_grid[i]).view(1, self.targetH, self.targetW).data.cuda().add_(-v_min).mul_(1./(v_max-v_min))
+        #         else:
+        #             img_offsets = (offsets_grid[i]).view(1, self.targetH, self.targetW).data.cpu().add_(-v_min).mul_(1./(v_max-v_min))
+        #         img_offsets = to_pil_image(img_offsets)
+        #         img_offsets = np.array(img_offsets)
+        #         color_map = np.empty([self.targetH, self.targetW, 3], dtype=int)
+        #         for h_i in range(self.targetH):
+        #             for w_i in range(self.targetW):
+        #                 color_map[h_i][w_i] = rgb_colors[int(img_offsets[h_i, w_i]/256.*density_range)]
+        #         color_map = color_map.astype(np.uint8)
+        #         cv2.addWeighted(color_map, alpha, img_copy, 1-alpha, 0, img_copy)
+
+        #         img_processed = x_rectified[i].data.cpu().mul_(0.5).add_(0.5)
+        #         img_processed = to_pil_image(img_processed)
+        #         img_processed = np.array(img_processed)
+        #         if len(img_processed.shape) == 2:
+        #             img_processed = cv2.merge([img_processed.copy()]*3)
+
+        #         total_img = np.ones([self.targetH, self.targetW*3+10, 3], dtype=int)*255
+        #         total_img[0:self.targetH, 0:self.targetW] = img
+        #         total_img[0:self.targetH, self.targetW+5:2*self.targetW+5] = img_copy
+        #         total_img[0:self.targetH, self.targetW*2+10:3*self.targetW+10] = img_processed
+        #         total_img = cv2.resize(total_img.astype(np.uint8), (300, 50))
+        #         # cv2.imshow("Input_Offsets_Output", total_img)
+        #         # cv2.waitKey()
+
+        #     return x_rectified, total_img
+
+        return x_rectified

openocr/openrec/modeling/transforms/tps.py

@@ -0,0 +1,246 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from openrec.modeling.common import Activation
+
+
+class ConvBNLayer(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 groups=1,
+                 act=None):
+        super(ConvBNLayer, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            bias=False,
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.act = Activation(act) if act else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.act is not None:
+            x = self.act(x)
+        return x
+
+
+class LocalizationNetwork(nn.Module):
+
+    def __init__(self, in_channels, num_fiducial, loc_lr, model_name):
+        super(LocalizationNetwork, self).__init__()
+        self.F = num_fiducial
+        F = num_fiducial
+        if model_name == 'large':
+            num_filters_list = [64, 128, 256, 512]
+            fc_dim = 256
+        else:
+            num_filters_list = [16, 32, 64, 128]
+            fc_dim = 64
+
+        self.block_list = nn.ModuleList()
+        for fno in range(0, len(num_filters_list)):
+            num_filters = num_filters_list[fno]
+            conv = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=num_filters,
+                kernel_size=3,
+                act='relu',
+            )
+            self.block_list.append(conv)
+            if fno == len(num_filters_list) - 1:
+                pool = nn.AdaptiveAvgPool2d(1)
+            else:
+                pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+            in_channels = num_filters
+            self.block_list.append(pool)
+        self.fc1 = nn.Linear(in_channels, fc_dim)
+
+        # Init fc2 in LocalizationNetwork
+        self.fc2 = nn.Linear(fc_dim, F * 2)
+
+        initial_bias = self.get_initial_fiducials()
+        initial_bias = initial_bias.reshape(-1)
+        self.fc2.bias.data = torch.tensor(initial_bias, dtype=torch.float32)
+        nn.init.zeros_(self.fc2.weight.data)
+        self.out_channels = F * 2
+
+    def forward(self, x):
+        """
+        Estimating parameters of geometric transformation
+        Args:
+            image: input
+        Return:
+            batch_C_prime: the matrix of the geometric transformation
+        """
+        for block in self.block_list:
+            x = block(x)
+        x = x.squeeze(dim=2).squeeze(dim=2)
+        x = self.fc1(x)
+
+        x = F.relu(x)
+        x = self.fc2(x)
+        x = x.reshape(shape=[-1, self.F, 2])
+        return x
+
+    def get_initial_fiducials(self):
+        """see RARE paper Fig.
+
+        6 (a)
+        """
+        F = self.F
+        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)
+        return initial_bias
+
+
+class GridGenerator(nn.Module):
+
+    def __init__(self, in_channels, num_fiducial):
+        super(GridGenerator, self).__init__()
+        self.eps = 1e-6
+        self.F = num_fiducial
+
+        self.fc = nn.Linear(in_channels, 6)
+        nn.init.constant_(self.fc.weight, 0)
+        nn.init.constant_(self.fc.bias, 0)
+        self.fc.weight.requires_grad = False
+        self.fc.bias.requires_grad = False
+
+    def forward(self, batch_C_prime, I_r_size):
+        """Generate the grid for the grid_sampler.
+
+        Args:
+            batch_C_prime: the matrix of the geometric transformation
+            I_r_size: the shape of the input image
+        Return:
+            batch_P_prime: the grid for the grid_sampler
+        """
+        C = self.build_C_paddle()
+        P = self.build_P_paddle(I_r_size)
+
+        inv_delta_C_tensor = self.build_inv_delta_C_paddle(C).float()
+        P_hat_tensor = self.build_P_hat_paddle(C, torch.tensor(P)).float()
+
+        batch_C_ex_part_tensor = self.get_expand_tensor(batch_C_prime)
+
+        batch_C_prime_with_zeros = torch.cat(
+            [batch_C_prime, batch_C_ex_part_tensor], dim=1)
+        batch_T = torch.matmul(
+            inv_delta_C_tensor.to(batch_C_prime_with_zeros.device),
+            batch_C_prime_with_zeros,
+        )
+        batch_P_prime = torch.matmul(P_hat_tensor.to(batch_T.device), batch_T)
+        return batch_P_prime
+
+    def build_C_paddle(self):
+        """Return coordinates of fiducial points in I_r; C."""
+        F = self.F
+        ctrl_pts_x = torch.linspace(-1.0, 1.0, int(F / 2), dtype=torch.float64)
+        ctrl_pts_y_top = -1 * torch.ones([int(F / 2)], dtype=torch.float64)
+        ctrl_pts_y_bottom = torch.ones([int(F / 2)], dtype=torch.float64)
+        ctrl_pts_top = torch.stack([ctrl_pts_x, ctrl_pts_y_top], dim=1)
+        ctrl_pts_bottom = torch.stack([ctrl_pts_x, ctrl_pts_y_bottom], dim=1)
+        C = torch.cat([ctrl_pts_top, ctrl_pts_bottom], dim=0)
+        return C  # F x 2
+
+    def build_P_paddle(self, I_r_size):
+        I_r_height, I_r_width = I_r_size
+        I_r_grid_x = (torch.arange(-I_r_width, I_r_width, 2) +
+                      1.0) / torch.tensor(np.array([I_r_width]))
+
+        I_r_grid_y = (torch.arange(-I_r_height, I_r_height, 2) +
+                      1.0) / torch.tensor(np.array([I_r_height]))
+
+        # P: self.I_r_width x self.I_r_height x 2
+        P = torch.stack(torch.meshgrid(I_r_grid_x, I_r_grid_y), dim=2)
+        P = torch.permute(P, [1, 0, 2])
+        # n (= self.I_r_width x self.I_r_height) x 2
+        return P.reshape([-1, 2])
+
+    def build_inv_delta_C_paddle(self, C):
+        """Return inv_delta_C which is needed to calculate T."""
+        F = self.F
+        hat_eye = torch.eye(F)  # F x F
+        hat_C = torch.norm(C.reshape([1, F, 2]) - C.reshape([F, 1, 2]),
+                           dim=2) + hat_eye
+        hat_C = (hat_C**2) * torch.log(hat_C)
+        delta_C = torch.cat(  # F+3 x F+3
+            [
+                torch.cat([torch.ones((F, 1)), C, hat_C], dim=1),  # F x F+3
+                torch.concat([torch.zeros(
+                    (2, 3)), C.transpose(0, 1)], dim=1),  # 2 x F+3
+                torch.concat([torch.zeros(
+                    (1, 3)), torch.ones((1, F))], dim=1),  # 1 x F+3
+            ],
+            axis=0,
+        )
+        inv_delta_C = torch.inverse(delta_C)
+        return inv_delta_C  # F+3 x F+3
+
+    def build_P_hat_paddle(self, C, P):
+        F = self.F
+        eps = self.eps
+        n = P.shape[0]  # n (= self.I_r_width x self.I_r_height)
+        # P_tile: n x 2 -> n x 1 x 2 -> n x F x 2
+        P_tile = torch.tile(torch.unsqueeze(P, dim=1), (1, F, 1))
+        C_tile = torch.unsqueeze(C, dim=0)  # 1 x F x 2
+        P_diff = P_tile - C_tile  # n x F x 2
+        # rbf_norm: n x F
+        rbf_norm = torch.norm(P_diff, p=2, dim=2, keepdim=False)
+
+        # rbf: n x F
+        rbf = torch.multiply(torch.square(rbf_norm), torch.log(rbf_norm + eps))
+        P_hat = torch.cat([torch.ones((n, 1)), P, rbf], dim=1)
+        return P_hat  # n x F+3
+
+    def get_expand_tensor(self, batch_C_prime):
+        B, H, C = batch_C_prime.shape
+        batch_C_prime = batch_C_prime.reshape([B, H * C])
+        batch_C_ex_part_tensor = self.fc(batch_C_prime)
+        batch_C_ex_part_tensor = batch_C_ex_part_tensor.reshape([-1, 3, 2])
+        return batch_C_ex_part_tensor
+
+
+class TPS(nn.Module):
+
+    def __init__(self, in_channels, num_fiducial, loc_lr, model_name):
+        super(TPS, self).__init__()
+        self.loc_net = LocalizationNetwork(in_channels, num_fiducial, loc_lr,
+                                           model_name)
+        self.grid_generator = GridGenerator(self.loc_net.out_channels,
+                                            num_fiducial)
+        self.out_channels = in_channels
+
+    def forward(self, image):
+        image.stop_gradient = False
+        batch_C_prime = self.loc_net(image)
+        batch_P_prime = self.grid_generator(batch_C_prime, image.shape[2:])
+        batch_P_prime = batch_P_prime.reshape(
+            [-1, image.shape[2], image.shape[3], 2])
+        is_fp16 = False
+        if batch_P_prime.dtype != torch.float32:
+            data_type = batch_P_prime.dtype
+            image = image.float()
+            batch_P_prime = batch_P_prime.float()
+            is_fp16 = True
+        batch_I_r = F.grid_sample(image, grid=batch_P_prime)
+        if is_fp16:
+            batch_I_r = batch_I_r.astype(data_type)
+
+        return batch_I_r