Spaces:
Runtime error
Runtime error
File size: 2,642 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 |
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import numpy as np
from mmocr.datasets.pipelines import LoadImageFromNdarray, LoadTextAnnotations
def _create_dummy_ann():
results = {}
results['img_info'] = {}
results['img_info']['height'] = 1000
results['img_info']['width'] = 1000
results['ann_info'] = {}
results['ann_info']['masks'] = []
results['mask_fields'] = []
results['ann_info']['masks_ignore'] = [
[[499, 94, 531, 94, 531, 124, 499, 124]],
[[3, 156, 81, 155, 78, 181, 0, 182]],
[[11, 223, 59, 221, 59, 234, 11, 236]],
[[500, 156, 551, 156, 550, 165, 499, 165]]
]
return results
def test_loadtextannotation():
results = _create_dummy_ann()
with_bbox = True
with_label = True
with_mask = True
with_seg = False
poly2mask = False
# If no 'ori_shape' in result but use_img_shape=True,
# result['img_info']['height'] and result['img_info']['width']
# will be used to generate mask.
loader = LoadTextAnnotations(
with_bbox,
with_label,
with_mask,
with_seg,
poly2mask,
use_img_shape=True)
tmp_results = copy.deepcopy(results)
output = loader._load_masks(tmp_results)
assert len(output['gt_masks_ignore']) == 4
assert np.allclose(output['gt_masks_ignore'].masks[0],
[[499, 94, 531, 94, 531, 124, 499, 124]])
assert output['gt_masks_ignore'].height == results['img_info']['height']
# If 'ori_shape' in result and use_img_shape=True,
# result['ori_shape'] will be used to generate mask.
loader = LoadTextAnnotations(
with_bbox,
with_label,
with_mask,
with_seg,
poly2mask=True,
use_img_shape=True)
tmp_results = copy.deepcopy(results)
tmp_results['ori_shape'] = (640, 640, 3)
output = loader._load_masks(tmp_results)
assert output['img_info']['height'] == 640
assert output['gt_masks_ignore'].height == 640
def test_load_img_from_numpy():
result = {'img': np.ones((32, 100, 3), dtype=np.uint8)}
load = LoadImageFromNdarray(color_type='color')
output = load(result)
assert output['img'].shape[2] == 3
assert len(output['img'].shape) == 3
result = {'img': np.ones((32, 100, 1), dtype=np.uint8)}
load = LoadImageFromNdarray(color_type='color')
output = load(result)
assert output['img'].shape[2] == 3
result = {'img': np.ones((32, 100, 3), dtype=np.uint8)}
load = LoadImageFromNdarray(color_type='grayscale', to_float32=True)
output = load(result)
assert output['img'].shape[2] == 1
|