113 lines
3.6 KiB
Python
113 lines
3.6 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os.path as osp
|
|
from typing import List
|
|
|
|
import mmengine
|
|
import numpy as np
|
|
from mmengine.dataset import BaseDataset
|
|
from pycocotools.coco import COCO
|
|
|
|
from mmpretrain.registry import DATASETS
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class RefCOCO(BaseDataset):
|
|
"""RefCOCO dataset.
|
|
|
|
RefCOCO is a popular dataset used for the task of visual grounding.
|
|
Here are the steps for accessing and utilizing the
|
|
RefCOCO dataset.
|
|
|
|
You can access the RefCOCO dataset from the official source:
|
|
https://github.com/lichengunc/refer
|
|
|
|
The RefCOCO dataset is organized in a structured format: ::
|
|
|
|
FeaturesDict({
|
|
'coco_annotations': Sequence({
|
|
'area': int64,
|
|
'bbox': BBoxFeature(shape=(4,), dtype=float32),
|
|
'id': int64,
|
|
'label': int64,
|
|
}),
|
|
'image': Image(shape=(None, None, 3), dtype=uint8),
|
|
'image/id': int64,
|
|
'objects': Sequence({
|
|
'area': int64,
|
|
'bbox': BBoxFeature(shape=(4,), dtype=float32),
|
|
'gt_box_index': int64,
|
|
'id': int64,
|
|
'label': int64,
|
|
'refexp': Sequence({
|
|
'raw': Text(shape=(), dtype=string),
|
|
'refexp_id': int64,
|
|
}),
|
|
}),
|
|
})
|
|
|
|
Args:
|
|
ann_file (str): Annotation file path.
|
|
data_root (str): The root directory for ``data_prefix`` and
|
|
``ann_file``. Defaults to ''.
|
|
data_prefix (str): Prefix for training data.
|
|
pipeline (Sequence): Processing pipeline. Defaults to an empty tuple.
|
|
**kwargs: Other keyword arguments in :class:`BaseDataset`.
|
|
"""
|
|
|
|
def __init__(self,
|
|
data_root,
|
|
ann_file,
|
|
data_prefix,
|
|
split_file,
|
|
split='train',
|
|
**kwargs):
|
|
self.split_file = split_file
|
|
self.split = split
|
|
|
|
super().__init__(
|
|
data_root=data_root,
|
|
data_prefix=dict(img_path=data_prefix),
|
|
ann_file=ann_file,
|
|
**kwargs,
|
|
)
|
|
|
|
def _join_prefix(self):
|
|
if not mmengine.is_abs(self.split_file) and self.split_file:
|
|
self.split_file = osp.join(self.data_root, self.split_file)
|
|
|
|
return super()._join_prefix()
|
|
|
|
def load_data_list(self) -> List[dict]:
|
|
"""Load data list."""
|
|
with mmengine.get_local_path(self.ann_file) as ann_file:
|
|
coco = COCO(ann_file)
|
|
splits = mmengine.load(self.split_file, file_format='pkl')
|
|
img_prefix = self.data_prefix['img_path']
|
|
|
|
data_list = []
|
|
join_path = mmengine.fileio.get_file_backend(img_prefix).join_path
|
|
for refer in splits:
|
|
if refer['split'] != self.split:
|
|
continue
|
|
|
|
ann = coco.anns[refer['ann_id']]
|
|
img = coco.imgs[ann['image_id']]
|
|
sentences = refer['sentences']
|
|
bbox = np.array(ann['bbox'], dtype=np.float32)
|
|
bbox[2:4] = bbox[0:2] + bbox[2:4] # XYWH -> XYXY
|
|
|
|
for sent in sentences:
|
|
data_info = {
|
|
'img_path': join_path(img_prefix, img['file_name']),
|
|
'image_id': ann['image_id'],
|
|
'ann_id': ann['id'],
|
|
'text': sent['sent'],
|
|
'gt_bboxes': bbox[None, :],
|
|
}
|
|
data_list.append(data_info)
|
|
|
|
if len(data_list) == 0:
|
|
raise ValueError(f'No sample in split "{self.split}".')
|
|
|
|
return data_list
|