mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-02-19 07:44:47 +08:00
init
This commit is contained in:
28
backend/ppocr/modeling/transforms/__init__.py
Executable file
28
backend/ppocr/modeling/transforms/__init__.py
Executable file
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# 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.
|
||||
|
||||
__all__ = ['build_transform']
|
||||
|
||||
|
||||
def build_transform(config):
|
||||
from .tps import TPS
|
||||
from .stn import STN_ON
|
||||
|
||||
support_dict = ['TPS', 'STN_ON']
|
||||
|
||||
module_name = config.pop('name')
|
||||
assert module_name in support_dict, Exception(
|
||||
'transform only support {}'.format(support_dict))
|
||||
module_class = eval(module_name)(**config)
|
||||
return module_class
|
||||
135
backend/ppocr/modeling/transforms/stn.py
Normal file
135
backend/ppocr/modeling/transforms/stn.py
Normal file
@@ -0,0 +1,135 @@
|
||||
# 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.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/ayumiymk/aster.pytorch/blob/master/lib/models/stn_head.py
|
||||
"""
|
||||
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
|
||||
|
||||
from .tps_spatial_transformer import TPSSpatialTransformer
|
||||
|
||||
|
||||
def conv3x3_block(in_channels, out_channels, stride=1):
|
||||
n = 3 * 3 * out_channels
|
||||
w = math.sqrt(2. / n)
|
||||
conv_layer = nn.Conv2D(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
stride=stride,
|
||||
padding=1,
|
||||
weight_attr=nn.initializer.Normal(
|
||||
mean=0.0, std=w),
|
||||
bias_attr=nn.initializer.Constant(0))
|
||||
block = nn.Sequential(conv_layer, nn.BatchNorm2D(out_channels), nn.ReLU())
|
||||
return block
|
||||
|
||||
|
||||
class STN(nn.Layer):
|
||||
def __init__(self, in_channels, num_ctrlpoints, activation='none'):
|
||||
super(STN, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.num_ctrlpoints = num_ctrlpoints
|
||||
self.activation = activation
|
||||
self.stn_convnet = nn.Sequential(
|
||||
conv3x3_block(in_channels, 32), #32x64
|
||||
nn.MaxPool2D(
|
||||
kernel_size=2, stride=2),
|
||||
conv3x3_block(32, 64), #16x32
|
||||
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,
|
||||
weight_attr=nn.initializer.Normal(0, 0.001),
|
||||
bias_attr=nn.initializer.Constant(0)),
|
||||
nn.BatchNorm1D(512),
|
||||
nn.ReLU())
|
||||
fc2_bias = self.init_stn()
|
||||
self.stn_fc2 = nn.Linear(
|
||||
512,
|
||||
num_ctrlpoints * 2,
|
||||
weight_attr=nn.initializer.Constant(0.0),
|
||||
bias_attr=nn.initializer.Assign(fc2_bias))
|
||||
|
||||
def init_stn(self):
|
||||
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.)
|
||||
ctrl_points = paddle.to_tensor(ctrl_points)
|
||||
fc2_bias = paddle.reshape(
|
||||
ctrl_points, shape=[ctrl_points.shape[0] * ctrl_points.shape[1]])
|
||||
return fc2_bias
|
||||
|
||||
def forward(self, x):
|
||||
x = self.stn_convnet(x)
|
||||
batch_size, _, h, w = x.shape
|
||||
x = paddle.reshape(x, shape=(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 = paddle.reshape(x, shape=[-1, self.num_ctrlpoints, 2])
|
||||
return img_feat, x
|
||||
|
||||
|
||||
class STN_ON(nn.Layer):
|
||||
def __init__(self, in_channels, tps_inputsize, tps_outputsize,
|
||||
num_control_points, tps_margins, stn_activation):
|
||||
super(STN_ON, self).__init__()
|
||||
self.tps = TPSSpatialTransformer(
|
||||
output_image_size=tuple(tps_outputsize),
|
||||
num_control_points=num_control_points,
|
||||
margins=tuple(tps_margins))
|
||||
self.stn_head = STN(in_channels=in_channels,
|
||||
num_ctrlpoints=num_control_points,
|
||||
activation=stn_activation)
|
||||
self.tps_inputsize = tps_inputsize
|
||||
self.out_channels = in_channels
|
||||
|
||||
def forward(self, image):
|
||||
stn_input = paddle.nn.functional.interpolate(
|
||||
image, self.tps_inputsize, mode="bilinear", align_corners=True)
|
||||
stn_img_feat, ctrl_points = self.stn_head(stn_input)
|
||||
x, _ = self.tps(image, ctrl_points)
|
||||
return x
|
||||
308
backend/ppocr/modeling/transforms/tps.py
Normal file
308
backend/ppocr/modeling/transforms/tps.py
Normal file
@@ -0,0 +1,308 @@
|
||||
# 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.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/clovaai/deep-text-recognition-benchmark/blob/master/modules/transformation.py
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
|
||||
class ConvBNLayer(nn.Layer):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
groups=1,
|
||||
act=None,
|
||||
name=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,
|
||||
weight_attr=ParamAttr(name=name + "_weights"),
|
||||
bias_attr=False)
|
||||
bn_name = "bn_" + name
|
||||
self.bn = nn.BatchNorm(
|
||||
out_channels,
|
||||
act=act,
|
||||
param_attr=ParamAttr(name=bn_name + '_scale'),
|
||||
bias_attr=ParamAttr(bn_name + '_offset'),
|
||||
moving_mean_name=bn_name + '_mean',
|
||||
moving_variance_name=bn_name + '_variance')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
return x
|
||||
|
||||
|
||||
class LocalizationNetwork(nn.Layer):
|
||||
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 = []
|
||||
for fno in range(0, len(num_filters_list)):
|
||||
num_filters = num_filters_list[fno]
|
||||
name = "loc_conv%d" % fno
|
||||
conv = self.add_sublayer(
|
||||
name,
|
||||
ConvBNLayer(
|
||||
in_channels=in_channels,
|
||||
out_channels=num_filters,
|
||||
kernel_size=3,
|
||||
act='relu',
|
||||
name=name))
|
||||
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)
|
||||
name = "loc_fc1"
|
||||
stdv = 1.0 / math.sqrt(num_filters_list[-1] * 1.0)
|
||||
self.fc1 = nn.Linear(
|
||||
in_channels,
|
||||
fc_dim,
|
||||
weight_attr=ParamAttr(
|
||||
learning_rate=loc_lr,
|
||||
name=name + "_w",
|
||||
initializer=nn.initializer.Uniform(-stdv, stdv)),
|
||||
bias_attr=ParamAttr(name=name + '.b_0'),
|
||||
name=name)
|
||||
|
||||
# Init fc2 in LocalizationNetwork
|
||||
initial_bias = self.get_initial_fiducials()
|
||||
initial_bias = initial_bias.reshape(-1)
|
||||
name = "loc_fc2"
|
||||
param_attr = ParamAttr(
|
||||
learning_rate=loc_lr,
|
||||
initializer=nn.initializer.Assign(np.zeros([fc_dim, F * 2])),
|
||||
name=name + "_w")
|
||||
bias_attr = ParamAttr(
|
||||
learning_rate=loc_lr,
|
||||
initializer=nn.initializer.Assign(initial_bias),
|
||||
name=name + "_b")
|
||||
self.fc2 = nn.Linear(
|
||||
fc_dim,
|
||||
F * 2,
|
||||
weight_attr=param_attr,
|
||||
bias_attr=bias_attr,
|
||||
name=name)
|
||||
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
|
||||
"""
|
||||
B = x.shape[0]
|
||||
i = 0
|
||||
for block in self.block_list:
|
||||
x = block(x)
|
||||
x = x.squeeze(axis=2).squeeze(axis=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.Layer):
|
||||
def __init__(self, in_channels, num_fiducial):
|
||||
super(GridGenerator, self).__init__()
|
||||
self.eps = 1e-6
|
||||
self.F = num_fiducial
|
||||
|
||||
name = "ex_fc"
|
||||
initializer = nn.initializer.Constant(value=0.0)
|
||||
param_attr = ParamAttr(
|
||||
learning_rate=0.0, initializer=initializer, name=name + "_w")
|
||||
bias_attr = ParamAttr(
|
||||
learning_rate=0.0, initializer=initializer, name=name + "_b")
|
||||
self.fc = nn.Linear(
|
||||
in_channels,
|
||||
6,
|
||||
weight_attr=param_attr,
|
||||
bias_attr=bias_attr,
|
||||
name=name)
|
||||
|
||||
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).astype('float32')
|
||||
P_hat_tensor = self.build_P_hat_paddle(
|
||||
C, paddle.to_tensor(P)).astype('float32')
|
||||
|
||||
inv_delta_C_tensor.stop_gradient = True
|
||||
P_hat_tensor.stop_gradient = True
|
||||
|
||||
batch_C_ex_part_tensor = self.get_expand_tensor(batch_C_prime)
|
||||
|
||||
batch_C_ex_part_tensor.stop_gradient = True
|
||||
|
||||
batch_C_prime_with_zeros = paddle.concat(
|
||||
[batch_C_prime, batch_C_ex_part_tensor], axis=1)
|
||||
batch_T = paddle.matmul(inv_delta_C_tensor, batch_C_prime_with_zeros)
|
||||
batch_P_prime = paddle.matmul(P_hat_tensor, 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 = paddle.linspace(-1.0, 1.0, int(F / 2), dtype='float64')
|
||||
ctrl_pts_y_top = -1 * paddle.ones([int(F / 2)], dtype='float64')
|
||||
ctrl_pts_y_bottom = paddle.ones([int(F / 2)], dtype='float64')
|
||||
ctrl_pts_top = paddle.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
|
||||
ctrl_pts_bottom = paddle.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
|
||||
C = paddle.concat([ctrl_pts_top, ctrl_pts_bottom], axis=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 = (paddle.arange(
|
||||
-I_r_width, I_r_width, 2, dtype='float64') + 1.0
|
||||
) / paddle.to_tensor(np.array([I_r_width]))
|
||||
|
||||
I_r_grid_y = (paddle.arange(
|
||||
-I_r_height, I_r_height, 2, dtype='float64') + 1.0
|
||||
) / paddle.to_tensor(np.array([I_r_height]))
|
||||
|
||||
# P: self.I_r_width x self.I_r_height x 2
|
||||
P = paddle.stack(paddle.meshgrid(I_r_grid_x, I_r_grid_y), axis=2)
|
||||
P = paddle.transpose(P, perm=[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 = paddle.eye(F, dtype='float64') # F x F
|
||||
hat_C = paddle.norm(
|
||||
C.reshape([1, F, 2]) - C.reshape([F, 1, 2]), axis=2) + hat_eye
|
||||
hat_C = (hat_C**2) * paddle.log(hat_C)
|
||||
delta_C = paddle.concat( # F+3 x F+3
|
||||
[
|
||||
paddle.concat(
|
||||
[paddle.ones(
|
||||
(F, 1), dtype='float64'), C, hat_C], axis=1), # F x F+3
|
||||
paddle.concat(
|
||||
[
|
||||
paddle.zeros(
|
||||
(2, 3), dtype='float64'), paddle.transpose(
|
||||
C, perm=[1, 0])
|
||||
],
|
||||
axis=1), # 2 x F+3
|
||||
paddle.concat(
|
||||
[
|
||||
paddle.zeros(
|
||||
(1, 3), dtype='float64'), paddle.ones(
|
||||
(1, F), dtype='float64')
|
||||
],
|
||||
axis=1) # 1 x F+3
|
||||
],
|
||||
axis=0)
|
||||
inv_delta_C = paddle.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 = paddle.tile(paddle.unsqueeze(P, axis=1), (1, F, 1))
|
||||
C_tile = paddle.unsqueeze(C, axis=0) # 1 x F x 2
|
||||
P_diff = P_tile - C_tile # n x F x 2
|
||||
# rbf_norm: n x F
|
||||
rbf_norm = paddle.norm(P_diff, p=2, axis=2, keepdim=False)
|
||||
|
||||
# rbf: n x F
|
||||
rbf = paddle.multiply(
|
||||
paddle.square(rbf_norm), paddle.log(rbf_norm + eps))
|
||||
P_hat = paddle.concat(
|
||||
[paddle.ones(
|
||||
(n, 1), dtype='float64'), P, rbf], axis=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.Layer):
|
||||
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])
|
||||
batch_I_r = F.grid_sample(x=image, grid=batch_P_prime)
|
||||
return batch_I_r
|
||||
156
backend/ppocr/modeling/transforms/tps_spatial_transformer.py
Normal file
156
backend/ppocr/modeling/transforms/tps_spatial_transformer.py
Normal file
@@ -0,0 +1,156 @@
|
||||
# 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.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/ayumiymk/aster.pytorch/blob/master/lib/models/tps_spatial_transformer.py
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
def grid_sample(input, grid, canvas=None):
|
||||
input.stop_gradient = False
|
||||
output = F.grid_sample(input, grid)
|
||||
if canvas is None:
|
||||
return output
|
||||
else:
|
||||
input_mask = paddle.ones(shape=input.shape)
|
||||
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.shape[0]
|
||||
M = control_points.shape[0]
|
||||
pairwise_diff = paddle.reshape(
|
||||
input_points, shape=[N, 1, 2]) - paddle.reshape(
|
||||
control_points, shape=[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 * paddle.log(pairwise_dist)
|
||||
# fix numerical error for 0 * log(0), substitute all nan with 0
|
||||
mask = np.array(repr_matrix != repr_matrix)
|
||||
repr_matrix[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)
|
||||
output_ctrl_pts_arr = np.concatenate(
|
||||
[ctrl_pts_top, ctrl_pts_bottom], axis=0)
|
||||
output_ctrl_pts = paddle.to_tensor(output_ctrl_pts_arr)
|
||||
return output_ctrl_pts
|
||||
|
||||
|
||||
class TPSSpatialTransformer(nn.Layer):
|
||||
def __init__(self,
|
||||
output_image_size=None,
|
||||
num_control_points=None,
|
||||
margins=None):
|
||||
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
|
||||
|
||||
# create padded kernel matrix
|
||||
forward_kernel = paddle.zeros(shape=[N + 3, N + 3])
|
||||
target_control_partial_repr = compute_partial_repr(
|
||||
target_control_points, target_control_points)
|
||||
target_control_partial_repr = paddle.cast(target_control_partial_repr,
|
||||
forward_kernel.dtype)
|
||||
forward_kernel[:N, :N] = target_control_partial_repr
|
||||
forward_kernel[:N, -3] = 1
|
||||
forward_kernel[-3, :N] = 1
|
||||
target_control_points = paddle.cast(target_control_points,
|
||||
forward_kernel.dtype)
|
||||
forward_kernel[:N, -2:] = target_control_points
|
||||
forward_kernel[-2:, :N] = paddle.transpose(
|
||||
target_control_points, perm=[1, 0])
|
||||
# compute inverse matrix
|
||||
inverse_kernel = paddle.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 = paddle.to_tensor(target_coordinate) # HW x 2
|
||||
Y, X = paddle.split(
|
||||
target_coordinate, target_coordinate.shape[1], axis=1)
|
||||
Y = Y / (self.target_height - 1)
|
||||
X = X / (self.target_width - 1)
|
||||
target_coordinate = paddle.concat(
|
||||
[X, Y], axis=1) # convert from (y, x) to (x, y)
|
||||
target_coordinate_partial_repr = compute_partial_repr(
|
||||
target_coordinate, target_control_points)
|
||||
target_coordinate_repr = paddle.concat(
|
||||
[
|
||||
target_coordinate_partial_repr, paddle.ones(shape=[HW, 1]),
|
||||
target_coordinate
|
||||
],
|
||||
axis=1)
|
||||
|
||||
# register precomputed matrices
|
||||
self.inverse_kernel = inverse_kernel
|
||||
self.padding_matrix = paddle.zeros(shape=[3, 2])
|
||||
self.target_coordinate_repr = target_coordinate_repr
|
||||
self.target_control_points = target_control_points
|
||||
|
||||
def forward(self, input, source_control_points):
|
||||
assert source_control_points.ndimension() == 3
|
||||
assert source_control_points.shape[1] == self.num_control_points
|
||||
assert source_control_points.shape[2] == 2
|
||||
batch_size = paddle.shape(source_control_points)[0]
|
||||
|
||||
padding_matrix = paddle.expand(
|
||||
self.padding_matrix, shape=[batch_size, 3, 2])
|
||||
Y = paddle.concat([source_control_points, padding_matrix], 1)
|
||||
mapping_matrix = paddle.matmul(self.inverse_kernel, Y)
|
||||
source_coordinate = paddle.matmul(self.target_coordinate_repr,
|
||||
mapping_matrix)
|
||||
|
||||
grid = paddle.reshape(
|
||||
source_coordinate,
|
||||
shape=[-1, self.target_height, self.target_width, 2])
|
||||
grid = paddle.clip(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, source_coordinate
|
||||
Reference in New Issue
Block a user