Spaces:
Runtime error
Runtime error
File size: 5,158 Bytes
2366e36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from mmocr.models.textdet.modules import GCN, LocalGraphs, ProposalLocalGraphs
from mmocr.models.textdet.modules.utils import (feature_embedding,
normalize_adjacent_matrix)
def test_local_graph_forward_train():
geo_feat_len = 24
pooling_h, pooling_w = pooling_out_size = (2, 2)
num_rois = 32
local_graph_generator = LocalGraphs((4, 4), 3, geo_feat_len, 1.0,
pooling_out_size, 0.5)
feature_maps = torch.randn((2, 3, 128, 128), dtype=torch.float)
x = np.random.randint(4, 124, (num_rois, 1))
y = np.random.randint(4, 124, (num_rois, 1))
h = 4 * np.ones((num_rois, 1))
w = 4 * np.ones((num_rois, 1))
angle = (np.random.random_sample((num_rois, 1)) * 2 - 1) * np.pi / 2
cos, sin = np.cos(angle), np.sin(angle)
comp_labels = np.random.randint(1, 3, (num_rois, 1))
num_rois = num_rois * np.ones((num_rois, 1))
comp_attribs = np.hstack([num_rois, x, y, h, w, cos, sin, comp_labels])
comp_attribs = comp_attribs.astype(np.float32)
comp_attribs_ = comp_attribs.copy()
comp_attribs = np.stack([comp_attribs, comp_attribs_])
(node_feats, adjacent_matrix, knn_inds,
linkage_labels) = local_graph_generator(feature_maps, comp_attribs)
feat_len = geo_feat_len + feature_maps.size()[1] * pooling_h * pooling_w
assert node_feats.dim() == adjacent_matrix.dim() == 3
assert node_feats.size()[-1] == feat_len
assert knn_inds.size()[-1] == 4
assert linkage_labels.size()[-1] == 4
assert (node_feats.size()[0] == adjacent_matrix.size()[0] ==
knn_inds.size()[0] == linkage_labels.size()[0])
assert (node_feats.size()[1] == adjacent_matrix.size()[1] ==
adjacent_matrix.size()[2])
def test_local_graph_forward_test():
geo_feat_len = 24
pooling_h, pooling_w = pooling_out_size = (2, 2)
local_graph_generator = ProposalLocalGraphs(
(4, 4), 2, geo_feat_len, 1., pooling_out_size, 0.1, 3., 6., 1., 0.5,
0.3, 0.5, 0.5, 2)
maps = torch.zeros((1, 6, 224, 224), dtype=torch.float)
maps[:, 0:2, :, :] = -10.
maps[:, 0, 60:100, 50:170] = 10.
maps[:, 1, 75:85, 60:160] = 10.
maps[:, 2, 75:85, 60:160] = 0.
maps[:, 3, 75:85, 60:160] = 1.
maps[:, 4, 75:85, 60:160] = 10.
maps[:, 5, 75:85, 60:160] = 10.
feature_maps = torch.randn((2, 6, 224, 224), dtype=torch.float)
feat_len = geo_feat_len + feature_maps.size()[1] * pooling_h * pooling_w
none_flag, graph_data = local_graph_generator(maps, feature_maps)
(node_feats, adjacent_matrices, knn_inds, local_graphs,
text_comps) = graph_data
assert none_flag is False
assert text_comps.ndim == 2
assert text_comps.shape[0] > 0
assert text_comps.shape[1] == 9
assert (node_feats.size()[0] == adjacent_matrices.size()[0] ==
knn_inds.size()[0] == local_graphs.size()[0] ==
text_comps.shape[0])
assert (node_feats.size()[1] == adjacent_matrices.size()[1] ==
adjacent_matrices.size()[2] == local_graphs.size()[1])
assert node_feats.size()[-1] == feat_len
# test proposal local graphs with area of center region less than threshold
maps[:, 1, 75:85, 60:160] = -10.
maps[:, 1, 80, 80] = 10.
none_flag, _ = local_graph_generator(maps, feature_maps)
assert none_flag
# test proposal local graphs with one text component
local_graph_generator = ProposalLocalGraphs(
(4, 4), 2, geo_feat_len, 1., pooling_out_size, 0.1, 8., 20., 1., 0.5,
0.3, 0.5, 0.5, 2)
maps[:, 1, 78:82, 78:82] = 10.
none_flag, _ = local_graph_generator(maps, feature_maps)
assert none_flag
# test proposal local graphs with text components out of text region
maps[:, 0, 60:100, 50:170] = -10.
maps[:, 0, 78:82, 78:82] = 10.
none_flag, _ = local_graph_generator(maps, feature_maps)
assert none_flag
def test_gcn():
num_local_graphs = 32
num_max_graph_nodes = 16
input_feat_len = 512
k = 8
gcn = GCN(input_feat_len)
node_feat = torch.randn(
(num_local_graphs, num_max_graph_nodes, input_feat_len))
adjacent_matrix = torch.rand(
(num_local_graphs, num_max_graph_nodes, num_max_graph_nodes))
knn_inds = torch.randint(1, num_max_graph_nodes, (num_local_graphs, k))
output = gcn(node_feat, adjacent_matrix, knn_inds)
assert output.size() == (num_local_graphs * k, 2)
def test_normalize_adjacent_matrix():
adjacent_matrix = np.random.randint(0, 2, (16, 16))
normalized_matrix = normalize_adjacent_matrix(adjacent_matrix)
assert normalized_matrix.shape == adjacent_matrix.shape
def test_feature_embedding():
out_feat_len = 48
# test without residue dimensions
feats = np.random.randn(10, 8)
embed_feats = feature_embedding(feats, out_feat_len)
assert embed_feats.shape == (10, out_feat_len)
# test with residue dimensions
feats = np.random.randn(10, 9)
embed_feats = feature_embedding(feats, out_feat_len)
assert embed_feats.shape == (10, out_feat_len)
|