|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
This code is refer from: |
|
https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textdet/modules/local_graph.py |
|
""" |
|
|
|
from __future__ import absolute_import |
|
from __future__ import division |
|
from __future__ import print_function |
|
|
|
import numpy as np |
|
import paddle |
|
import paddle.nn as nn |
|
from ppocr.ext_op import RoIAlignRotated |
|
|
|
|
|
def normalize_adjacent_matrix(A): |
|
assert A.ndim == 2 |
|
assert A.shape[0] == A.shape[1] |
|
|
|
A = A + np.eye(A.shape[0]) |
|
d = np.sum(A, axis=0) |
|
d = np.clip(d, 0, None) |
|
d_inv = np.power(d, -0.5).flatten() |
|
d_inv[np.isinf(d_inv)] = 0.0 |
|
d_inv = np.diag(d_inv) |
|
G = A.dot(d_inv).transpose().dot(d_inv) |
|
return G |
|
|
|
|
|
def euclidean_distance_matrix(A, B): |
|
"""Calculate the Euclidean distance matrix. |
|
|
|
Args: |
|
A (ndarray): The point sequence. |
|
B (ndarray): The point sequence with the same dimensions as A. |
|
|
|
returns: |
|
D (ndarray): The Euclidean distance matrix. |
|
""" |
|
assert A.ndim == 2 |
|
assert B.ndim == 2 |
|
assert A.shape[1] == B.shape[1] |
|
|
|
m = A.shape[0] |
|
n = B.shape[0] |
|
|
|
A_dots = (A * A).sum(axis=1).reshape((m, 1)) * np.ones(shape=(1, n)) |
|
B_dots = (B * B).sum(axis=1) * np.ones(shape=(m, 1)) |
|
D_squared = A_dots + B_dots - 2 * A.dot(B.T) |
|
|
|
zero_mask = np.less(D_squared, 0.0) |
|
D_squared[zero_mask] = 0.0 |
|
D = np.sqrt(D_squared) |
|
return D |
|
|
|
|
|
def feature_embedding(input_feats, out_feat_len): |
|
"""Embed features. This code was partially adapted from |
|
https://github.com/GXYM/DRRG licensed under the MIT license. |
|
|
|
Args: |
|
input_feats (ndarray): The input features of shape (N, d), where N is |
|
the number of nodes in graph, d is the input feature vector length. |
|
out_feat_len (int): The length of output feature vector. |
|
|
|
Returns: |
|
embedded_feats (ndarray): The embedded features. |
|
""" |
|
assert input_feats.ndim == 2 |
|
assert isinstance(out_feat_len, int) |
|
assert out_feat_len >= input_feats.shape[1] |
|
|
|
num_nodes = input_feats.shape[0] |
|
feat_dim = input_feats.shape[1] |
|
feat_repeat_times = out_feat_len // feat_dim |
|
residue_dim = out_feat_len % feat_dim |
|
|
|
if residue_dim > 0: |
|
embed_wave = np.array([ |
|
np.power(1000, 2.0 * (j // 2) / feat_repeat_times + 1) |
|
for j in range(feat_repeat_times + 1) |
|
]).reshape((feat_repeat_times + 1, 1, 1)) |
|
repeat_feats = np.repeat( |
|
np.expand_dims( |
|
input_feats, axis=0), feat_repeat_times, axis=0) |
|
residue_feats = np.hstack([ |
|
input_feats[:, 0:residue_dim], np.zeros( |
|
(num_nodes, feat_dim - residue_dim)) |
|
]) |
|
residue_feats = np.expand_dims(residue_feats, axis=0) |
|
repeat_feats = np.concatenate([repeat_feats, residue_feats], axis=0) |
|
embedded_feats = repeat_feats / embed_wave |
|
embedded_feats[:, 0::2] = np.sin(embedded_feats[:, 0::2]) |
|
embedded_feats[:, 1::2] = np.cos(embedded_feats[:, 1::2]) |
|
embedded_feats = np.transpose(embedded_feats, (1, 0, 2)).reshape( |
|
(num_nodes, -1))[:, 0:out_feat_len] |
|
else: |
|
embed_wave = np.array([ |
|
np.power(1000, 2.0 * (j // 2) / feat_repeat_times) |
|
for j in range(feat_repeat_times) |
|
]).reshape((feat_repeat_times, 1, 1)) |
|
repeat_feats = np.repeat( |
|
np.expand_dims( |
|
input_feats, axis=0), feat_repeat_times, axis=0) |
|
embedded_feats = repeat_feats / embed_wave |
|
embedded_feats[:, 0::2] = np.sin(embedded_feats[:, 0::2]) |
|
embedded_feats[:, 1::2] = np.cos(embedded_feats[:, 1::2]) |
|
embedded_feats = np.transpose(embedded_feats, (1, 0, 2)).reshape( |
|
(num_nodes, -1)).astype(np.float32) |
|
|
|
return embedded_feats |
|
|
|
|
|
class LocalGraphs: |
|
def __init__(self, k_at_hops, num_adjacent_linkages, node_geo_feat_len, |
|
pooling_scale, pooling_output_size, local_graph_thr): |
|
|
|
assert len(k_at_hops) == 2 |
|
assert all(isinstance(n, int) for n in k_at_hops) |
|
assert isinstance(num_adjacent_linkages, int) |
|
assert isinstance(node_geo_feat_len, int) |
|
assert isinstance(pooling_scale, float) |
|
assert all(isinstance(n, int) for n in pooling_output_size) |
|
assert isinstance(local_graph_thr, float) |
|
|
|
self.k_at_hops = k_at_hops |
|
self.num_adjacent_linkages = num_adjacent_linkages |
|
self.node_geo_feat_dim = node_geo_feat_len |
|
self.pooling = RoIAlignRotated(pooling_output_size, pooling_scale) |
|
self.local_graph_thr = local_graph_thr |
|
|
|
def generate_local_graphs(self, sorted_dist_inds, gt_comp_labels): |
|
"""Generate local graphs for GCN to predict which instance a text |
|
component belongs to. |
|
|
|
Args: |
|
sorted_dist_inds (ndarray): The complete graph node indices, which |
|
is sorted according to the Euclidean distance. |
|
gt_comp_labels(ndarray): The ground truth labels define the |
|
instance to which the text components (nodes in graphs) belong. |
|
|
|
Returns: |
|
pivot_local_graphs(list[list[int]]): The list of local graph |
|
neighbor indices of pivots. |
|
pivot_knns(list[list[int]]): The list of k-nearest neighbor indices |
|
of pivots. |
|
""" |
|
|
|
assert sorted_dist_inds.ndim == 2 |
|
assert (sorted_dist_inds.shape[0] == sorted_dist_inds.shape[1] == |
|
gt_comp_labels.shape[0]) |
|
|
|
knn_graph = sorted_dist_inds[:, 1:self.k_at_hops[0] + 1] |
|
pivot_local_graphs = [] |
|
pivot_knns = [] |
|
for pivot_ind, knn in enumerate(knn_graph): |
|
|
|
local_graph_neighbors = set(knn) |
|
|
|
for neighbor_ind in knn: |
|
local_graph_neighbors.update( |
|
set(sorted_dist_inds[neighbor_ind, 1:self.k_at_hops[1] + |
|
1])) |
|
|
|
local_graph_neighbors.discard(pivot_ind) |
|
pivot_local_graph = list(local_graph_neighbors) |
|
pivot_local_graph.insert(0, pivot_ind) |
|
pivot_knn = [pivot_ind] + list(knn) |
|
|
|
if pivot_ind < 1: |
|
pivot_local_graphs.append(pivot_local_graph) |
|
pivot_knns.append(pivot_knn) |
|
else: |
|
add_flag = True |
|
for graph_ind, added_knn in enumerate(pivot_knns): |
|
added_pivot_ind = added_knn[0] |
|
added_local_graph = pivot_local_graphs[graph_ind] |
|
|
|
union = len( |
|
set(pivot_local_graph[1:]).union( |
|
set(added_local_graph[1:]))) |
|
intersect = len( |
|
set(pivot_local_graph[1:]).intersection( |
|
set(added_local_graph[1:]))) |
|
local_graph_iou = intersect / (union + 1e-8) |
|
|
|
if (local_graph_iou > self.local_graph_thr and |
|
pivot_ind in added_knn and |
|
gt_comp_labels[added_pivot_ind] == |
|
gt_comp_labels[pivot_ind] and |
|
gt_comp_labels[pivot_ind] != 0): |
|
add_flag = False |
|
break |
|
if add_flag: |
|
pivot_local_graphs.append(pivot_local_graph) |
|
pivot_knns.append(pivot_knn) |
|
|
|
return pivot_local_graphs, pivot_knns |
|
|
|
def generate_gcn_input(self, node_feat_batch, node_label_batch, |
|
local_graph_batch, knn_batch, sorted_dist_ind_batch): |
|
"""Generate graph convolution network input data. |
|
|
|
Args: |
|
node_feat_batch (List[Tensor]): The batched graph node features. |
|
node_label_batch (List[ndarray]): The batched text component |
|
labels. |
|
local_graph_batch (List[List[list[int]]]): The local graph node |
|
indices of image batch. |
|
knn_batch (List[List[list[int]]]): The knn graph node indices of |
|
image batch. |
|
sorted_dist_ind_batch (list[ndarray]): The node indices sorted |
|
according to the Euclidean distance. |
|
|
|
Returns: |
|
local_graphs_node_feat (Tensor): The node features of graph. |
|
adjacent_matrices (Tensor): The adjacent matrices of local graphs. |
|
pivots_knn_inds (Tensor): The k-nearest neighbor indices in |
|
local graph. |
|
gt_linkage (Tensor): The surpervision signal of GCN for linkage |
|
prediction. |
|
""" |
|
assert isinstance(node_feat_batch, list) |
|
assert isinstance(node_label_batch, list) |
|
assert isinstance(local_graph_batch, list) |
|
assert isinstance(knn_batch, list) |
|
assert isinstance(sorted_dist_ind_batch, list) |
|
|
|
num_max_nodes = max([ |
|
len(pivot_local_graph) |
|
for pivot_local_graphs in local_graph_batch |
|
for pivot_local_graph in pivot_local_graphs |
|
]) |
|
|
|
local_graphs_node_feat = [] |
|
adjacent_matrices = [] |
|
pivots_knn_inds = [] |
|
pivots_gt_linkage = [] |
|
|
|
for batch_ind, sorted_dist_inds in enumerate(sorted_dist_ind_batch): |
|
node_feats = node_feat_batch[batch_ind] |
|
pivot_local_graphs = local_graph_batch[batch_ind] |
|
pivot_knns = knn_batch[batch_ind] |
|
node_labels = node_label_batch[batch_ind] |
|
|
|
for graph_ind, pivot_knn in enumerate(pivot_knns): |
|
pivot_local_graph = pivot_local_graphs[graph_ind] |
|
num_nodes = len(pivot_local_graph) |
|
pivot_ind = pivot_local_graph[0] |
|
node2ind_map = {j: i for i, j in enumerate(pivot_local_graph)} |
|
|
|
knn_inds = paddle.to_tensor( |
|
[node2ind_map[i] for i in pivot_knn[1:]]) |
|
pivot_feats = node_feats[pivot_ind] |
|
normalized_feats = node_feats[paddle.to_tensor( |
|
pivot_local_graph)] - pivot_feats |
|
|
|
adjacent_matrix = np.zeros( |
|
(num_nodes, num_nodes), dtype=np.float32) |
|
for node in pivot_local_graph: |
|
neighbors = sorted_dist_inds[node, 1: |
|
self.num_adjacent_linkages + 1] |
|
for neighbor in neighbors: |
|
if neighbor in pivot_local_graph: |
|
|
|
adjacent_matrix[node2ind_map[node], node2ind_map[ |
|
neighbor]] = 1 |
|
adjacent_matrix[node2ind_map[neighbor], |
|
node2ind_map[node]] = 1 |
|
|
|
adjacent_matrix = normalize_adjacent_matrix(adjacent_matrix) |
|
pad_adjacent_matrix = paddle.zeros( |
|
(num_max_nodes, num_max_nodes)) |
|
pad_adjacent_matrix[:num_nodes, :num_nodes] = paddle.cast( |
|
paddle.to_tensor(adjacent_matrix), 'float32') |
|
|
|
pad_normalized_feats = paddle.concat( |
|
[ |
|
normalized_feats, paddle.zeros( |
|
(num_max_nodes - num_nodes, |
|
normalized_feats.shape[1])) |
|
], |
|
axis=0) |
|
local_graph_labels = node_labels[pivot_local_graph] |
|
knn_labels = local_graph_labels[knn_inds.numpy()] |
|
link_labels = ((node_labels[pivot_ind] == knn_labels) & |
|
(node_labels[pivot_ind] > 0)).astype(np.int64) |
|
link_labels = paddle.to_tensor(link_labels) |
|
|
|
local_graphs_node_feat.append(pad_normalized_feats) |
|
adjacent_matrices.append(pad_adjacent_matrix) |
|
pivots_knn_inds.append(knn_inds) |
|
pivots_gt_linkage.append(link_labels) |
|
|
|
local_graphs_node_feat = paddle.stack(local_graphs_node_feat, 0) |
|
adjacent_matrices = paddle.stack(adjacent_matrices, 0) |
|
pivots_knn_inds = paddle.stack(pivots_knn_inds, 0) |
|
pivots_gt_linkage = paddle.stack(pivots_gt_linkage, 0) |
|
|
|
return (local_graphs_node_feat, adjacent_matrices, pivots_knn_inds, |
|
pivots_gt_linkage) |
|
|
|
def __call__(self, feat_maps, comp_attribs): |
|
"""Generate local graphs as GCN input. |
|
|
|
Args: |
|
feat_maps (Tensor): The feature maps to extract the content |
|
features of text components. |
|
comp_attribs (ndarray): The text component attributes. |
|
|
|
Returns: |
|
local_graphs_node_feat (Tensor): The node features of graph. |
|
adjacent_matrices (Tensor): The adjacent matrices of local graphs. |
|
pivots_knn_inds (Tensor): The k-nearest neighbor indices in local |
|
graph. |
|
gt_linkage (Tensor): The surpervision signal of GCN for linkage |
|
prediction. |
|
""" |
|
|
|
assert isinstance(feat_maps, paddle.Tensor) |
|
assert comp_attribs.ndim == 3 |
|
assert comp_attribs.shape[2] == 8 |
|
|
|
sorted_dist_inds_batch = [] |
|
local_graph_batch = [] |
|
knn_batch = [] |
|
node_feat_batch = [] |
|
node_label_batch = [] |
|
|
|
for batch_ind in range(comp_attribs.shape[0]): |
|
num_comps = int(comp_attribs[batch_ind, 0, 0]) |
|
comp_geo_attribs = comp_attribs[batch_ind, :num_comps, 1:7] |
|
node_labels = comp_attribs[batch_ind, :num_comps, 7].astype( |
|
np.int32) |
|
|
|
comp_centers = comp_geo_attribs[:, 0:2] |
|
distance_matrix = euclidean_distance_matrix(comp_centers, |
|
comp_centers) |
|
|
|
batch_id = np.zeros( |
|
(comp_geo_attribs.shape[0], 1), dtype=np.float32) * batch_ind |
|
comp_geo_attribs[:, -2] = np.clip(comp_geo_attribs[:, -2], -1, 1) |
|
angle = np.arccos(comp_geo_attribs[:, -2]) * np.sign( |
|
comp_geo_attribs[:, -1]) |
|
angle = angle.reshape((-1, 1)) |
|
rotated_rois = np.hstack( |
|
[batch_id, comp_geo_attribs[:, :-2], angle]) |
|
rois = paddle.to_tensor(rotated_rois) |
|
content_feats = self.pooling(feat_maps[batch_ind].unsqueeze(0), |
|
rois) |
|
|
|
content_feats = content_feats.reshape([content_feats.shape[0], -1]) |
|
geo_feats = feature_embedding(comp_geo_attribs, |
|
self.node_geo_feat_dim) |
|
geo_feats = paddle.to_tensor(geo_feats) |
|
node_feats = paddle.concat([content_feats, geo_feats], axis=-1) |
|
|
|
sorted_dist_inds = np.argsort(distance_matrix, axis=1) |
|
pivot_local_graphs, pivot_knns = self.generate_local_graphs( |
|
sorted_dist_inds, node_labels) |
|
|
|
node_feat_batch.append(node_feats) |
|
node_label_batch.append(node_labels) |
|
local_graph_batch.append(pivot_local_graphs) |
|
knn_batch.append(pivot_knns) |
|
sorted_dist_inds_batch.append(sorted_dist_inds) |
|
|
|
(node_feats, adjacent_matrices, knn_inds, gt_linkage) = \ |
|
self.generate_gcn_input(node_feat_batch, |
|
node_label_batch, |
|
local_graph_batch, |
|
knn_batch, |
|
sorted_dist_inds_batch) |
|
|
|
return node_feats, adjacent_matrices, knn_inds, gt_linkage |
|
|