Mountchicken's picture
Upload 704 files
9bf4bd7
raw
history blame
3.45 kB
# Copyright (c) OpenMMLab. All rights reserved.
import xml.etree.ElementTree as ET
from typing import List, Tuple
import numpy as np
from mmocr.datasets.preparers.data_preparer import DATA_PARSERS
from mmocr.datasets.preparers.parsers.base import BaseParser
from mmocr.utils import list_from_file
@DATA_PARSERS.register_module()
class CTW1500AnnParser(BaseParser):
"""SCUT-CTW1500 dataset parser.
Args:
ignore (str): The text of the ignored instances. Defaults to
'###'.
"""
def __init__(self, ignore: str = '###', **kwargs) -> None:
self.ignore = ignore
super().__init__(**kwargs)
def parse_file(self, img_path: str, ann_path: str) -> Tuple:
"""Convert annotation for a single image.
Args:
img_path (str): The path of image.
ann_path (str): The path of annotation.
Returns:
Tuple: A tuple of (img_path, instance).
- img_path (str): The path of image file, which can be read
directly by opencv.
- instance: instance is a list of dict containing parsed
annotations, which should contain the following keys:
- 'poly' or 'box' (textdet or textspotting)
- 'text' (textspotting or textrecog)
- 'ignore' (all task)
Examples:
An example of returned values:
>>> ('imgs/train/xxx.jpg',
>>> dict(
>>> poly=[[[0, 1], [1, 1], [1, 0], [0, 0]]],
>>> text='hello',
>>> ignore=False)
>>> )
"""
if self.split == 'train':
instances = self.load_xml_info(ann_path)
elif self.split == 'test':
instances = self.load_txt_info(ann_path)
return img_path, instances
def load_txt_info(self, anno_dir: str) -> List:
"""Load the annotation of the SCUT-CTW dataset (test split).
Args:
anno_dir (str): Path to the annotation file.
Returns:
list[Dict]: List of instances.
"""
instances = list()
for line in list_from_file(anno_dir):
# each line has one ploygen (n vetices), and one text.
# e.g., 695,885,866,888,867,1146,696,1143,####Latin 9
line = line.strip()
strs = line.split(',')
assert strs[28][0] == '#'
xy = [int(x) for x in strs[0:28]]
assert len(xy) == 28
poly = np.array(xy).reshape(-1).tolist()
text = strs[28][4:]
instances.append(
dict(poly=poly, text=text, ignore=text == self.ignore))
return instances
def load_xml_info(self, anno_dir: str) -> List:
"""Load the annotation of the SCUT-CTW dataset (train split).
Args:
anno_dir (str): Path to the annotation file.
Returns:
list[Dict]: List of instances.
"""
obj = ET.parse(anno_dir)
instances = list()
for image in obj.getroot(): # image
for box in image: # image
text = box[0].text
segs = box[1].text
pts = segs.strip().split(',')
pts = [int(x) for x in pts]
assert len(pts) == 28
poly = np.array(pts).reshape(-1).tolist()
instances.append(dict(poly=poly, text=text, ignore=0))
return instances