def test_polygon_mask_to_bitmap():
# polygon masks contain 3 instances to bitmap
raw_masks = dummy_raw_polygon_masks((3, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
bitmap_masks = polygon_masks.to_bitmap()
assert (polygon_masks.to_ndarray() == bitmap_masks.to_ndarray()).all()
def test_polygon_mask_iter():
raw_masks = dummy_raw_polygon_masks((3, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
for i, polygon_mask in enumerate(polygon_masks):
assert np.equal(polygon_mask, raw_masks[i]).all()
def test_polygon_mask_expand():
with pytest.raises(NotImplementedError):
raw_masks = dummy_raw_polygon_masks((0, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
polygon_masks.expand(56, 56, 10, 17)
def test_polygon_mask_flip():
# flip with empty polygon masks
raw_masks = dummy_raw_polygon_masks((0, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
flipped_masks = polygon_masks.flip(flip_direction='horizontal')
assert len(flipped_masks) == 0
assert flipped_masks.height == 28
assert flipped_masks.width == 28
assert flipped_masks.to_ndarray().shape == (0, 28, 28)
# TODO: fixed flip correctness checking after v2.0_coord is merged
# horizontally flip with polygon masks contain 3 instances
raw_masks = dummy_raw_polygon_masks((3, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
flipped_masks = polygon_masks.flip(flip_direction='horizontal')
flipped_flipped_masks = flipped_masks.flip(flip_direction='horizontal')
assert len(flipped_masks) == 3
assert flipped_masks.height == 28
assert flipped_masks.width == 28
assert flipped_masks.to_ndarray().shape == (3, 28, 28)
assert (polygon_masks.to_ndarray() == flipped_flipped_masks.to_ndarray()
).all()
# vertically flip with polygon masks contain 3 instances
raw_masks = dummy_raw_polygon_masks((3, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
flipped_masks = polygon_masks.flip(flip_direction='vertical')
flipped_flipped_masks = flipped_masks.flip(flip_direction='vertical')
assert len(flipped_masks) == 3
assert flipped_masks.height == 28
assert flipped_masks.width == 28
assert flipped_masks.to_ndarray().shape == (3, 28, 28)
assert (polygon_masks.to_ndarray() == flipped_flipped_masks.to_ndarray()
).all()
def test_polygon_mask_resize():
# resize with empty polygon masks
raw_masks = dummy_raw_polygon_masks((0, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
resized_masks = polygon_masks.resize((56, 72))
assert len(resized_masks) == 0
assert resized_masks.height == 56
assert resized_masks.width == 72
assert resized_masks.to_ndarray().shape == (0, 56, 72)
# resize with polygon masks contain 1 instance 1 part
raw_masks1 = [[np.array([1, 1, 3, 1, 4, 3, 2, 4, 1, 3], dtype=np.float)]]
polygon_masks1 = PolygonMasks(raw_masks1, 5, 5)
resized_masks1 = polygon_masks1.resize((10, 10))
assert len(resized_masks1) == 1
assert resized_masks1.height == 10
assert resized_masks1.width == 10
assert resized_masks1.to_ndarray().shape == (1, 10, 10)
truth1 = np.array(
[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 1, 1, 1, 0, 0, 0, 0], [0, 0, 1, 1, 1, 1, 1, 0, 0, 0],
[0, 0, 1, 1, 1, 1, 1, 0, 0, 0], [0, 0, 1, 1, 1, 1, 1, 1, 0, 0],
[0, 0, 0, 1, 1, 1, 1, 0, 0, 0], [0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
np.uint8)
assert (resized_masks1.to_ndarray() == truth1).all()
# resize with polygon masks contain 1 instance 2 part
raw_masks2 = [[
np.array([0., 0., 1., 0., 1., 1.]),
np.array([1., 1., 2., 1., 2., 2., 1., 2.])
]]
polygon_masks2 = PolygonMasks(raw_masks2, 3, 3)
resized_masks2 = polygon_masks2.resize((6, 6))
assert len(resized_masks2) == 1
assert resized_masks2.height == 6
assert resized_masks2.width == 6
assert resized_masks2.to_ndarray().shape == (1, 6, 6)
truth2 = np.array(
[[0, 1, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 1, 1, 0, 0],
[0, 0, 1, 1, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]], np.uint8)
assert (resized_masks2.to_ndarray() == truth2).all()
# resize with polygon masks contain 2 instances
raw_masks3 = [raw_masks1[0], raw_masks2[0]]
polygon_masks3 = PolygonMasks(raw_masks3, 5, 5)
resized_masks3 = polygon_masks3.resize((10, 10))
assert len(resized_masks3) == 2
assert resized_masks3.height == 10
assert resized_masks3.width == 10
assert resized_masks3.to_ndarray().shape == (2, 10, 10)
truth3 = np.stack([truth1, np.pad(truth2, ((0, 4), (0, 4)), 'constant')])
assert (resized_masks3.to_ndarray() == truth3).all()
def test_gen_drrg_targets():
target_generator = textdet_targets.DRRGTargets()
assert np.allclose(target_generator.orientation_thr, 2.0)
assert np.allclose(target_generator.resample_step, 8.0)
assert target_generator.num_min_comps == 9
assert target_generator.num_max_comps == 600
assert np.allclose(target_generator.min_width, 8.0)
assert np.allclose(target_generator.max_width, 24.0)
assert np.allclose(target_generator.center_region_shrink_ratio, 0.3)
assert np.allclose(target_generator.comp_shrink_ratio, 1.0)
assert np.allclose(target_generator.comp_w_h_ratio, 0.3)
assert np.allclose(target_generator.text_comp_nms_thr, 0.25)
assert np.allclose(target_generator.min_rand_half_height, 8.0)
assert np.allclose(target_generator.max_rand_half_height, 24.0)
assert np.allclose(target_generator.jitter_level, 0.2)
# test generate_targets
target_generator = textdet_targets.DRRGTargets(min_width=2.,
max_width=4.,
min_rand_half_height=3.,
max_rand_half_height=5.)
results = {}
results['img'] = np.zeros((64, 64, 3), np.uint8)
text_polys = [[np.array([4, 2, 30, 2, 30, 10, 4, 10])],
[np.array([36, 12, 8, 12, 8, 22, 36, 22])],
[np.array([48, 20, 52, 20, 52, 50, 48, 50])],
[np.array([44, 50, 38, 50, 38, 20, 44, 20])]]
results['gt_masks'] = PolygonMasks(text_polys, 20, 30)
results['gt_masks_ignore'] = PolygonMasks([], 64, 64)
results['img_shape'] = (64, 64, 3)
results['mask_fields'] = []
output = target_generator(results)
assert len(output['gt_text_mask']) == 1
assert len(output['gt_center_region_mask']) == 1
assert len(output['gt_mask']) == 1
assert len(output['gt_top_height_map']) == 1
assert len(output['gt_bot_height_map']) == 1
assert len(output['gt_sin_map']) == 1
assert len(output['gt_cos_map']) == 1
assert output['gt_comp_attribs'].shape[-1] == 8
# test generate_targets with the number of proposed text components exceeds
# num_max_comps
target_generator = textdet_targets.DRRGTargets(min_width=2.,
max_width=4.,
min_rand_half_height=3.,
max_rand_half_height=5.,
num_max_comps=6)
output = target_generator(results)
assert output['gt_comp_attribs'].ndim == 2
assert output['gt_comp_attribs'].shape[0] == 6
# test generate_targets with blank polygon masks
target_generator = textdet_targets.DRRGTargets(min_width=2.,
max_width=4.,
min_rand_half_height=3.,
max_rand_half_height=5.)
results = {}
results['img'] = np.zeros((20, 30, 3), np.uint8)
results['gt_masks'] = PolygonMasks([], 20, 30)
results['gt_masks_ignore'] = PolygonMasks([], 20, 30)
results['img_shape'] = (20, 30, 3)
results['mask_fields'] = []
output = target_generator(results)
assert output['gt_comp_attribs'][0, 0] > 8
# test generate_targets with one proposed text component
text_polys = [[np.array([13, 6, 17, 6, 17, 14, 13, 14])]]
target_generator = textdet_targets.DRRGTargets(min_width=4.,
max_width=8.,
min_rand_half_height=3.,
max_rand_half_height=5.)
results['gt_masks'] = PolygonMasks(text_polys, 20, 30)
output = target_generator(results)
assert output['gt_comp_attribs'][0, 0] > 8
# test generate_targets with shrunk margin in generate_rand_comp_attribs
target_generator = textdet_targets.DRRGTargets(min_width=2.,
max_width=30.,
min_rand_half_height=3.,
max_rand_half_height=30.)
output = target_generator(results)
assert output['gt_comp_attribs'][0, 0] > 8
def test_gen_textsnake_targets(mock_show_feature):
target_generator = textdet_targets.TextSnakeTargets()
assert np.allclose(target_generator.orientation_thr, 2.0)
assert np.allclose(target_generator.resample_step, 4.0)
assert np.allclose(target_generator.center_region_shrink_ratio, 0.3)
# test find_head_tail for quadrangle
polygon = np.array([[1.0, 1.0], [5.0, 1.0], [5.0, 3.0], [1.0, 3.0]])
head_inds, tail_inds = target_generator.find_head_tail(polygon, 2.0)
assert np.allclose(head_inds, [3, 0])
assert np.allclose(tail_inds, [1, 2])
# test find_head_tail for polygon
polygon = np.array([[0., 10.], [3., 3.], [10., 0.], [17., 3.], [20., 10.],
[15., 10.], [13.5, 6.5], [10., 5.], [6.5, 6.5],
[5., 10.]])
head_inds, tail_inds = target_generator.find_head_tail(polygon, 2.0)
assert np.allclose(head_inds, [9, 0])
assert np.allclose(tail_inds, [4, 5])
# test generate_text_region_mask
img_size = (3, 10)
text_polys = [[np.array([0, 0, 1, 0, 1, 1, 0, 1])],
[np.array([2, 0, 3, 0, 3, 1, 2, 1])]]
output = target_generator.generate_text_region_mask(img_size, text_polys)
target = np.array([[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
assert np.allclose(output, target)
# test generate_center_region_mask
target_generator.center_region_shrink_ratio = 1.0
(center_region_mask, radius_map, sin_map,
cos_map) = target_generator.generate_center_mask_attrib_maps(
img_size, text_polys)
target = np.array([[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
assert np.allclose(center_region_mask, target)
assert np.allclose(sin_map, np.zeros(img_size))
assert np.allclose(cos_map, target)
# test generate_effective_mask
polys_ignore = text_polys
output = target_generator.generate_effective_mask(img_size, polys_ignore)
target = np.array([[0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
assert np.allclose(output, target)
# test generate_targets
results = {}
results['img'] = np.zeros((3, 10, 3), np.uint8)
results['gt_masks'] = PolygonMasks(text_polys, 3, 10)
results['gt_masks_ignore'] = PolygonMasks([], 3, 10)
results['img_shape'] = (3, 10, 3)
results['mask_fields'] = []
output = target_generator(results)
assert len(output['gt_text_mask']) == 1
assert len(output['gt_center_region_mask']) == 1
assert len(output['gt_mask']) == 1
assert len(output['gt_radius_map']) == 1
assert len(output['gt_sin_map']) == 1
assert len(output['gt_cos_map']) == 1
bundle = cf_bundle.CustomFormatBundle(keys=[
'gt_text_mask', 'gt_center_region_mask', 'gt_mask', 'gt_radius_map',
'gt_sin_map', 'gt_cos_map'
],
visualize=dict(
flag=True,
boundary_key='gt_text_mask'))
bundle(output)
assert 'gt_text_mask' in output.keys()
assert 'gt_center_region_mask' in output.keys()
assert 'gt_mask' in output.keys()
assert 'gt_radius_map' in output.keys()
assert 'gt_sin_map' in output.keys()
assert 'gt_cos_map' in output.keys()
mock_show_feature.assert_called_once()
