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)