# Copyright (c) OpenMMLab. All rights reserved. import pickle from typing import List, Optional import mmengine.dist as dist import numpy as np from mmengine import FileClient from mmcls.registry import DATASETS from .base_dataset import BaseDataset from .categories import CIFAR10_CATEGORIES, CIFAR100_CATEGORIES from .utils import check_md5, download_and_extract_archive @DATASETS.register_module() class CIFAR10(BaseDataset): """`CIFAR10 `_ Dataset. This implementation is modified from https://github.com/pytorch/vision/blob/master/torchvision/datasets/cifar.py Args: data_prefix (str): Prefix for data. test_mode (bool): ``test_mode=True`` means in test phase. It determines to use the training set or test set. metainfo (dict, optional): Meta information for dataset, such as categories information. Defaults to None. data_root (str, optional): The root directory for ``data_prefix``. Defaults to None. download (bool): Whether to download the dataset if not exists. Defaults to True. **kwargs: Other keyword arguments in :class:`BaseDataset`. """ # noqa: E501 base_folder = 'cifar-10-batches-py' url = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz' filename = 'cifar-10-python.tar.gz' tgz_md5 = 'c58f30108f718f92721af3b95e74349a' train_list = [ ['data_batch_1', 'c99cafc152244af753f735de768cd75f'], ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'], ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'], ['data_batch_4', '634d18415352ddfa80567beed471001a'], ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'], ] test_list = [ ['test_batch', '40351d587109b95175f43aff81a1287e'], ] meta = { 'filename': 'batches.meta', 'key': 'label_names', 'md5': '5ff9c542aee3614f3951f8cda6e48888', } METAINFO = {'classes': CIFAR10_CATEGORIES} def __init__(self, data_prefix: str, test_mode: bool, metainfo: Optional[dict] = None, data_root: Optional[str] = None, download: bool = True, **kwargs): self.download = download super().__init__( # The CIFAR dataset doesn't need specify annotation file ann_file='', metainfo=metainfo, data_root=data_root, data_prefix=dict(root=data_prefix), test_mode=test_mode, **kwargs) def load_data_list(self): """Load images and ground truth labels.""" root_prefix = self.data_prefix['root'] file_client = FileClient.infer_client(uri=root_prefix) if dist.is_main_process() and not self._check_integrity(): if file_client.name != 'HardDiskBackend': raise RuntimeError( f'The dataset on {root_prefix} is not integrated, ' f'please manually handle it.') if self.download: download_and_extract_archive( self.url, root_prefix, filename=self.filename, md5=self.tgz_md5) else: raise RuntimeError( f'Cannot find {self.__class__.__name__} dataset in ' f"{self.data_prefix['root']}, you can specify " '`download=True` to download automatically.') dist.barrier() assert self._check_integrity(), \ 'Download failed or shared storage is unavailable. Please ' \ f'download the dataset manually through {self.url}.' if not self.test_mode: downloaded_list = self.train_list else: downloaded_list = self.test_list imgs = [] gt_labels = [] # load the picked numpy arrays for file_name, _ in downloaded_list: file_path = file_client.join_path(root_prefix, self.base_folder, file_name) content = file_client.get(file_path) entry = pickle.loads(content, encoding='latin1') imgs.append(entry['data']) if 'labels' in entry: gt_labels.extend(entry['labels']) else: gt_labels.extend(entry['fine_labels']) imgs = np.vstack(imgs).reshape(-1, 3, 32, 32) imgs = imgs.transpose((0, 2, 3, 1)) # convert to HWC if self.CLASSES is None: # The metainfo in the file has the lowest priority, therefore # we only need to load it if classes is not specified. self._load_meta() data_list = [] for img, gt_label in zip(imgs, gt_labels): info = {'img': img, 'gt_label': int(gt_label)} data_list.append(info) return data_list def _load_meta(self): """Load categories information from metafile.""" root = self.data_prefix['root'] file_client = FileClient.infer_client(uri=root) path = file_client.join_path(root, self.base_folder, self.meta['filename']) md5 = self.meta.get('md5', None) if not file_client.exists(path) or (md5 is not None and not check_md5(path, md5)): raise RuntimeError( 'Dataset metadata file not found or corrupted.' + ' You can use `download=True` to download it') content = file_client.get(path) data = pickle.loads(content, encoding='latin1') self._metainfo.setdefault('classes', data[self.meta['key']]) def _check_integrity(self): """Check the integrity of data files.""" root = self.data_prefix['root'] file_client = FileClient.infer_client(uri=root) for fentry in (self.train_list + self.test_list): filename, md5 = fentry[0], fentry[1] fpath = file_client.join_path(root, self.base_folder, filename) if not file_client.exists(fpath): return False if md5 is not None and not check_md5( fpath, md5, file_client=file_client): return False return True def extra_repr(self) -> List[str]: """The extra repr information of the dataset.""" body = [f"Prefix of data: \t{self.data_prefix['root']}"] return body @DATASETS.register_module() class CIFAR100(CIFAR10): """`CIFAR100 `_ Dataset. Args: data_prefix (str): Prefix for data. test_mode (bool): ``test_mode=True`` means in test phase. It determines to use the training set or test set. metainfo (dict, optional): Meta information for dataset, such as categories information. Defaults to None. data_root (str, optional): The root directory for ``data_prefix``. Defaults to None. download (bool): Whether to download the dataset if not exists. Defaults to True. **kwargs: Other keyword arguments in :class:`BaseDataset`. """ base_folder = 'cifar-100-python' url = 'https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz' filename = 'cifar-100-python.tar.gz' tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85' train_list = [ ['train', '16019d7e3df5f24257cddd939b257f8d'], ] test_list = [ ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'], ] meta = { 'filename': 'meta', 'key': 'fine_label_names', 'md5': '7973b15100ade9c7d40fb424638fde48', } METAINFO = {'classes': CIFAR100_CATEGORIES}