import os.path as osp import xml.etree.ElementTree as ET import mmcv from .builder import DATASETS from .xml_style import XMLDataset @DATASETS.register_module() class WIDERFaceDataset(XMLDataset): """Reader for the WIDER Face dataset in PASCAL VOC format. Conversion scripts can be found in https://github.com/sovrasov/wider-face-pascal-voc-annotations """ CLASSES = ('face', ) def __init__(self, **kwargs): super(WIDERFaceDataset, self).__init__(**kwargs) def load_annotations(self, ann_file): """Load annotation from WIDERFace XML style annotation file. Args: ann_file (str): Path of XML file. Returns: list[dict]: Annotation info from XML file. """ data_infos = [] img_ids = mmcv.list_from_file(ann_file) for img_id in img_ids: filename = f'{img_id}.jpg' xml_path = osp.join(self.img_prefix, 'Annotations', f'{img_id}.xml') tree = ET.parse(xml_path) root = tree.getroot() size = root.find('size') width = int(size.find('width').text) height = int(size.find('height').text) folder = root.find('folder').text data_infos.append( dict( id=img_id, filename=osp.join(folder, filename), width=width, height=height)) return data_infos