Spaces:
Running
Running
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import os.path as osp | |
| from typing import List | |
| from mmdet.datasets.api_wrappers import COCO | |
| from mmocr.datasets.preparers.parsers.base import BaseParser | |
| from mmocr.registry import DATA_PARSERS | |
| class COCOTextDetAnnParser(BaseParser): | |
| """COCO-like Format Text Detection Parser. | |
| Args: | |
| data_root (str): The root path of the dataset. Defaults to None. | |
| nproc (int): The number of processes to parse the annotation. Defaults | |
| to 1. | |
| variant (str): Variant of COCO dataset, options are ['standard', | |
| 'cocotext', 'textocr']. Defaults to 'standard'. | |
| """ | |
| def __init__(self, | |
| split: str, | |
| nproc: int = 1, | |
| variant: str = 'standard') -> None: | |
| super().__init__(nproc=nproc, split=split) | |
| assert variant in ['standard', 'cocotext', 'textocr'], \ | |
| f'variant {variant} is not supported' | |
| self.variant = variant | |
| def parse_files(self, img_dir: str, ann_path: str) -> List: | |
| """Parse single annotation.""" | |
| samples = list() | |
| coco = COCO(ann_path) | |
| if self.variant == 'cocotext' or self.variant == 'textocr': | |
| # cocotext stores both 'train' and 'val' split in one annotation | |
| # file, and uses the 'set' field to distinguish them. | |
| if self.variant == 'cocotext': | |
| for img in coco.dataset['imgs']: | |
| if self.split == coco.dataset['imgs'][img]['set']: | |
| coco.imgs[img] = coco.dataset['imgs'][img] | |
| # textocr stores 'train' and 'val'split separately | |
| elif self.variant == 'textocr': | |
| coco.imgs = coco.dataset['imgs'] | |
| # both cocotext and textocr stores the annotation ID in the | |
| # 'imgToAnns' field, so we need to convert it to the 'anns' field | |
| for img in coco.dataset['imgToAnns']: | |
| ann_ids = coco.dataset['imgToAnns'][img] | |
| anns = [ | |
| coco.dataset['anns'][str(ann_id)] for ann_id in ann_ids | |
| ] | |
| coco.dataset['imgToAnns'][img] = anns | |
| coco.imgToAnns = coco.dataset['imgToAnns'] | |
| coco.anns = coco.dataset['anns'] | |
| img_ids = coco.get_img_ids() | |
| total_ann_ids = [] | |
| for img_id in img_ids: | |
| img_info = coco.load_imgs([img_id])[0] | |
| img_info['img_id'] = img_id | |
| img_path = img_info['file_name'] | |
| ann_ids = coco.get_ann_ids(img_ids=[img_id]) | |
| if len(ann_ids) == 0: | |
| continue | |
| ann_ids = [str(ann_id) for ann_id in ann_ids] | |
| ann_info = coco.load_anns(ann_ids) | |
| total_ann_ids.extend(ann_ids) | |
| instances = list() | |
| for ann in ann_info: | |
| if self.variant == 'standard': | |
| # standard coco format use 'segmentation' field to store | |
| # the polygon and 'iscrowd' field to store the ignore flag, | |
| # and the 'text' field to store the text content. | |
| instances.append( | |
| dict( | |
| poly=ann['segmentation'][0], | |
| text=ann.get('text', None), | |
| ignore=ann.get('iscrowd', False))) | |
| elif self.variant == 'cocotext': | |
| # cocotext use 'utf8_string' field to store the text and | |
| # 'legibility' field to store the ignore flag, and the | |
| # 'mask' field to store the polygon. | |
| instances.append( | |
| dict( | |
| poly=ann['mask'], | |
| text=ann.get('utf8_string', None), | |
| ignore=ann['legibility'] == 'illegible')) | |
| elif self.variant == 'textocr': | |
| # textocr use 'utf8_string' field to store the text and | |
| # the 'points' field to store the polygon, '.' is used to | |
| # represent the ignored text. | |
| text = ann.get('utf8_string', None) | |
| instances.append( | |
| dict( | |
| poly=ann['points'], text=text, ignore=text == '.')) | |
| samples.append((osp.join(img_dir, | |
| osp.basename(img_path)), instances)) | |
| return samples | |