add dataset alias for PaddleX (#10528)

This commit is contained in:
wuyefeilin 2023-08-02 10:59:32 +08:00 committed by GitHub
parent 51aa02f97d
commit 233658adb1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -39,7 +39,16 @@ from ppocr.data.pgnet_dataset import PGDataSet
from ppocr.data.pubtab_dataset import PubTabDataSet
from ppocr.data.multi_scale_sampler import MultiScaleSampler
__all__ = ['build_dataloader', 'transform', 'create_operators', 'set_signal_handlers']
# for PaddleX dataset_type
TextDetDataset = SimpleDataSet
TextRecDataset = SimpleDataSet
MSTextRecDataset = MultiScaleDataSet
PubTabTableRecDataset = PubTabDataSet
KieDataset = SimpleDataSet
__all__ = [
'build_dataloader', 'transform', 'create_operators', 'set_signal_handlers'
]
def term_mp(sig_num, frame):
@ -76,8 +85,18 @@ def build_dataloader(config, mode, device, logger, seed=None):
config = copy.deepcopy(config)
support_dict = [
'SimpleDataSet', 'LMDBDataSet', 'PGDataSet', 'PubTabDataSet',
'LMDBDataSetSR', 'LMDBDataSetTableMaster', 'MultiScaleDataSet'
'SimpleDataSet',
'LMDBDataSet',
'PGDataSet',
'PubTabDataSet',
'LMDBDataSetSR',
'LMDBDataSetTableMaster',
'MultiScaleDataSet',
'TextDetDataset',
'TextRecDataset',
'MSTextRecDataset',
'PubTabTableRecDataset',
'KieDataset',
]
module_name = config[mode]['dataset']['name']
assert module_name in support_dict, Exception(