|
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 |
|
|
