Spaces:
Runtime error
Runtime error
File size: 3,291 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 |
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from mmocr.models.textdet.dense_heads import DRRGHead
def test_drrg_head():
in_channels = 10
drrg_head = DRRGHead(in_channels)
assert drrg_head.in_channels == in_channels
assert drrg_head.k_at_hops == (8, 4)
assert drrg_head.num_adjacent_linkages == 3
assert drrg_head.node_geo_feat_len == 120
assert np.allclose(drrg_head.pooling_scale, 1.0)
assert drrg_head.pooling_output_size == (4, 3)
assert np.allclose(drrg_head.nms_thr, 0.3)
assert np.allclose(drrg_head.min_width, 8.0)
assert np.allclose(drrg_head.max_width, 24.0)
assert np.allclose(drrg_head.comp_shrink_ratio, 1.03)
assert np.allclose(drrg_head.comp_ratio, 0.4)
assert np.allclose(drrg_head.comp_score_thr, 0.3)
assert np.allclose(drrg_head.text_region_thr, 0.2)
assert np.allclose(drrg_head.center_region_thr, 0.2)
assert drrg_head.center_region_area_thr == 50
assert np.allclose(drrg_head.local_graph_thr, 0.7)
# test forward train
num_rois = 16
feature_maps = torch.randn((2, 10, 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_])
pred_maps, gcn_data = drrg_head(feature_maps, comp_attribs)
pred_labels, gt_labels = gcn_data
assert pred_maps.size() == (2, 6, 128, 128)
assert pred_labels.ndim == gt_labels.ndim == 2
assert gt_labels.size()[0] * gt_labels.size()[1] == pred_labels.size()[0]
assert pred_labels.size()[1] == 2
# test forward test
with torch.no_grad():
feat_maps = torch.zeros((1, 10, 128, 128))
drrg_head.out_conv.bias.data.fill_(-10)
preds = drrg_head.single_test(feat_maps)
assert all([pred is None for pred in preds])
# test get_boundary
edges = np.stack([np.arange(0, 10), np.arange(1, 11)]).transpose()
edges = np.vstack([edges, np.array([1, 0])])
scores = np.ones(11, dtype=np.float32) * 0.9
x1 = np.arange(2, 22, 2)
x2 = x1 + 2
y1 = np.ones(10) * 2
y2 = y1 + 2
comp_scores = np.ones(10, dtype=np.float32) * 0.9
text_comps = np.stack([x1, y1, x2, y1, x2, y2, x1, y2,
comp_scores]).transpose()
outlier = np.array([50, 50, 52, 50, 52, 52, 50, 52, 0.9])
text_comps = np.vstack([text_comps, outlier])
(C, H, W) = (10, 128, 128)
img_metas = [{
'img_shape': (H, W, C),
'ori_shape': (H, W, C),
'pad_shape': (H, W, C),
'filename': '<demo>.png',
'scale_factor': np.array([1, 1, 1, 1]),
'flip': False,
}]
results = drrg_head.get_boundary(
edges, scores, text_comps, img_metas, rescale=True)
assert 'boundary_result' in results.keys()
|