Spaces:
Sleeping
Sleeping
File size: 4,496 Bytes
9bf4bd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
# 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
@DATA_PARSERS.register_module()
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
|