# copyright (c) 2023 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
try:
from collections import Callable
except:
from collections.abc import Callable
import numpy as np
import paddle
from paddle import nn
from paddle.nn import functional as F
from ppocr.modeling.heads.rec_nrtr_head import Embeddings
from ppocr.modeling.backbones.rec_svtrnet import (
DropPath,
Identity,
trunc_normal_,
zeros_,
ones_,
Mlp,
)
class Attention(nn.Layer):
def __init__(
self,
dim,
num_heads=8,
qkv_bias=False,
qk_scale=None,
attn_drop=0.0,
proj_drop=0.0,
):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim**-0.5
self.q = nn.Linear(dim, dim, bias_attr=qkv_bias)
self.kv = nn.Linear(dim, dim * 2, bias_attr=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, q, kv):
N, C = kv.shape[1:]
QN = q.shape[1]
q = (
self.q(q)
.reshape([-1, QN, self.num_heads, C // self.num_heads])
.transpose([0, 2, 1, 3])
)
k, v = (
self.kv(kv)
.reshape([-1, N, 2, self.num_heads, C // self.num_heads])
.transpose((2, 0, 3, 1, 4))
)
attn = q.matmul(k.transpose((0, 1, 3, 2))) * self.scale
attn = F.softmax(attn, axis=-1)
attn = self.attn_drop(attn)
x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((-1, QN, C))
x = self.proj(x)
x = self.proj_drop(x)
return x
class EdgeDecoderLayer(nn.Layer):
def __init__(
self,
dim,
num_heads,
mlp_ratio=4.0,
qkv_bias=False,
qk_scale=None,
drop=0.0,
attn_drop=0.0,
drop_path=[0.0, 0.0],
act_layer=nn.GELU,
norm_layer="nn.LayerNorm",
epsilon=1e-6,
):
super().__init__()
self.head_dim = dim // num_heads
self.scale = qk_scale or self.head_dim**-0.5
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path1 = DropPath(drop_path[0]) if drop_path[0] > 0.0 else Identity()
self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
self.p = nn.Linear(dim, dim)
self.cv = nn.Linear(dim, dim)
self.pv = nn.Linear(dim, dim)
self.dim = dim
self.num_heads = num_heads
self.p_proj = nn.Linear(dim, dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp_ratio = mlp_ratio
self.mlp = Mlp(
in_features=dim,
hidden_features=mlp_hidden_dim,
act_layer=act_layer,
drop=drop,
)
def forward(self, p, cv, pv):
pN = p.shape[1]
vN = cv.shape[1]
p_shortcut = p
p1 = (
self.p(p)
.reshape([-1, pN, self.num_heads, self.dim // self.num_heads])
.transpose([0, 2, 1, 3])
)
cv1 = (
self.cv(cv)
.reshape([-1, vN, self.num_heads, self.dim // self.num_heads])
.transpose([0, 2, 1, 3])
)
pv1 = (
self.pv(pv)
.reshape([-1, vN, self.num_heads, self.dim // self.num_heads])
.transpose([0, 2, 1, 3])
)
edge = F.softmax(p1.matmul(pv1.transpose((0, 1, 3, 2))), -1) # B h N N
p_c = (edge @ cv1).transpose((0, 2, 1, 3)).reshape((-1, pN, self.dim))
x1 = self.norm1(p_shortcut + self.drop_path1(self.p_proj(p_c)))
x = self.norm2(x1 + self.drop_path1(self.mlp(x1)))
return x
class DecoderLayer(nn.Layer):
def __init__(
self,
dim,
num_heads,
mlp_ratio=4.0,
qkv_bias=False,
qk_scale=None,
drop=0.0,
attn_drop=0.0,
drop_path=0.0,
act_layer=nn.GELU,
norm_layer="nn.LayerNorm",
epsilon=1e-6,
):
super().__init__()
if isinstance(norm_layer, str):
self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
self.normkv = eval(norm_layer)(dim, epsilon=epsilon)
elif isinstance(norm_layer, Callable):
self.norm1 = norm_layer(dim)
self.normkv = norm_layer(dim)
else:
raise TypeError("The norm_layer must be str or paddle.nn.LayerNorm class")
self.mixer = Attention(
dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
attn_drop=attn_drop,
proj_drop=drop,
)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
if isinstance(norm_layer, str):
self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
elif isinstance(norm_layer, Callable):
self.norm2 = norm_layer(dim)
else:
raise TypeError("The norm_layer must be str or paddle.nn.layer.Layer class")
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp_ratio = mlp_ratio
self.mlp = Mlp(
in_features=dim,
hidden_features=mlp_hidden_dim,
act_layer=act_layer,
drop=drop,
)
def forward(self, q, kv):
x1 = self.norm1(q + self.drop_path(self.mixer(q, kv)))
x = self.norm2(x1 + self.drop_path(self.mlp(x1)))
return x
class CPPDHead(nn.Layer):
def __init__(
self,
in_channels,
dim,
out_channels,
num_layer=2,
drop_path_rate=0.1,
max_len=25,
vis_seq=50,
ch=False,
**kwargs
):
super(CPPDHead, self).__init__()
self.out_channels = out_channels # none + 26 + 10
self.dim = dim
self.ch = ch
self.max_len = max_len + 1 # max_len + eos
self.char_node_embed = Embeddings(
d_model=dim, vocab=self.out_channels, scale_embedding=True
)
self.pos_node_embed = Embeddings(
d_model=dim, vocab=self.max_len, scale_embedding=True
)
dpr = np.linspace(0, drop_path_rate, num_layer + 1)
self.char_node_decoder = nn.LayerList(
[
DecoderLayer(
dim=dim,
num_heads=dim // 32,
mlp_ratio=4.0,
qkv_bias=True,
drop_path=dpr[i],
)
for i in range(num_layer)
]
)
self.pos_node_decoder = nn.LayerList(
[
DecoderLayer(
dim=dim,
num_heads=dim // 32,
mlp_ratio=4.0,
qkv_bias=True,
drop_path=dpr[i],
)
for i in range(num_layer)
]
)
self.edge_decoder = EdgeDecoderLayer(
dim=dim,
num_heads=dim // 32,
mlp_ratio=4.0,
qkv_bias=True,
drop_path=dpr[num_layer : num_layer + 1],
)
self.char_pos_embed = self.create_parameter(
shape=[1, self.max_len, dim], default_initializer=zeros_
)
self.add_parameter("char_pos_embed", self.char_pos_embed)
self.vis_pos_embed = self.create_parameter(
shape=[1, vis_seq, dim], default_initializer=zeros_
)
self.add_parameter("vis_pos_embed", self.vis_pos_embed)
self.char_node_fc1 = nn.Linear(dim, max_len)
self.pos_node_fc1 = nn.Linear(dim, self.max_len)
self.edge_fc = nn.Linear(dim, self.out_channels)
trunc_normal_(self.char_pos_embed)
trunc_normal_(self.vis_pos_embed)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight)
if isinstance(m, nn.Linear) and m.bias is not None:
zeros_(m.bias)
elif isinstance(m, nn.LayerNorm):
zeros_(m.bias)
ones_(m.weight)
def forward(self, x, targets=None, epoch=0):
if self.training:
return self.forward_train(x, targets, epoch)
else:
return self.forward_test(x)
def forward_test(self, x):
visual_feats = x + self.vis_pos_embed
bs = visual_feats.shape[0]
pos_node_embed = (
self.pos_node_embed(paddle.arange(self.max_len)).unsqueeze(0)
+ self.char_pos_embed
)
pos_node_embed = paddle.tile(pos_node_embed, [bs, 1, 1])
char_vis_node_query = visual_feats
pos_vis_node_query = paddle.concat([pos_node_embed, visual_feats], 1)
for char_decoder_layer, pos_decoder_layer in zip(
self.char_node_decoder, self.pos_node_decoder
):
char_vis_node_query = char_decoder_layer(
char_vis_node_query, char_vis_node_query
)
pos_vis_node_query = pos_decoder_layer(
pos_vis_node_query, pos_vis_node_query[:, self.max_len :, :]
)
pos_node_query = pos_vis_node_query[:, : self.max_len, :]
char_vis_feats = char_vis_node_query
pos_node_feats = self.edge_decoder(
pos_node_query, char_vis_feats, char_vis_feats
) # B, 26, dim
edge_feats = self.edge_fc(pos_node_feats) # B, 26, 37
edge_logits = F.softmax(edge_feats, -1)
return edge_logits
def forward_train(self, x, targets=None, epoch=0):
visual_feats = x + self.vis_pos_embed
bs = visual_feats.shape[0]
if self.ch:
char_node_embed = self.char_node_embed(targets[-2])
else:
char_node_embed = self.char_node_embed(
paddle.arange(self.out_channels)
).unsqueeze(0)
char_node_embed = paddle.tile(char_node_embed, [bs, 1, 1])
counting_char_num = char_node_embed.shape[1]
pos_node_embed = (
self.pos_node_embed(paddle.arange(self.max_len)).unsqueeze(0)
+ self.char_pos_embed
)
pos_node_embed = paddle.tile(pos_node_embed, [bs, 1, 1])
node_feats = []
char_vis_node_query = paddle.concat([char_node_embed, visual_feats], 1)
pos_vis_node_query = paddle.concat([pos_node_embed, visual_feats], 1)
for char_decoder_layer, pos_decoder_layer in zip(
self.char_node_decoder, self.pos_node_decoder
):
char_vis_node_query = char_decoder_layer(
char_vis_node_query, char_vis_node_query[:, counting_char_num:, :]
)
pos_vis_node_query = pos_decoder_layer(
pos_vis_node_query, pos_vis_node_query[:, self.max_len :, :]
)
char_node_query = char_vis_node_query[:, :counting_char_num, :]
pos_node_query = pos_vis_node_query[:, : self.max_len, :]
char_vis_feats = char_vis_node_query[:, counting_char_num:, :]
char_node_feats1 = self.char_node_fc1(char_node_query)
pos_node_feats1 = self.pos_node_fc1(pos_node_query)
diag_mask = (
paddle.eye(pos_node_feats1.shape[1])
.unsqueeze(0)
.tile([pos_node_feats1.shape[0], 1, 1])
)
pos_node_feats1 = (pos_node_feats1 * diag_mask).sum(-1)
node_feats.append(char_node_feats1)
node_feats.append(pos_node_feats1)
pos_node_feats = self.edge_decoder(
pos_node_query, char_vis_feats, char_vis_feats
) # B, 26, dim
edge_feats = self.edge_fc(pos_node_feats) # B, 26, 37
return node_feats, edge_feats