cifar.py 5.85 KB
Newer Older
Soumith Chintala's avatar
Soumith Chintala committed
1
2
3
4
from PIL import Image
import os
import os.path
import numpy as np
5
import pickle
Philip Meier's avatar
Philip Meier committed
6
from typing import Any, Callable, Optional, Tuple
Soumith Chintala's avatar
Soumith Chintala committed
7

8
from .vision import VisionDataset
9
from .utils import check_integrity, download_and_extract_archive
10

11

12
class CIFAR10(VisionDataset):
13
14
15
16
    """`CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.

    Args:
        root (string): Root directory of dataset where directory
17
            ``cifar-10-batches-py`` exists or will be saved to if download is set to True.
18
19
        train (bool, optional): If True, creates dataset from training set, otherwise
            creates from test set.
Tongzhou Wang's avatar
Tongzhou Wang committed
20
        transform (callable, optional): A function/transform that takes in an PIL image
21
22
23
24
25
26
27
28
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.

    """
Soumith Chintala's avatar
Soumith Chintala committed
29
    base_folder = 'cifar-10-batches-py'
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
30
    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
Soumith Chintala's avatar
Soumith Chintala committed
31
    filename = "cifar-10-python.tar.gz"
zhoumingjun's avatar
zhoumingjun committed
32
    tgz_md5 = 'c58f30108f718f92721af3b95e74349a'
Soumith Chintala's avatar
Soumith Chintala committed
33
    train_list = [
34
35
36
37
38
        ['data_batch_1', 'c99cafc152244af753f735de768cd75f'],
        ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'],
        ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'],
        ['data_batch_4', '634d18415352ddfa80567beed471001a'],
        ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'],
Soumith Chintala's avatar
Soumith Chintala committed
39
40
41
    ]

    test_list = [
42
        ['test_batch', '40351d587109b95175f43aff81a1287e'],
Soumith Chintala's avatar
Soumith Chintala committed
43
    ]
44
45
46
47
48
49
    meta = {
        'filename': 'batches.meta',
        'key': 'label_names',
        'md5': '5ff9c542aee3614f3951f8cda6e48888',
    }

Philip Meier's avatar
Philip Meier committed
50
51
52
53
54
55
56
57
    def __init__(
            self,
            root: str,
            train: bool = True,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            download: bool = False,
    ) -> None:
58

59
60
        super(CIFAR10, self).__init__(root, transform=transform,
                                      target_transform=target_transform)
61

62
63
        self.train = train  # training set or test set

Soumith Chintala's avatar
Soumith Chintala committed
64
65
66
67
        if download:
            self.download()

        if not self._check_integrity():
68
69
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can use download=True to download it')
70

71
        if self.train:
72
73
74
75
            downloaded_list = self.train_list
        else:
            downloaded_list = self.test_list

Philip Meier's avatar
Philip Meier committed
76
        self.data: Any = []
77
78
79
80
81
82
        self.targets = []

        # now load the picked numpy arrays
        for file_name, checksum in downloaded_list:
            file_path = os.path.join(self.root, self.base_folder, file_name)
            with open(file_path, 'rb') as f:
83
                entry = pickle.load(f, encoding='latin1')
84
                self.data.append(entry['data'])
85
                if 'labels' in entry:
86
                    self.targets.extend(entry['labels'])
87
                else:
88
                    self.targets.extend(entry['fine_labels'])
89

90
91
        self.data = np.vstack(self.data).reshape(-1, 3, 32, 32)
        self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
Soumith Chintala's avatar
Soumith Chintala committed
92

93
94
        self._load_meta()

Philip Meier's avatar
Philip Meier committed
95
    def _load_meta(self) -> None:
96
97
98
99
100
        path = os.path.join(self.root, self.base_folder, self.meta['filename'])
        if not check_integrity(path, self.meta['md5']):
            raise RuntimeError('Dataset metadata file not found or corrupted.' +
                               ' You can use download=True to download it')
        with open(path, 'rb') as infile:
101
            data = pickle.load(infile, encoding='latin1')
102
103
104
            self.classes = data[self.meta['key']]
        self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)}

Philip Meier's avatar
Philip Meier committed
105
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
106
107
108
109
110
111
112
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
113
        img, target = self.data[index], self.targets[index]
114

115
116
        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
117
        img = Image.fromarray(img)
Soumith Chintala's avatar
Soumith Chintala committed
118
119
120
121
122
123
124
125
126

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

Philip Meier's avatar
Philip Meier committed
127
    def __len__(self) -> int:
128
        return len(self.data)
Soumith Chintala's avatar
Soumith Chintala committed
129

Philip Meier's avatar
Philip Meier committed
130
    def _check_integrity(self) -> bool:
Soumith Chintala's avatar
Soumith Chintala committed
131
        root = self.root
132
        for fentry in (self.train_list + self.test_list):
Soumith Chintala's avatar
Soumith Chintala committed
133
134
            filename, md5 = fentry[0], fentry[1]
            fpath = os.path.join(root, self.base_folder, filename)
soumith's avatar
soumith committed
135
            if not check_integrity(fpath, md5):
Soumith Chintala's avatar
Soumith Chintala committed
136
137
138
                return False
        return True

Philip Meier's avatar
Philip Meier committed
139
    def download(self) -> None:
Soumith Chintala's avatar
Soumith Chintala committed
140
141
142
        if self._check_integrity():
            print('Files already downloaded and verified')
            return
143
        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)
Soumith Chintala's avatar
Soumith Chintala committed
144

Philip Meier's avatar
Philip Meier committed
145
    def extra_repr(self) -> str:
146
        return "Split: {}".format("Train" if self.train is True else "Test")
147

Soumith Chintala's avatar
Soumith Chintala committed
148
149

class CIFAR100(CIFAR10):
jvmancuso's avatar
jvmancuso committed
150
151
152
153
    """`CIFAR100 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.

    This is a subclass of the `CIFAR10` Dataset.
    """
Soumith Chintala's avatar
Soumith Chintala committed
154
    base_folder = 'cifar-100-python'
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
155
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
Soumith Chintala's avatar
Soumith Chintala committed
156
157
158
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
    train_list = [
159
        ['train', '16019d7e3df5f24257cddd939b257f8d'],
Soumith Chintala's avatar
Soumith Chintala committed
160
161
162
    ]

    test_list = [
163
        ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
Soumith Chintala's avatar
Soumith Chintala committed
164
    ]
165
166
167
168
169
    meta = {
        'filename': 'meta',
        'key': 'fine_label_names',
        'md5': '7973b15100ade9c7d40fb424638fde48',
    }