| |
| import os.path as osp |
| from typing import List |
|
|
| from mmengine.fileio import get_local_path |
|
|
| from mmdet.datasets import BaseDetDataset |
| from mmdet.registry import DATASETS |
| from .api_wrappers import COCO |
|
|
|
|
| @DATASETS.register_module() |
| class MDETRStyleRefCocoDataset(BaseDetDataset): |
| """RefCOCO dataset. |
| |
| Only support evaluation now. |
| """ |
|
|
| def load_data_list(self) -> List[dict]: |
| with get_local_path( |
| self.ann_file, backend_args=self.backend_args) as local_path: |
| coco = COCO(local_path) |
|
|
| img_ids = coco.get_img_ids() |
|
|
| data_infos = [] |
| for img_id in img_ids: |
| raw_img_info = coco.load_imgs([img_id])[0] |
| ann_ids = coco.get_ann_ids(img_ids=[img_id]) |
| raw_ann_info = coco.load_anns(ann_ids) |
|
|
| data_info = {} |
| img_path = osp.join(self.data_prefix['img'], |
| raw_img_info['file_name']) |
| data_info['img_path'] = img_path |
| data_info['img_id'] = img_id |
| data_info['height'] = raw_img_info['height'] |
| data_info['width'] = raw_img_info['width'] |
| data_info['dataset_mode'] = raw_img_info['dataset_name'] |
|
|
| data_info['text'] = raw_img_info['caption'] |
| data_info['custom_entities'] = False |
| data_info['tokens_positive'] = -1 |
|
|
| instances = [] |
| for i, ann in enumerate(raw_ann_info): |
| instance = {} |
| x1, y1, w, h = ann['bbox'] |
| bbox = [x1, y1, x1 + w, y1 + h] |
| instance['bbox'] = bbox |
| instance['bbox_label'] = ann['category_id'] |
| instance['ignore_flag'] = 0 |
| instances.append(instance) |
|
|
| data_info['instances'] = instances |
| data_infos.append(data_info) |
| return data_infos |
|
|
