| | |
| | from mmdet.datasets.builder import DATASETS |
| |
|
| | from mmocr.core.evaluation.ocr_metric import eval_ocr_metric |
| | from mmocr.datasets.base_dataset import BaseDataset |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class OCRDataset(BaseDataset): |
| |
|
| | def pre_pipeline(self, results): |
| | results['img_prefix'] = self.img_prefix |
| | results['text'] = results['img_info']['text'] |
| |
|
| | def evaluate(self, results, metric='acc', 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: |
| | dict[str: float] |
| | """ |
| | gt_texts = [] |
| | pred_texts = [] |
| | for i in range(len(self)): |
| | item_info = self.data_infos[i] |
| | text = item_info['text'] |
| | gt_texts.append(text) |
| | pred_texts.append(results[i]['text']) |
| |
|
| | eval_results = eval_ocr_metric(pred_texts, gt_texts) |
| |
|
| | return eval_results |
| |
|
