mmpretrain.datasets.refcoco 源代码

# 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: 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
Read the Docs v: stable
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.