| | |
| | from mmdet.datasets.builder import DATASETS |
| |
|
| | from mmocr.core.evaluation.ner_metric import eval_ner_f1 |
| | from mmocr.datasets.base_dataset import BaseDataset |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class NerDataset(BaseDataset): |
| | """Custom dataset for named entity recognition tasks. |
| | |
| | Args: |
| | ann_file (txt): Annotation file path. |
| | loader (dict): Dictionary to construct loader |
| | to load annotation infos. |
| | pipeline (list[dict]): Processing pipeline. |
| | test_mode (bool, optional): If True, try...except will |
| | be turned off in __getitem__. |
| | """ |
| |
|
| | def prepare_train_img(self, index): |
| | """Get training data and annotations after pipeline. |
| | |
| | Args: |
| | index (int): Index of data. |
| | |
| | Returns: |
| | dict: Training data and annotation after pipeline with new keys \ |
| | introduced by pipeline. |
| | """ |
| | ann_info = self.data_infos[index] |
| |
|
| | return self.pipeline(ann_info) |
| |
|
| | def evaluate(self, results, metric=None, logger=None, **kwargs): |
| | """Evaluate the dataset. |
| | |
| | Args: |
| | results (list): Testing results of the dataset. |
| | metric (str | list[str]): Metrics to be evaluated. |
| | logger (logging.Logger | str | None): Logger used for printing |
| | related information during evaluation. Default: None. |
| | Returns: |
| | info (dict): A dict containing the following keys: |
| | 'acc', 'recall', 'f1-score'. |
| | """ |
| | gt_infos = list(self.data_infos) |
| | eval_results = eval_ner_f1(results, gt_infos) |
| | return eval_results |
| |
|
