mmagic.datasets.mscoco_dataset 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import os
import random
from typing import Optional, Sequence, Union

import mmengine
from mmengine.fileio import get_file_backend

from mmagic.registry import DATASETS
from .basic_conditional_dataset import BasicConditionalDataset

[文档]class MSCoCoDataset(BasicConditionalDataset): """MSCoCo 2014 dataset. Args: ann_file (str): Annotation file path. Defaults to ''. metainfo (dict, optional): Meta information for dataset, such as class information. Defaults to None. data_root (str): The root directory for ``data_prefix`` and ``ann_file``. Defaults to ''. drop_caption_rate (float, optional): Rate of dropping caption, used for training. Defaults to 0.0. phase (str, optional): Subdataset used for certain phase, can be set to `train`, `test` and `val`. Defaults to 'train'. year (int, optional): Version of CoCo dataset, can be set to 2014 and 2017. Defaults to 2014. data_prefix (str | dict): Prefix for the data. Defaults to ''. extensions (Sequence[str]): A sequence of allowed extensions. Defaults to ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif'). lazy_init (bool): Whether to load annotation during instantiation. In some cases, such as visualization, only the meta information of the dataset is needed, which is not necessary to load annotation file. ``Basedataset`` can skip load annotations to save time by set ``lazy_init=False``. Defaults to False. caption_style (str): If you want to add a style description for each caption, you can set caption_style to your style prompt. For example, 'realistic style'. Defaults to empty str. **kwargs: Other keyword arguments in :class:`BaseDataset`. """
[文档] METAINFO = dict(dataset_type='text_image_dataset', task_name='editing')
def __init__(self, ann_file: str = '', metainfo: Optional[dict] = None, data_root: str = '', drop_caption_rate=0.0, phase='train', year=2014, data_prefix: Union[str, dict] = '', extensions: Sequence[str] = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif'), lazy_init: bool = False, classes: Union[str, Sequence[str], None] = None, caption_style: str = '', **kwargs): ann_file = os.path.join('annotations', 'captions_' + phase + f'{year}.json') if ann_file == '' else ann_file self.year = year assert self.year == 2014 or self.year == 2017, \ 'Caption is only supported in 2014 or 2017.' self.image_prename = '' if self.year == 2014: self.image_prename = 'COCO_' + phase + f'{year}_' self.phase = phase self.drop_rate = drop_caption_rate self.caption_style = caption_style super().__init__( ann_file=ann_file, metainfo=metainfo, data_root=data_root, data_prefix=data_prefix, extensions=extensions, lazy_init=lazy_init, classes=classes, **kwargs)
[文档] def load_data_list(self): """Load image paths and gt_labels.""" if self.img_prefix: file_backend = get_file_backend(uri=self.img_prefix) json_file = def add_prefix(filename, prefix=''): if not prefix: return filename else: return file_backend.join_path(prefix, filename) data_list = [] for item in json_file['annotations']: image_name = self.image_prename + str( item['image_id']).zfill(12) + '.jpg' img_path = add_prefix( os.path.join(self.phase + str(self.year), image_name), self.img_prefix) caption = item['caption'].lower() if self.caption_style != '': caption = caption + ' ' + self.caption_style info = { 'img_path': img_path, 'gt_prompt': caption if (self.phase != 'train' or self.drop_rate < 1e-6 or random.random() >= self.drop_rate) else '' } data_list.append(info) return data_list
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.