56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import List
|
|
|
|
import mmengine
|
|
from mmengine.dataset import BaseDataset
|
|
|
|
from mmpretrain.registry import DATASETS
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class VSR(BaseDataset):
|
|
"""VSR: Visual Spatial Reasoning dataset.
|
|
|
|
Args:
|
|
data_root (str): The root directory for ``data_prefix``, ``ann_file``
|
|
and ``question_file``.
|
|
data_prefix (str): The directory of images.
|
|
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,
|
|
ann_file: str = '',
|
|
**kwarg):
|
|
super().__init__(
|
|
data_root=data_root,
|
|
data_prefix=dict(img_path=data_prefix),
|
|
ann_file=ann_file,
|
|
**kwarg,
|
|
)
|
|
|
|
def load_data_list(self) -> List[dict]:
|
|
"""Load data list."""
|
|
annotations = mmengine.load(self.ann_file)
|
|
|
|
data_list = []
|
|
for ann in annotations:
|
|
# ann example
|
|
# {
|
|
# "image": "train2017/000000372029.jpg",
|
|
# "question": "The dog is on the surfboard.",
|
|
# "answer": true
|
|
# }
|
|
data_info = dict()
|
|
data_info['img_path'] = mmengine.join_path(
|
|
self.data_prefix['img_path'], ann['image'])
|
|
data_info['question'] = ann['question']
|
|
data_info['gt_answer'] = 'yes' if ann['answer'] else 'no'
|
|
|
|
data_list.append(data_info)
|
|
|
|
return data_list
|