| | |
| | import os.path as osp |
| | import xml.etree.ElementTree as ET |
| | from typing import List, Optional, Union |
| |
|
| | import mmcv |
| | from mmengine.fileio import get, get_local_path, list_from_file |
| |
|
| | from mmdet.registry import DATASETS |
| | from .base_det_dataset import BaseDetDataset |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class XMLDataset(BaseDetDataset): |
| | """XML dataset for detection. |
| | |
| | Args: |
| | img_subdir (str): Subdir where images are stored. Default: JPEGImages. |
| | ann_subdir (str): Subdir where annotations are. Default: Annotations. |
| | backend_args (dict, optional): Arguments to instantiate the |
| | corresponding backend. Defaults to None. |
| | """ |
| |
|
| | def __init__(self, |
| | img_subdir: str = 'JPEGImages', |
| | ann_subdir: str = 'Annotations', |
| | **kwargs) -> None: |
| | self.img_subdir = img_subdir |
| | self.ann_subdir = ann_subdir |
| | super().__init__(**kwargs) |
| |
|
| | @property |
| | def sub_data_root(self) -> str: |
| | """Return the sub data root.""" |
| | return self.data_prefix.get('sub_data_root', '') |
| |
|
| | 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) |
| | for img_id in img_ids: |
| | file_name = osp.join(self.img_subdir, f'{img_id}.jpg') |
| | xml_path = osp.join(self.sub_data_root, self.ann_subdir, |
| | f'{img_id}.xml') |
| |
|
| | raw_img_info = {} |
| | raw_img_info['img_id'] = img_id |
| | raw_img_info['file_name'] = file_name |
| | raw_img_info['xml_path'] = xml_path |
| |
|
| | parsed_data_info = self.parse_data_info(raw_img_info) |
| | data_list.append(parsed_data_info) |
| | return data_list |
| |
|
| | @property |
| | def bbox_min_size(self) -> Optional[int]: |
| | """Return the minimum size of bounding boxes in the images.""" |
| | if self.filter_cfg is not None: |
| | return self.filter_cfg.get('bbox_min_size', None) |
| | else: |
| | return None |
| |
|
| | 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_path = osp.join(self.sub_data_root, img_info['file_name']) |
| | data_info['img_path'] = img_path |
| | data_info['img_id'] = img_info['img_id'] |
| | data_info['xml_path'] = img_info['xml_path'] |
| |
|
| | |
| | with get_local_path( |
| | img_info['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') |
| | if size is not None: |
| | width = int(size.find('width').text) |
| | height = int(size.find('height').text) |
| | else: |
| | img_bytes = get(img_path, backend_args=self.backend_args) |
| | img = mmcv.imfrombytes(img_bytes, backend='cv2') |
| | height, width = img.shape[:2] |
| | del img, img_bytes |
| |
|
| | data_info['height'] = height |
| | data_info['width'] = width |
| |
|
| | data_info['instances'] = self._parse_instance_info( |
| | raw_ann_info, minus_one=True) |
| |
|
| | return data_info |
| |
|
| | def _parse_instance_info(self, |
| | raw_ann_info: ET, |
| | minus_one: bool = True) -> List[dict]: |
| | """parse instance information. |
| | |
| | Args: |
| | raw_ann_info (ElementTree): ElementTree object. |
| | minus_one (bool): Whether to subtract 1 from the coordinates. |
| | Defaults to True. |
| | |
| | Returns: |
| | List[dict]: List of instances. |
| | """ |
| | instances = [] |
| | for obj in raw_ann_info.findall('object'): |
| | instance = {} |
| | name = obj.find('name').text |
| | if name not in self._metainfo['classes']: |
| | continue |
| | difficult = obj.find('difficult') |
| | difficult = 0 if difficult is None else int(difficult.text) |
| | bnd_box = obj.find('bndbox') |
| | bbox = [ |
| | int(float(bnd_box.find('xmin').text)), |
| | int(float(bnd_box.find('ymin').text)), |
| | int(float(bnd_box.find('xmax').text)), |
| | int(float(bnd_box.find('ymax').text)) |
| | ] |
| |
|
| | |
| | if minus_one: |
| | bbox = [x - 1 for x in bbox] |
| |
|
| | ignore = False |
| | if self.bbox_min_size is not None: |
| | assert not self.test_mode |
| | w = bbox[2] - bbox[0] |
| | h = bbox[3] - bbox[1] |
| | if w < self.bbox_min_size or h < self.bbox_min_size: |
| | ignore = True |
| | if difficult or ignore: |
| | instance['ignore_flag'] = 1 |
| | else: |
| | instance['ignore_flag'] = 0 |
| | instance['bbox'] = bbox |
| | instance['bbox_label'] = self.cat2label[name] |
| | instances.append(instance) |
| | return instances |
| |
|
| | def filter_data(self) -> List[dict]: |
| | """Filter annotations according to filter_cfg. |
| | |
| | Returns: |
| | List[dict]: Filtered results. |
| | """ |
| | if self.test_mode: |
| | return self.data_list |
| |
|
| | filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False) \ |
| | if self.filter_cfg is not None else False |
| | min_size = self.filter_cfg.get('min_size', 0) \ |
| | if self.filter_cfg is not None else 0 |
| |
|
| | valid_data_infos = [] |
| | for i, data_info in enumerate(self.data_list): |
| | width = data_info['width'] |
| | height = data_info['height'] |
| | if filter_empty_gt and len(data_info['instances']) == 0: |
| | continue |
| | if min(width, height) >= min_size: |
| | valid_data_infos.append(data_info) |
| |
|
| | return valid_data_infos |
| |
|
