Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import json | |
| import math | |
| import os.path as osp | |
| import tempfile | |
| import torch | |
| from mmocr.datasets.openset_kie_dataset import OpensetKIEDataset | |
| from mmocr.utils import list_to_file | |
| def _create_dummy_ann_file(ann_file): | |
| ann_info1 = { | |
| 'file_name': | |
| '1.png', | |
| 'height': | |
| 200, | |
| 'width': | |
| 200, | |
| 'annotations': [{ | |
| 'text': 'store', | |
| 'box': [11.0, 0.0, 22.0, 0.0, 12.0, 12.0, 0.0, 12.0], | |
| 'label': 1, | |
| 'edge': 1 | |
| }, { | |
| 'text': 'MyFamily', | |
| 'box': [23.0, 2.0, 31.0, 1.0, 24.0, 11.0, 16.0, 11.0], | |
| 'label': 2, | |
| 'edge': 1 | |
| }] | |
| } | |
| list_to_file(ann_file, [json.dumps(ann_info1)]) | |
| return ann_info1 | |
| def _create_dummy_dict_file(dict_file): | |
| dict_str = '0123' | |
| list_to_file(dict_file, list(dict_str)) | |
| def _create_dummy_loader(): | |
| loader = dict( | |
| type='HardDiskLoader', | |
| repeat=1, | |
| parser=dict( | |
| type='LineJsonParser', | |
| keys=['file_name', 'height', 'width', 'annotations'])) | |
| return loader | |
| def test_openset_kie_dataset(): | |
| with tempfile.TemporaryDirectory() as tmp_dir_name: | |
| # create dummy data | |
| ann_file = osp.join(tmp_dir_name, 'fake_data.txt') | |
| ann_info1 = _create_dummy_ann_file(ann_file) | |
| dict_file = osp.join(tmp_dir_name, 'fake_dict.txt') | |
| _create_dummy_dict_file(dict_file) | |
| # test initialization | |
| loader = _create_dummy_loader() | |
| dataset = OpensetKIEDataset(ann_file, loader, dict_file, pipeline=[]) | |
| dataset.prepare_train_img(0) | |
| # test pre_pipeline | |
| img_ann_info = dataset.data_infos[0] | |
| img_info = { | |
| 'filename': img_ann_info['file_name'], | |
| 'height': img_ann_info['height'], | |
| 'width': img_ann_info['width'] | |
| } | |
| ann_info = dataset._parse_anno_info(img_ann_info['annotations']) | |
| results = dict(img_info=img_info, ann_info=ann_info) | |
| dataset.pre_pipeline(results) | |
| assert results['img_prefix'] == dataset.img_prefix | |
| assert 'ori_texts' in results | |
| # test evaluation | |
| result = { | |
| 'img_metas': [{ | |
| 'filename': ann_info1['file_name'], | |
| 'ori_filename': ann_info1['file_name'], | |
| 'ori_texts': [], | |
| 'ori_boxes': [] | |
| }] | |
| } | |
| for anno in ann_info1['annotations']: | |
| result['img_metas'][0]['ori_texts'].append(anno['text']) | |
| result['img_metas'][0]['ori_boxes'].append(anno['box']) | |
| result['nodes'] = torch.tensor([[0.01, 0.8, 0.01, 0.18], | |
| [0.01, 0.01, 0.9, 0.08]]) | |
| result['edges'] = torch.Tensor([[0.01, 0.99] for _ in range(4)]) | |
| eval_res = dataset.evaluate([result]) | |
| assert math.isclose(eval_res['edge_openset_f1'], 1.0, abs_tol=1e-4) | |