115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os.path as osp
|
|
import re
|
|
from collections import Counter
|
|
from typing import List
|
|
|
|
import mmengine
|
|
from mmengine.dataset import BaseDataset
|
|
|
|
from mmpretrain.registry import DATASETS
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class COCOVQA(BaseDataset):
|
|
"""VQAv2 dataset.
|
|
|
|
Args:
|
|
data_root (str): The root directory for ``data_prefix``, ``ann_file``
|
|
and ``question_file``.
|
|
data_prefix (str): The directory of images.
|
|
question_file (str): Question file path.
|
|
ann_file (str, optional): Annotation file path for training and
|
|
validation. Defaults to an empty string.
|
|
**kwargs: Other keyword arguments in :class:`BaseDataset`.
|
|
"""
|
|
|
|
def __init__(self,
|
|
data_root: str,
|
|
data_prefix: str,
|
|
question_file: str,
|
|
ann_file: str = '',
|
|
**kwarg):
|
|
self.question_file = question_file
|
|
super().__init__(
|
|
data_root=data_root,
|
|
data_prefix=dict(img_path=data_prefix),
|
|
ann_file=ann_file,
|
|
**kwarg,
|
|
)
|
|
|
|
def _join_prefix(self):
|
|
if not mmengine.is_abs(self.question_file) and self.question_file:
|
|
self.question_file = osp.join(self.data_root, self.question_file)
|
|
|
|
return super()._join_prefix()
|
|
|
|
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{12}', 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."""
|
|
questions = mmengine.load(self.question_file)['questions']
|
|
if self.ann_file:
|
|
annotations = mmengine.load(self.ann_file)['annotations']
|
|
assert len(questions) == len(annotations)
|
|
else:
|
|
annotations = [None] * len(questions)
|
|
|
|
# The original VQAv2 annotation file and question file includes
|
|
# only image id but no image file paths.
|
|
self.image_index = self._create_image_index()
|
|
|
|
data_list = []
|
|
for question, ann in zip(questions, annotations):
|
|
# question example
|
|
# {
|
|
# 'image_id': 262144,
|
|
# 'question': "Is the ball flying towards the batter?",
|
|
# 'question_id': 262144000
|
|
# }
|
|
#
|
|
# ann example
|
|
# {
|
|
# 'question_type': "what are the",
|
|
# 'answer_type': "other",
|
|
# 'answers': [
|
|
# {'answer': 'watching',
|
|
# 'answer_id': 1,
|
|
# 'answer_confidence': 'yes'},
|
|
# ...
|
|
# ],
|
|
# 'image_id': 262148,
|
|
# 'question_id': 262148000,
|
|
# 'multiple_choice_answer': 'watching',
|
|
# 'answer_type': 'other',
|
|
# }
|
|
|
|
data_info = question
|
|
data_info['img_path'] = self.image_index[question['image_id']]
|
|
|
|
if ann is not None:
|
|
assert ann['question_id'] == question['question_id']
|
|
|
|
# add answer_weight & answer_count, delete duplicate answer
|
|
answers = [item['answer'] for item in ann.pop('answers')]
|
|
count = Counter(answers)
|
|
answer_weight = [i / len(answers) for i in count.values()]
|
|
data_info['gt_answer'] = list(count.keys())
|
|
data_info['gt_answer_weight'] = answer_weight
|
|
data_info.update(ann)
|
|
|
|
data_list.append(data_info)
|
|
|
|
return data_list
|