stl10.py 7.28 KB
Newer Older
Elad Hoffer's avatar
Elad Hoffer committed
1
2
3
4
from PIL import Image
import os
import os.path
import numpy as np
5
from typing import Any, Callable, Optional, Tuple
Elad Hoffer's avatar
Elad Hoffer committed
6

Francisco Massa's avatar
Francisco Massa committed
7
from .vision import VisionDataset
8
from .utils import check_integrity, download_and_extract_archive, verify_str_arg
Elad Hoffer's avatar
Elad Hoffer committed
9

Francisco Massa's avatar
Francisco Massa committed
10
11

class STL10(VisionDataset):
12
13
14
15
16
17
18
    """`STL10 <https://cs.stanford.edu/~acoates/stl10/>`_ Dataset.

    Args:
        root (string): Root directory of dataset where directory
            ``stl10_binary`` exists.
        split (string): One of {'train', 'test', 'unlabeled', 'train+unlabeled'}.
            Accordingly dataset is selected.
19
20
        folds (int, optional): One of {0-9} or None.
            For training, loads one of the 10 pre-defined folds of 1k samples for the
21
            standard evaluation procedure. If no value is passed, loads the 5k samples.
22
23
24
25
26
27
28
29
        transform (callable, optional): A function/transform that  takes in an PIL image
            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.
    """
Elad Hoffer's avatar
Elad Hoffer committed
30
31
32
33
34
    base_folder = 'stl10_binary'
    url = "http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz"
    filename = "stl10_binary.tar.gz"
    tgz_md5 = '91f7769df0f17e558f3565bffb0c7dfb'
    class_names_file = 'class_names.txt'
35
    folds_list_file = 'fold_indices.txt'
Elad Hoffer's avatar
Elad Hoffer committed
36
37
38
39
40
41
42
43
44
45
    train_list = [
        ['train_X.bin', '918c2871b30a85fa023e0c44e0bee87f'],
        ['train_y.bin', '5a34089d4802c674881badbb80307741'],
        ['unlabeled_X.bin', '5242ba1fed5e4be9e1e742405eb56ca4']
    ]

    test_list = [
        ['test_X.bin', '7f263ba9f9e0b06b93213547f721ac82'],
        ['test_y.bin', '36f9794fa4beb8a2c72628de14fa638e']
    ]
46
    splits = ('train', 'train+unlabeled', 'unlabeled', 'test')
Elad Hoffer's avatar
Elad Hoffer committed
47

48
49
50
51
52
53
54
55
56
    def __init__(
            self,
            root: str,
            split: str = "train",
            folds: Optional[int] = None,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None,
            download: bool = False,
    ) -> None:
57
58
        super(STL10, self).__init__(root, transform=transform,
                                    target_transform=target_transform)
59
        self.split = verify_str_arg(split, "split", self.splits)
60
        self.folds = self._verify_folds(folds)
Elad Hoffer's avatar
Elad Hoffer committed
61
62
63

        if download:
            self.download()
64
        elif not self._check_integrity():
Elad Hoffer's avatar
Elad Hoffer committed
65
            raise RuntimeError(
soumith's avatar
soumith committed
66
67
                'Dataset not found or corrupted. '
                'You can use download=True to download it')
Elad Hoffer's avatar
Elad Hoffer committed
68
69

        # now load the picked numpy arrays
Vasilis Vryniotis's avatar
Vasilis Vryniotis committed
70
        self.labels: Optional[np.ndarray]
Elad Hoffer's avatar
Elad Hoffer committed
71
72
73
        if self.split == 'train':
            self.data, self.labels = self.__loadfile(
                self.train_list[0][0], self.train_list[1][0])
74
75
            self.__load_folds(folds)

Elad Hoffer's avatar
Elad Hoffer committed
76
77
78
        elif self.split == 'train+unlabeled':
            self.data, self.labels = self.__loadfile(
                self.train_list[0][0], self.train_list[1][0])
79
            self.__load_folds(folds)
