| | |
| | import os.path as osp |
| | from typing import List, Optional |
| |
|
| | import numpy as np |
| |
|
| | from mmdet.registry import DATASETS |
| | from .base_det_dataset import BaseDetDataset |
| |
|
| | try: |
| | from d_cube import D3 |
| | except ImportError: |
| | D3 = None |
| | from .api_wrappers import COCO |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class DODDataset(BaseDetDataset): |
| |
|
| | def __init__(self, |
| | *args, |
| | data_root: Optional[str] = '', |
| | data_prefix: dict = dict(img_path=''), |
| | **kwargs) -> None: |
| | if D3 is None: |
| | raise ImportError( |
| | 'Please install d3 by `pip install ddd-dataset`.') |
| | pkl_anno_path = osp.join(data_root, data_prefix['anno']) |
| | self.img_root = osp.join(data_root, data_prefix['img']) |
| | self.d3 = D3(self.img_root, pkl_anno_path) |
| |
|
| | sent_infos = self.d3.load_sents() |
| | classes = tuple([sent_info['raw_sent'] for sent_info in sent_infos]) |
| | super().__init__( |
| | *args, |
| | data_root=data_root, |
| | data_prefix=data_prefix, |
| | metainfo={'classes': classes}, |
| | **kwargs) |
| |
|
| | def load_data_list(self) -> List[dict]: |
| | coco = COCO(self.ann_file) |
| | data_list = [] |
| | img_ids = self.d3.get_img_ids() |
| | for img_id in img_ids: |
| | data_info = {} |
| |
|
| | img_info = self.d3.load_imgs(img_id)[0] |
| | file_name = img_info['file_name'] |
| | img_path = osp.join(self.img_root, file_name) |
| | data_info['img_path'] = img_path |
| | data_info['img_id'] = img_id |
| | data_info['height'] = img_info['height'] |
| | data_info['width'] = img_info['width'] |
| |
|
| | group_ids = self.d3.get_group_ids(img_ids=[img_id]) |
| | sent_ids = self.d3.get_sent_ids(group_ids=group_ids) |
| | sent_list = self.d3.load_sents(sent_ids=sent_ids) |
| | text_list = [sent['raw_sent'] for sent in sent_list] |
| | ann_ids = coco.get_ann_ids(img_ids=[img_id]) |
| | anno = coco.load_anns(ann_ids) |
| |
|
| | data_info['text'] = text_list |
| | data_info['sent_ids'] = np.array([s for s in sent_ids]) |
| | data_info['custom_entities'] = True |
| |
|
| | instances = [] |
| | for i, ann in enumerate(anno): |
| | instance = {} |
| | x1, y1, w, h = ann['bbox'] |
| | bbox = [x1, y1, x1 + w, y1 + h] |
| | instance['ignore_flag'] = 0 |
| | instance['bbox'] = bbox |
| | instance['bbox_label'] = ann['category_id'] - 1 |
| | instances.append(instance) |
| | data_info['instances'] = instances |
| | data_list.append(data_info) |
| | return data_list |
| |
|
