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