Elad Hoffer's avatar
Elad Hoffer committed
80
81
82
83
84
85
86
            unlabeled_data, _ = self.__loadfile(self.train_list[2][0])
            self.data = np.concatenate((self.data, unlabeled_data))
            self.labels = np.concatenate(
                (self.labels, np.asarray([-1] * unlabeled_data.shape[0])))

        elif self.split == 'unlabeled':
            self.data, _ = self.__loadfile(self.train_list[2][0])
87
            self.labels = np.asarray([-1] * self.data.shape[0])
Elad Hoffer's avatar
Elad Hoffer committed
88
89
90
91
92
        else:  # self.split == 'test':
            self.data, self.labels = self.__loadfile(
                self.test_list[0][0], self.test_list[1][0])

        class_file = os.path.join(
moskomule's avatar
moskomule committed
93
            self.root, self.base_folder, self.class_names_file)
Elad Hoffer's avatar
Elad Hoffer committed
94
95
96
97
        if os.path.isfile(class_file):
            with open(class_file) as f:
                self.classes = f.read().splitlines()

98
    def _verify_folds(self, folds: Optional[int]) -> Optional[int]:
99
100
101
102
103
104
105
106
107
108
109
110
        if folds is None:
            return folds
        elif isinstance(folds, int):
            if folds in range(10):
                return folds
            msg = ("Value for argument folds should be in the range [0, 10), "
                   "but got {}.")
            raise ValueError(msg.format(folds))
        else:
            msg = "Expected type None or int for argument folds, but got type {}."
            raise ValueError(msg.format(type(folds)))

111
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
112
113
114
115
116
117
118
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
119
        target: Optional[int]
Elad Hoffer's avatar
Elad Hoffer committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
        if self.labels is not None:
            img, target = self.data[index], int(self.labels[index])
        else:
            img, target = self.data[index], None

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(np.transpose(img, (1, 2, 0)))

        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

137
    def __len__(self) -> int:
Elad Hoffer's avatar
Elad Hoffer committed
138
139
        return self.data.shape[0]

140
    def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]:
Elad Hoffer's avatar
Elad Hoffer committed
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        labels = None
        if labels_file:
            path_to_labels = os.path.join(
                self.root, self.base_folder, labels_file)
            with open(path_to_labels, 'rb') as f:
                labels = np.fromfile(f, dtype=np.uint8) - 1  # 0-based

        path_to_data = os.path.join(self.root, self.base_folder, data_file)
        with open(path_to_data, 'rb') as f:
            # read whole file in uint8 chunks
            everything = np.fromfile(f, dtype=np.uint8)
            images = np.reshape(everything, (-1, 3, 96, 96))
            images = np.transpose(images, (0, 1, 3, 2))

        return images, labels
156

157
    def _check_integrity(self) -> bool:
Francisco Massa's avatar
Francisco Massa committed
158
159
160
161
162
163
164
165
        root = self.root
        for fentry in (self.train_list + self.test_list):
            filename, md5 = fentry[0], fentry[1]
            fpath = os.path.join(root, self.base_folder, filename)
            if not check_integrity(fpath, md5):
                return False
        return True

166
    def download(self) -> None:
Francisco Massa's avatar
Francisco Massa committed
167
168
169
        if self._check_integrity():
            print('Files already downloaded and verified')
            return
170
        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)
171
        self._check_integrity()
Francisco Massa's avatar
Francisco Massa committed
172

173
    def extra_repr(self) -> str:
174
        return "Split: {split}".format(**self.__dict__)
175

176
    def __load_folds(self, folds: Optional[int]) -> None:
177
        # loads one of the folds if specified
178
179
180
181
182
183
        if folds is None:
            return
        path_to_folds = os.path.join(
            self.root, self.base_folder, self.folds_list_file)
        with open(path_to_folds, 'r') as f:
            str_idx = f.read().splitlines()[folds]
Philip Meier's avatar
Philip Meier committed
184
            list_idx = np.fromstring(str_idx, dtype=np.int64, sep=' ')
Vasilis Vryniotis's avatar
Vasilis Vryniotis committed
185
186
187
            self.data = self.data[list_idx, :, :, :]
            if self.labels is not None:
                self.labels = self.labels[list_idx]