| | |
| | import os.path as osp |
| | from typing import List |
| |
|
| | from pycocotools.coco import COCO |
| |
|
| | from mmdet.registry import DATASETS |
| | from .base_det_dataset import BaseDetDataset |
| |
|
| |
|
| | def convert_phrase_ids(phrase_ids: list) -> list: |
| | unique_elements = sorted(set(phrase_ids)) |
| | element_to_new_label = { |
| | element: label |
| | for label, element in enumerate(unique_elements) |
| | } |
| | phrase_ids = [element_to_new_label[element] for element in phrase_ids] |
| | return phrase_ids |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class Flickr30kDataset(BaseDetDataset): |
| | """Flickr30K Dataset.""" |
| |
|
| | def load_data_list(self) -> List[dict]: |
| |
|
| | self.coco = COCO(self.ann_file) |
| |
|
| | self.ids = sorted(list(self.coco.imgs.keys())) |
| |
|
| | data_list = [] |
| | for img_id in self.ids: |
| | if isinstance(img_id, str): |
| | ann_ids = self.coco.getAnnIds(imgIds=[img_id], iscrowd=None) |
| | else: |
| | ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=None) |
| |
|
| | coco_img = self.coco.loadImgs(img_id)[0] |
| |
|
| | caption = coco_img['caption'] |
| | file_name = coco_img['file_name'] |
| | img_path = osp.join(self.data_prefix['img'], file_name) |
| | width = coco_img['width'] |
| | height = coco_img['height'] |
| | tokens_positive = coco_img['tokens_positive_eval'] |
| | phrases = [caption[i[0][0]:i[0][1]] for i in tokens_positive] |
| | phrase_ids = [] |
| |
|
| | instances = [] |
| | annos = self.coco.loadAnns(ann_ids) |
| | for anno in annos: |
| | instance = { |
| | 'bbox': [ |
| | anno['bbox'][0], anno['bbox'][1], |
| | anno['bbox'][0] + anno['bbox'][2], |
| | anno['bbox'][1] + anno['bbox'][3] |
| | ], |
| | 'bbox_label': |
| | anno['category_id'], |
| | 'ignore_flag': |
| | anno['iscrowd'] |
| | } |
| | phrase_ids.append(anno['phrase_ids']) |
| | instances.append(instance) |
| |
|
| | phrase_ids = convert_phrase_ids(phrase_ids) |
| |
|
| | data_list.append( |
| | dict( |
| | img_path=img_path, |
| | img_id=img_id, |
| | height=height, |
| | width=width, |
| | instances=instances, |
| | text=caption, |
| | phrase_ids=phrase_ids, |
| | tokens_positive=tokens_positive, |
| | phrases=phrases, |
| | )) |
| |
|
| | return data_list |
| |
|
