149 lines
5.4 KiB
Python
149 lines
5.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import json
|
|
import os.path as osp
|
|
from collections import OrderedDict
|
|
from os import PathLike
|
|
from typing import List, Sequence, Union
|
|
|
|
from mmengine import get_file_backend
|
|
|
|
from mmpretrain.registry import DATASETS, TRANSFORMS
|
|
from .base_dataset import BaseDataset
|
|
|
|
|
|
def expanduser(data_prefix):
|
|
if isinstance(data_prefix, (str, PathLike)):
|
|
return osp.expanduser(data_prefix)
|
|
else:
|
|
return data_prefix
|
|
|
|
|
|
@DATASETS.register_module()
|
|
class COCORetrieval(BaseDataset):
|
|
"""COCO Retrieval dataset.
|
|
|
|
COCO (Common Objects in Context): The COCO dataset contains more than
|
|
330K images,each of which has approximately 5 descriptive annotations.
|
|
This dataset was releasedin collaboration between Microsoft and Carnegie
|
|
Mellon University
|
|
|
|
COCO_2014 dataset directory: ::
|
|
|
|
COCO_2014
|
|
├── val2014
|
|
├── train2014
|
|
├── annotations
|
|
├── instances_train2014.json
|
|
├── instances_val2014.json
|
|
├── person_keypoints_train2014.json
|
|
├── person_keypoints_val2014.json
|
|
├── captions_train2014.json
|
|
├── captions_val2014.json
|
|
|
|
Args:
|
|
ann_file (str): Annotation file path.
|
|
test_mode (bool): Whether dataset is used for evaluation. This will
|
|
decide the annotation format in data list annotations.
|
|
Defaults to False.
|
|
data_root (str): The root directory for ``data_prefix`` and
|
|
``ann_file``. Defaults to ''.
|
|
data_prefix (str | dict): Prefix for training data. Defaults to ''.
|
|
pipeline (Sequence): Processing pipeline. Defaults to an empty tuple.
|
|
**kwargs: Other keyword arguments in :class:`BaseDataset`.
|
|
|
|
Examples:
|
|
>>> from mmpretrain.datasets import COCORetrieval
|
|
>>> train_dataset=COCORetrieval(data_root='coco2014/')
|
|
>>> train_dataset
|
|
Dataset COCORetrieval
|
|
Number of samples: 414113
|
|
Annotation file: /coco2014/annotations/captions_train2014.json
|
|
Prefix of images: /coco2014/
|
|
>>> from mmpretrain.datasets import COCORetrieval
|
|
>>> val_dataset = COCORetrieval(data_root='coco2014/')
|
|
>>> val_dataset
|
|
Dataset COCORetrieval
|
|
Number of samples: 202654
|
|
Annotation file: /coco2014/annotations/captions_val2014.json
|
|
Prefix of images: /coco2014/
|
|
"""
|
|
|
|
def __init__(self,
|
|
ann_file: str,
|
|
test_mode: bool = False,
|
|
data_prefix: Union[str, dict] = '',
|
|
data_root: str = '',
|
|
pipeline: Sequence = (),
|
|
**kwargs):
|
|
|
|
if isinstance(data_prefix, str):
|
|
data_prefix = dict(img_path=expanduser(data_prefix))
|
|
|
|
ann_file = expanduser(ann_file)
|
|
transforms = []
|
|
for transform in pipeline:
|
|
if isinstance(transform, dict):
|
|
transforms.append(TRANSFORMS.build(transform))
|
|
else:
|
|
transforms.append(transform)
|
|
|
|
super().__init__(
|
|
data_root=data_root,
|
|
data_prefix=data_prefix,
|
|
test_mode=test_mode,
|
|
pipeline=transforms,
|
|
ann_file=ann_file,
|
|
**kwargs,
|
|
)
|
|
|
|
def load_data_list(self) -> List[dict]:
|
|
"""Load data list."""
|
|
# get file backend
|
|
img_prefix = self.data_prefix['img_path']
|
|
file_backend = get_file_backend(img_prefix)
|
|
|
|
anno_info = json.load(open(self.ann_file, 'r'))
|
|
# mapping img_id to img filename
|
|
img_dict = OrderedDict()
|
|
for idx, img in enumerate(anno_info['images']):
|
|
if img['id'] not in img_dict:
|
|
img_rel_path = img['coco_url'].rsplit('/', 2)[-2:]
|
|
img_path = file_backend.join_path(img_prefix, *img_rel_path)
|
|
|
|
# create new idx for image
|
|
img_dict[img['id']] = dict(
|
|
ori_id=img['id'],
|
|
image_id=idx, # will be used for evaluation
|
|
img_path=img_path,
|
|
text=[],
|
|
gt_text_id=[],
|
|
gt_image_id=[],
|
|
)
|
|
|
|
train_list = []
|
|
for idx, anno in enumerate(anno_info['annotations']):
|
|
anno['text'] = anno.pop('caption')
|
|
anno['ori_id'] = anno.pop('id')
|
|
anno['text_id'] = idx # will be used for evaluation
|
|
# 1. prepare train data list item
|
|
train_data = anno.copy()
|
|
train_image = img_dict[train_data['image_id']]
|
|
train_data['img_path'] = train_image['img_path']
|
|
train_data['image_ori_id'] = train_image['ori_id']
|
|
train_data['image_id'] = train_image['image_id']
|
|
train_data['is_matched'] = True
|
|
train_list.append(train_data)
|
|
# 2. prepare eval data list item based on img dict
|
|
img_dict[anno['image_id']]['gt_text_id'].append(anno['text_id'])
|
|
img_dict[anno['image_id']]['text'].append(anno['text'])
|
|
img_dict[anno['image_id']]['gt_image_id'].append(
|
|
train_image['image_id'])
|
|
|
|
self.img_size = len(img_dict)
|
|
self.text_size = len(anno_info['annotations'])
|
|
|
|
# return needed format data list
|
|
if self.test_mode:
|
|
return list(img_dict.values())
|
|
return train_list
|