|
|
|
|
|
import os.path as osp |
|
|
import xml.etree.ElementTree as ET |
|
|
|
|
|
from mmengine.dist import is_main_process |
|
|
from mmengine.fileio import get_local_path, list_from_file |
|
|
from mmengine.utils import ProgressBar |
|
|
|
|
|
from mmdet.registry import DATASETS |
|
|
from mmdet.utils.typing_utils import List, Union |
|
|
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 |
|
|
""" |
|
|
METAINFO = {'classes': ('face', ), 'palette': [(0, 255, 0)]} |
|
|
|
|
|
def load_data_list(self) -> List[dict]: |
|
|
"""Load annotation from XML style ann_file. |
|
|
|
|
|
Returns: |
|
|
list[dict]: Annotation info from XML file. |
|
|
""" |
|
|
assert self._metainfo.get('classes', None) is not None, \ |
|
|
'classes in `XMLDataset` can not be None.' |
|
|
self.cat2label = { |
|
|
cat: i |
|
|
for i, cat in enumerate(self._metainfo['classes']) |
|
|
} |
|
|
|
|
|
data_list = [] |
|
|
img_ids = list_from_file(self.ann_file, backend_args=self.backend_args) |
|
|
|
|
|
|
|
|
if is_main_process(): |
|
|
prog_bar = ProgressBar(len(img_ids)) |
|
|
|
|
|
for img_id in img_ids: |
|
|
raw_img_info = {} |
|
|
raw_img_info['img_id'] = img_id |
|
|
raw_img_info['file_name'] = f'{img_id}.jpg' |
|
|
parsed_data_info = self.parse_data_info(raw_img_info) |
|
|
data_list.append(parsed_data_info) |
|
|
|
|
|
if is_main_process(): |
|
|
prog_bar.update() |
|
|
return data_list |
|
|
|
|
|
def parse_data_info(self, img_info: dict) -> Union[dict, List[dict]]: |
|
|
"""Parse raw annotation to target format. |
|
|
|
|
|
Args: |
|
|
img_info (dict): Raw image information, usually it includes |
|
|
`img_id`, `file_name`, and `xml_path`. |
|
|
|
|
|
Returns: |
|
|
Union[dict, List[dict]]: Parsed annotation. |
|
|
""" |
|
|
data_info = {} |
|
|
img_id = img_info['img_id'] |
|
|
xml_path = osp.join(self.data_prefix['img'], 'Annotations', |
|
|
f'{img_id}.xml') |
|
|
data_info['img_id'] = img_id |
|
|
data_info['xml_path'] = xml_path |
|
|
|
|
|
|
|
|
with get_local_path( |
|
|
xml_path, backend_args=self.backend_args) as local_path: |
|
|
raw_ann_info = ET.parse(local_path) |
|
|
root = raw_ann_info.getroot() |
|
|
size = root.find('size') |
|
|
width = int(size.find('width').text) |
|
|
height = int(size.find('height').text) |
|
|
folder = root.find('folder').text |
|
|
img_path = osp.join(self.data_prefix['img'], folder, |
|
|
img_info['file_name']) |
|
|
data_info['img_path'] = img_path |
|
|
|
|
|
data_info['height'] = height |
|
|
data_info['width'] = width |
|
|
|
|
|
|
|
|
data_info['instances'] = self._parse_instance_info( |
|
|
raw_ann_info, minus_one=False) |
|
|
return data_info |
|
|
|
