mmcls.datasets.cub 源代码

# Copyright (c) OpenMMLab. All rights reserved.
from typing import List

from mmengine import get_file_backend, list_from_file

from mmcls.registry import DATASETS
from .base_dataset import BaseDataset
from .categories import CUB_CATEGORIES

[文档]@DATASETS.register_module() class CUB(BaseDataset): """The CUB-200-2011 Dataset. Support the `CUB-200-2011 <>`_ Dataset. Comparing with the `CUB-200 <>`_ Dataset, there are much more pictures in `CUB-200-2011`. After downloading and decompression, the dataset directory structure is as follows. CUB dataset directory: :: CUB-200-2011 (data_root)/ ├── images (data_prefix) │ ├── class_x │ │ ├── xx1.jpg │ │ ├── xx2.jpg │ │ └── ... │ ├── class_y │ │ ├── yy1.jpg │ │ ├── yy2.jpg │ │ └── ... │ └── ... ├── images.txt (ann_file) ├── image_class_labels.txt (image_class_labels_file) ├── train_test_split.txt (train_test_split_file) └── .... Args: data_root (str): The root directory for CUB-200-2011 dataset. test_mode (bool): ``test_mode=True`` means in test phase. It determines to use the training set or test set. ann_file (str, optional): Annotation file path, path relative to ``data_root``. Defaults to 'images.txt'. data_prefix (str): Prefix for iamges, path relative to ``data_root``. Defaults to 'images'. image_class_labels_file (str, optional): The label file, path relative to ``data_root``. Defaults to 'image_class_labels.txt'. train_test_split_file (str, optional): The split file to split train and test dataset, path relative to ``data_root``. Defaults to 'train_test_split_file.txt'. Examples: >>> from mmcls.datasets import CUB >>> cub_train_cfg = dict(data_root='data/CUB_200_2011', test_mode=True) >>> cub_train = CUB(**cub_train_cfg) >>> cub_train Dataset CUB Number of samples: 5994 Number of categories: 200 Root of dataset: data/CUB_200_2011 >>> cub_test_cfg = dict(data_root='data/CUB_200_2011', test_mode=True) >>> cub_test = CUB(**cub_test_cfg) >>> cub_test Dataset CUB Number of samples: 5794 Number of categories: 200 Root of dataset: data/CUB_200_2011 """ # noqa: E501 METAINFO = {'classes': CUB_CATEGORIES} def __init__(self, data_root: str, test_mode: bool, ann_file: str = 'images.txt', data_prefix: str = 'images', image_class_labels_file: str = 'image_class_labels.txt', train_test_split_file: str = 'train_test_split.txt', **kwargs): self.backend = get_file_backend(data_root, enable_singleton=True) self.image_class_labels_file = self.backend.join_path( data_root, image_class_labels_file) self.train_test_split_file = self.backend.join_path( data_root, train_test_split_file) super(CUB, self).__init__( ann_file=ann_file, data_root=data_root, data_prefix=data_prefix, test_mode=test_mode, **kwargs) def _load_data_from_txt(self, filepath): """load data from CUB txt file, the every line of the file is idx and a data item.""" pairs = list_from_file(filepath) data_dict = dict() for pair in pairs: idx, data_item = pair.split() # all the index starts from 1 in CUB files, # here we need to '- 1' to let them start from 0. data_dict[int(idx) - 1] = data_item return data_dict def load_data_list(self): """Load images and ground truth labels.""" sample_dict = self._load_data_from_txt(self.ann_file) label_dict = self._load_data_from_txt(self.image_class_labels_file) split_dict = self._load_data_from_txt(self.train_test_split_file) assert sample_dict.keys() == label_dict.keys() == split_dict.keys(),\ f'sample_ids should be same in files {self.ann_file}, ' \ f'{self.image_class_labels_file} and {self.train_test_split_file}' data_list = [] for sample_id in sample_dict.keys(): if split_dict[sample_id] == '1' and self.test_mode: # skip train samples when test_mode=True continue elif split_dict[sample_id] == '0' and not self.test_mode: # skip test samples when test_mode=False continue img_path = self.backend.join_path(self.img_prefix, sample_dict[sample_id]) gt_label = int(label_dict[sample_id]) - 1 info = dict(img_path=img_path, gt_label=gt_label) data_list.append(info) return data_list def extra_repr(self) -> List[str]: """The extra repr information of the dataset.""" body = [ f'Root of dataset: \t{self.data_root}', ] return body
Read the Docs v: mmcls-1.x
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.