mirror of https://github.com/open-mmlab/mmocr.git
52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from unittest import TestCase
|
|
|
|
from mmocr.datasets import RecogTextDataset
|
|
|
|
|
|
class TestRecogTextDataset(TestCase):
|
|
|
|
def test_txt_dataset(self):
|
|
|
|
# test initialization
|
|
dataset = RecogTextDataset(
|
|
ann_file='tests/data/rec_toy_dataset/old_label.txt',
|
|
data_prefix=dict(img_path='imgs'),
|
|
parser_cfg=dict(
|
|
type='LineStrParser',
|
|
keys=['filename', 'text'],
|
|
keys_idx=[0, 1]),
|
|
pipeline=[])
|
|
dataset.full_init()
|
|
self.assertEqual(len(dataset), 10)
|
|
self.assertEqual(len(dataset.load_data_list()), 10)
|
|
|
|
# test load_data_list
|
|
anno = dataset.load_data_list()
|
|
self.assertIn(anno[0]['img_path'],
|
|
['imgs/1223731.jpg', 'imgs\\1223731.jpg'])
|
|
self.assertEqual(anno[0]['instances'][0]['text'], 'GRAND')
|
|
self.assertIn(anno[1]['img_path'],
|
|
['imgs/1223733.jpg', 'imgs\\1223733.jpg'])
|
|
|
|
self.assertEqual(anno[1]['instances'][0]['text'], 'HOTEL')
|
|
|
|
def test_jsonl_dataset(self):
|
|
dataset = RecogTextDataset(
|
|
ann_file='tests/data/rec_toy_dataset/old_label.jsonl',
|
|
data_prefix=dict(img_path='imgs'),
|
|
parser_cfg=dict(type='LineJsonParser', keys=['filename', 'text']),
|
|
pipeline=[])
|
|
dataset.full_init()
|
|
self.assertEqual(len(dataset), 10)
|
|
self.assertEqual(len(dataset.load_data_list()), 10)
|
|
|
|
# test load_data_list
|
|
anno = dataset.load_data_list()
|
|
self.assertIn(anno[0]['img_path'],
|
|
['imgs/1223731.jpg', 'imgs\\1223731.jpg'])
|
|
self.assertEqual(anno[0]['instances'][0]['text'], 'GRAND')
|
|
self.assertIn(anno[1]['img_path'],
|
|
['imgs/1223733.jpg', 'imgs\\1223733.jpg'])
|
|
self.assertEqual(anno[1]['instances'][0]['text'], 'HOTEL')
|