96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import re
|
|
from itertools import chain
|
|
from typing import List
|
|
|
|
import mmengine
|
|
from mmengine.dataset import BaseDataset
|
|
|
|
from mmpretrain.registry import DATASETS
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class VisualGenomeQA(BaseDataset):
|
|
"""Visual Genome Question Answering dataset.
|
|
|
|
dataset structure: ::
|
|
|
|
data_root
|
|
├── image
|
|
│ ├── 1.jpg
|
|
│ ├── 2.jpg
|
|
│ └── ...
|
|
└── question_answers.json
|
|
|
|
Args:
|
|
data_root (str): The root directory for ``data_prefix``, ``ann_file``
|
|
and ``question_file``.
|
|
data_prefix (str): The directory of images. Defaults to ``"image"``.
|
|
ann_file (str, optional): Annotation file path for training and
|
|
validation. Defaults to ``"question_answers.json"``.
|
|
**kwargs: Other keyword arguments in :class:`BaseDataset`.
|
|
"""
|
|
|
|
def __init__(self,
|
|
data_root: str,
|
|
data_prefix: str = 'image',
|
|
ann_file: str = 'question_answers.json',
|
|
**kwarg):
|
|
super().__init__(
|
|
data_root=data_root,
|
|
data_prefix=dict(img_path=data_prefix),
|
|
ann_file=ann_file,
|
|
**kwarg,
|
|
)
|
|
|
|
def _create_image_index(self):
|
|
img_prefix = self.data_prefix['img_path']
|
|
|
|
files = mmengine.list_dir_or_file(img_prefix, list_dir=False)
|
|
image_index = {}
|
|
for file in files:
|
|
image_id = re.findall(r'\d+', file)
|
|
if len(image_id) > 0:
|
|
image_id = int(image_id[-1])
|
|
image_index[image_id] = mmengine.join_path(img_prefix, file)
|
|
|
|
return image_index
|
|
|
|
def load_data_list(self) -> List[dict]:
|
|
"""Load data list."""
|
|
annotations = mmengine.load(self.ann_file)
|
|
|
|
# The original Visual Genome annotation file and question file includes
|
|
# only image id but no image file paths.
|
|
self.image_index = self._create_image_index()
|
|
|
|
data_list = []
|
|
for qas in chain.from_iterable(ann['qas'] for ann in annotations):
|
|
# ann example
|
|
# {
|
|
# 'id': 1,
|
|
# 'qas': [
|
|
# {
|
|
# 'a_objects': [],
|
|
# 'question': 'What color is the clock?',
|
|
# 'image_id': 1,
|
|
# 'qa_id': 986768,
|
|
# 'answer': 'Two.',
|
|
# 'q_objects': [],
|
|
# }
|
|
# ...
|
|
# ]
|
|
# }
|
|
|
|
data_info = {
|
|
'img_path': self.image_index[qas['image_id']],
|
|
'quesiton': qas['quesiton'],
|
|
'question_id': qas['question_id'],
|
|
'image_id': qas['image_id'],
|
|
'gt_answer': [qas['answer']],
|
|
}
|
|
|
|
data_list.append(data_info)
|
|
|
|
return data_list
|