增加xfund RE任务

KevinNuNu · KevinNuNu · commit 201692166dc0 · 2023-03-25T23:33:56.000+08:00
diff --git a/configs/re/_base_/datasets/xfund_de.py b/configs/re/_base_/datasets/xfund_de.py
@@ -0,0 +1,15 @@
+xfund_de_re_data_root = 'data/xfund/de'
+
+xfund_de_re_train = dict(
+    type='REDataset',
+    data_root=xfund_de_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_de_re_test = dict(
+    type='REDataset',
+    data_root=xfund_de_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/configs/re/_base_/datasets/xfund_es.py b/configs/re/_base_/datasets/xfund_es.py
@@ -0,0 +1,15 @@
+xfund_es_re_data_root = 'data/xfund/es'
+
+xfund_es_re_train = dict(
+    type='REDataset',
+    data_root=xfund_es_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_es_re_test = dict(
+    type='REDataset',
+    data_root=xfund_es_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/configs/re/_base_/datasets/xfund_fr.py b/configs/re/_base_/datasets/xfund_fr.py
@@ -0,0 +1,15 @@
+xfund_fr_re_data_root = 'data/xfund/fr'
+
+xfund_fr_re_train = dict(
+    type='REDataset',
+    data_root=xfund_fr_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_fr_re_test = dict(
+    type='REDataset',
+    data_root=xfund_fr_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/configs/re/_base_/datasets/xfund_it.py b/configs/re/_base_/datasets/xfund_it.py
@@ -0,0 +1,15 @@
+xfund_it_re_data_root = 'data/xfund/it'
+
+xfund_it_re_train = dict(
+    type='REDataset',
+    data_root=xfund_it_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_it_re_test = dict(
+    type='REDataset',
+    data_root=xfund_it_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/configs/re/_base_/datasets/xfund_ja.py b/configs/re/_base_/datasets/xfund_ja.py
@@ -0,0 +1,15 @@
+xfund_ja_re_data_root = 'data/xfund/ja'
+
+xfund_ja_re_train = dict(
+    type='REDataset',
+    data_root=xfund_ja_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_ja_re_test = dict(
+    type='REDataset',
+    data_root=xfund_ja_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/configs/re/_base_/datasets/xfund_pt.py b/configs/re/_base_/datasets/xfund_pt.py
@@ -0,0 +1,15 @@
+xfund_pt_re_data_root = 'data/xfund/pt'
+
+xfund_pt_re_train = dict(
+    type='REDataset',
+    data_root=xfund_pt_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_pt_re_test = dict(
+    type='REDataset',
+    data_root=xfund_pt_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/configs/re/_base_/datasets/xfund_zh.py b/configs/re/_base_/datasets/xfund_zh.py
@@ -0,0 +1,15 @@
+xfund_zh_re_data_root = 'data/xfund/zh'
+
+xfund_zh_re_train = dict(
+    type='REDataset',
+    data_root=xfund_zh_re_data_root,
+    ann_file='re_train.json',
+    filter_cfg=dict(filter_empty_gt=True, min_size=32),
+    pipeline=None)
+
+xfund_zh_re_test = dict(
+    type='REDataset',
+    data_root=xfund_zh_re_data_root,
+    ann_file='re_test.json',
+    test_mode=True,
+    pipeline=None)
diff --git a/dataset_zoo/xfund/de/re.py b/dataset_zoo/xfund/de/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/dataset_zoo/xfund/es/re.py b/dataset_zoo/xfund/es/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/dataset_zoo/xfund/fr/re.py b/dataset_zoo/xfund/fr/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/dataset_zoo/xfund/it/re.py b/dataset_zoo/xfund/it/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/dataset_zoo/xfund/ja/re.py b/dataset_zoo/xfund/ja/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/dataset_zoo/xfund/pt/re.py b/dataset_zoo/xfund/pt/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/dataset_zoo/xfund/zh/re.py b/dataset_zoo/xfund/zh/re.py
@@ -0,0 +1,8 @@
+_base_ = ['ser.py']
+
+_base_.train_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.train_preparer.packer.type = 'REPacker'
+_base_.test_preparer.parser.type = 'XFUNDREAnnParser'
+_base_.test_preparer.packer.type = 'REPacker'
+
+config_generator = dict(type='REConfigGenerator')
diff --git a/mmocr/datasets/preparers/config_generators/__init__.py b/mmocr/datasets/preparers/config_generators/__init__.py
@@ -4,9 +4,10 @@
 from .textrecog_config_generator import TextRecogConfigGenerator
 from .textspotting_config_generator import TextSpottingConfigGenerator
 from .ser_config_generator import SERConfigGenerator
+from .re_config_generator import REConfigGenerator
 
 __all__ = [
     'BaseDatasetConfigGenerator', 'TextDetConfigGenerator',
     'TextRecogConfigGenerator', 'TextSpottingConfigGenerator',
-    'SERConfigGenerator'
+    'SERConfigGenerator', 'REConfigGenerator'
 ]
diff --git a/mmocr/datasets/preparers/config_generators/re_config_generator.py b/mmocr/datasets/preparers/config_generators/re_config_generator.py
@@ -0,0 +1,98 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List, Optional
+
+from mmocr.registry import CFG_GENERATORS
+from .base import BaseDatasetConfigGenerator
+
+
+@CFG_GENERATORS.register_module()
+class REConfigGenerator(BaseDatasetConfigGenerator):
+    """Text detection config generator.
+
+    Args:
+        data_root (str): The root path of the dataset.
+        dataset_name (str): The name of the dataset.
+        overwrite_cfg (bool): Whether to overwrite the dataset config file if
+            it already exists. If False, config generator will not generate new
+            config for datasets whose configs are already in base.
+        train_anns (List[Dict], optional): A list of train annotation files
+            to appear in the base configs. Defaults to
+            ``[dict(file='re_train.json', dataset_postfix='')]``.
+            Each element is typically a dict with the following fields:
+            - ann_file (str): The path to the annotation file relative to
+              data_root.
+            - dataset_postfix (str, optional): Affects the postfix of the
+              resulting variable in the generated config. If specified, the
+              dataset variable will be named in the form of
+              ``{dataset_name}_{dataset_postfix}_{task}_{split}``. Defaults to
+              None.
+        val_anns (List[Dict], optional): A list of val annotation files
+            to appear in the base configs, similar to ``train_anns``. Defaults
+            to [].
+        test_anns (List[Dict], optional): A list of test annotation files
+            to appear in the base configs, similar to ``train_anns``. Defaults
+            to ``[dict(file='re_test.json')]``.
+        config_path (str): Path to the configs. Defaults to 'configs/'.
+    """
+
+    def __init__(
+        self,
+        data_root: str,
+        dataset_name: str,
+        overwrite_cfg: bool = False,
+        train_anns: Optional[List[Dict]] = [
+            dict(ann_file='re_train.json', dataset_postfix='')
+        ],
+        val_anns: Optional[List[Dict]] = [],
+        test_anns: Optional[List[Dict]] = [
+            dict(ann_file='re_test.json', dataset_postfix='')
+        ],
+        config_path: str = 'configs/',
+    ) -> None:
+        if '/' in dataset_name:
+            dataset_name = '_'.join(dataset_name.split('/'))
+        super().__init__(
+            data_root=data_root,
+            task='re',
+            overwrite_cfg=overwrite_cfg,
+            dataset_name=dataset_name,
+            train_anns=train_anns,
+            val_anns=val_anns,
+            test_anns=test_anns,
+            config_path=config_path,
+        )
+
+    def _gen_dataset_config(self) -> str:
+        """Generate a full dataset config based on the annotation file
+        dictionary.
+
+        Args:
+            ann_dict (dict[str, dict(str, str)]): A nested dictionary that maps
+                a config variable name (such as icdar2015_textrecog_train) to
+                its corresponding annotation information dict. Each dict
+                contains following keys:
+                - ann_file (str): The path to the annotation file relative to
+                  data_root.
+                - dataset_postfix (str, optional): Affects the postfix of the
+                  resulting variable in the generated config. If specified, the
+                  dataset variable will be named in the form of
+                  ``{dataset_name}_{dataset_postfix}_{task}_{split}``. Defaults
+                  to None.
+                - split (str): The split the annotation belongs to. Usually
+                  it can be 'train', 'val' and 'test'.
+
+        Returns:
+            str: The generated dataset config.
+        """
+        cfg = ''
+        for key_name, ann_dict in self.anns.items():
+            cfg += f'\n{key_name} = dict(\n'
+            cfg += '    type=\'REDataset\',\n'
+            cfg += '    data_root=' + f'{self.dataset_name}_{self.task}_data_root,\n'  # noqa: E501
+            cfg += f'    ann_file=\'{ann_dict["ann_file"]}\',\n'
+            if ann_dict['split'] == 'train':
+                cfg += '    filter_cfg=dict(filter_empty_gt=True, min_size=32),\n'  # noqa: E501
+            elif ann_dict['split'] in ['test', 'val']:
+                cfg += '    test_mode=True,\n'
+            cfg += '    pipeline=None)\n'
+        return cfg
diff --git a/mmocr/datasets/preparers/packers/__init__.py b/mmocr/datasets/preparers/packers/__init__.py
@@ -5,8 +5,9 @@
 from .textspotting_packer import TextSpottingPacker
 from .wildreceipt_packer import WildReceiptPacker
 from .ser_packer import SERPacker
+from .re_packer import REPacker
 
 __all__ = [
     'BasePacker', 'TextDetPacker', 'TextRecogPacker', 'TextRecogCropPacker',
-    'TextSpottingPacker', 'WildReceiptPacker', 'SERPacker'
+    'TextSpottingPacker', 'WildReceiptPacker', 'SERPacker', 'REPacker'
 ]
diff --git a/mmocr/datasets/preparers/packers/re_packer.py b/mmocr/datasets/preparers/packers/re_packer.py
diff --git a/mmocr/datasets/preparers/packers/ser_packer.py b/mmocr/datasets/preparers/packers/ser_packer.py
