# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import from __future__ import division from __future__ import print_function import math import paddle from paddle import nn, ParamAttr from paddle.nn import functional as F import numpy as np import itertools import functools from .tps import GridGenerator '''This code is refer from: https://github.com/hikopensource/DAVAR-Lab-OCR/davarocr/davar_rcg/models/transformations/gaspin_transformation.py ''' class SP_TransformerNetwork(nn.Layer): """ Sturture-Preserving Transformation (SPT) as Equa. (2) in Ref. [1] Ref: [1] SPIN: Structure-Preserving Inner Offset Network for Scene Text Recognition. AAAI-2021. """ def __init__(self, nc=1, default_type=5): """ Based on SPIN Args: nc (int): number of input channels (usually in 1 or 3) default_type (int): the complexity of transformation intensities (by default set to 6 as the paper) """ super(SP_TransformerNetwork, self).__init__() self.power_list = self.cal_K(default_type) self.sigmoid = nn.Sigmoid() self.bn = nn.InstanceNorm2D(nc) def cal_K(self, k=5): """ Args: k (int): the complexity of transformation intensities (by default set to 6 as the paper) Returns: List: the normalized intensity of each pixel in [0,1], denoted as \beta [1x(2K+1)] """ from math import log x = [] if k != 0: for i in range(1, k+1): lower = round(log(1-(0.5/(k+1))*i)/log((0.5/(k+1))*i), 2) upper = round(1/lower, 2) x.append(lower) x.append(upper) x.append(1.00) return x def forward(self, batch_I, weights, offsets, lambda_color=None): """ Args: batch_I (torch.Tensor): batch of input images [batch_size x nc x I_height x I_width] weights: offsets: the predicted offset by AIN, a scalar lambda_color: the learnable update gate \alpha in Equa. (5) as g(x) = (1 - \alpha) \odot x + \alpha \odot x_{offsets} Returns: torch.Tensor: transformed images by SPN as Equa. (4) in Ref. [1] [batch_size x I_channel_num x I_r_height x I_r_width] """ batch_I = (batch_I + 1) * 0.5 if offsets is not None: batch_I = batch_I*(1-lambda_color) + offsets*lambda_color batch_weight_params = paddle.unsqueeze(paddle.unsqueeze(weights, -1), -1) batch_I_power = paddle.stack([batch_I.pow(p) for p in self.power_list], axis=1) batch_weight_sum = paddle.sum(batch_I_power * batch_weight_params, axis=1) batch_weight_sum = self.bn(batch_weight_sum) batch_weight_sum = self.sigmoid(batch_weight_sum) batch_weight_sum = batch_weight_sum * 2 - 1 return batch_weight_sum class GA_SPIN_Transformer(nn.Layer): """ Geometric-Absorbed SPIN Transformation (GA-SPIN) proposed in Ref. [1] Ref: [1] SPIN: Structure-Preserving Inner Offset Network for Scene Text Recognition. AAAI-2021. """ def __init__(self, in_channels=1, I_r_size=(32, 100), offsets=False, norm_type='BN', default_type=6, loc_lr=1, stn=True): """ Args: in_channels (int): channel of input features, set it to 1 if the grayscale images and 3 if RGB input I_r_size (tuple): size of rectified images (used in STN transformations) inputDataType (str): the type of input data, only support 'torch.cuda.FloatTensor' this version offsets (bool): set it to False if use SPN w.o. AIN, and set it to True if use SPIN (both with SPN and AIN) norm_type (str): the normalization type of the module, set it to 'BN' by default, 'IN' optionally default_type (int): the K chromatic space, set it to 3/5/6 depend on the complexity of transformation intensities loc_lr (float): learning rate of location network """ super(GA_SPIN_Transformer, self).__init__() self.nc = in_channels self.spt = True self.offsets = offsets self.stn = stn # set to True in GA-SPIN, while set it to False in SPIN self.I_r_size = I_r_size self.out_channels = in_channels if norm_type == 'BN': norm_layer = functools.partial(nn.BatchNorm2D, use_global_stats=True) elif norm_type == 'IN': norm_layer = functools.partial(nn.InstanceNorm2D, weight_attr=False, use_global_stats=False) else: raise NotImplementedError('normalization layer [%s] is not found' % norm_type) if self.spt: self.sp_net = SP_TransformerNetwork(in_channels, default_type) self.spt_convnet = nn.Sequential( # 32*100 nn.Conv2D(in_channels, 32, 3, 1, 1, bias_attr=False), norm_layer(32), nn.ReLU(), nn.MaxPool2D(kernel_size=2, stride=2), # 16*50 nn.Conv2D(32, 64, 3, 1, 1, bias_attr=False), norm_layer(64), nn.ReLU(), nn.MaxPool2D(kernel_size=2, stride=2), # 8*25 nn.Conv2D(64, 128, 3, 1, 1, bias_attr=False), norm_layer(128), nn.ReLU(), nn.MaxPool2D(kernel_size=2, stride=2), # 4*12 ) self.stucture_fc1 = nn.Sequential( nn.Conv2D(128, 256, 3, 1, 1, bias_attr=False), norm_layer(256), nn.ReLU(), nn.MaxPool2D(kernel_size=2, stride=2), nn.Conv2D(256, 256, 3, 1, 1, bias_attr=False), norm_layer(256), nn.ReLU(), # 2*6 nn.MaxPool2D(kernel_size=2, stride=2), nn.Conv2D(256, 512, 3, 1, 1, bias_attr=False), norm_layer(512), nn.ReLU(), # 1*3 nn.AdaptiveAvgPool2D(1), nn.Flatten(1, -1), # batch_size x 512 nn.Linear(512, 256, weight_attr=nn.initializer.Normal(0.001)), nn.BatchNorm1D(256), nn.ReLU() ) self.out_weight = 2*default_type+1 self.spt_length = 2*default_type+1 if offsets: self.out_weight += 1 if self.stn: self.F = 20 self.out_weight += self.F * 2 self.GridGenerator = GridGenerator(self.F*2, self.F) # self.out_weight*=nc # Init structure_fc2 in LocalizationNetwork initial_bias = self.init_spin(default_type*2) initial_bias = initial_bias.reshape(-1) param_attr = ParamAttr( learning_rate=loc_lr, initializer=nn.initializer.Assign(np.zeros([256, self.out_weight]))) bias_attr = ParamAttr( learning_rate=loc_lr, initializer=nn.initializer.Assign(initial_bias)) self.stucture_fc2 = nn.Linear(256, self.out_weight, weight_attr=param_attr, bias_attr=bias_attr) self.sigmoid = nn.Sigmoid() if offsets: self.offset_fc1 = nn.Sequential(nn.Conv2D(128, 16, 3, 1, 1, bias_attr=False), norm_layer(16), nn.ReLU(),) self.offset_fc2 = nn.Conv2D(16, in_channels, 3, 1, 1) self.pool = nn.MaxPool2D(2, 2) def init_spin(self, nz): """ Args: nz (int): number of paired \betas exponents, which means the value of K x 2 """ init_id = [0.00]*nz+[5.00] if self.offsets: init_id += [-5.00] # init_id *=3 init = np.array(init_id) if self.stn: 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) initial_bias = initial_bias.reshape(-1) init = np.concatenate([init, initial_bias], axis=0) return init def forward(self, x, return_weight=False): """ Args: x (torch.cuda.FloatTensor): input image batch return_weight (bool): set to False by default, if set to True return the predicted offsets of AIN, denoted as x_{offsets} Returns: torch.Tensor: rectified image [batch_size x I_channel_num x I_height x I_width], the same as the input size """ if self.spt: feat = self.spt_convnet(x) fc1 = self.stucture_fc1(feat) sp_weight_fusion = self.stucture_fc2(fc1) sp_weight_fusion = sp_weight_fusion.reshape([x.shape[0], self.out_weight, 1]) if self.offsets: # SPIN w. AIN lambda_color = sp_weight_fusion[:, self.spt_length, 0] lambda_color = self.sigmoid(lambda_color).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) sp_weight = sp_weight_fusion[:, :self.spt_length, :] offsets = self.pool(self.offset_fc2(self.offset_fc1(feat))) assert offsets.shape[2] == 2 # 2 assert offsets.shape[3] == 6 # 16 offsets = self.sigmoid(offsets) # v12 if return_weight: return offsets offsets = nn.functional.upsample(offsets, size=(x.shape[2], x.shape[3]), mode='bilinear') if self.stn: batch_C_prime = sp_weight_fusion[:, (self.spt_length + 1):, :].reshape([x.shape[0], self.F, 2]) build_P_prime = self.GridGenerator(batch_C_prime, self.I_r_size) build_P_prime_reshape = build_P_prime.reshape([build_P_prime.shape[0], self.I_r_size[0], self.I_r_size[1], 2]) else: # SPIN w.o. AIN sp_weight = sp_weight_fusion[:, :self.spt_length, :] lambda_color, offsets = None, None if self.stn: batch_C_prime = sp_weight_fusion[:, self.spt_length:, :].reshape([x.shape[0], self.F, 2]) build_P_prime = self.GridGenerator(batch_C_prime, self.I_r_size) build_P_prime_reshape = build_P_prime.reshape([build_P_prime.shape[0], self.I_r_size[0], self.I_r_size[1], 2]) x = self.sp_net(x, sp_weight, offsets, lambda_color) if self.stn: x = F.grid_sample(x=x, grid=build_P_prime_reshape, padding_mode='border') return x