stl10.py 7.06 KB
Newer Older
Elad Hoffer's avatar
Elad Hoffer committed
1
import os.path
2
from typing import Any, Callable, cast, Optional, Tuple
Elad Hoffer's avatar
Elad Hoffer committed
3

4
5
6
import numpy as np
from PIL import Image

7
from .utils import check_integrity, download_and_extract_archive, verify_str_arg
8
from .vision import VisionDataset
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
    """`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'}.
18
            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.
    """
30
31

    base_folder = "stl10_binary"
Elad Hoffer's avatar
Elad Hoffer committed
32
33
    url = "http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz"
    filename = "stl10_binary.tar.gz"
34
35
36
    tgz_md5 = "91f7769df0f17e558f3565bffb0c7dfb"
    class_names_file = "class_names.txt"
    folds_list_file = "fold_indices.txt"
Elad Hoffer's avatar
Elad Hoffer committed
37
    train_list = [
38
39
40
        ["train_X.bin", "918c2871b30a85fa023e0c44e0bee87f"],
        ["train_y.bin", "5a34089d4802c674881badbb80307741"],
        ["unlabeled_X.bin", "5242ba1fed5e4be9e1e742405eb56ca4"],
Elad Hoffer's avatar
Elad Hoffer committed
41
42
    ]

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

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

        if download:
            self.download()
61
        elif not self._check_integrity():
62
            raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it")
Elad Hoffer's avatar
Elad Hoffer committed
63
64

        # now load the picked numpy arrays
Vasilis Vryniotis's avatar
Vasilis Vryniotis committed
65
        self.labels: Optional[np.ndarray]
66
67
        if self.split == "train":
            self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
68
            self.labels = cast(np.ndarray, self.labels)
69
70
            self.__load_folds(folds)

71
72
        elif self.split == "train+unlabeled":
            self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
73
            self.labels = cast(np.ndarray, self.labels)
74
            self.__load_folds(folds)
Elad Hoffer's avatar
Elad Hoffer committed
75
76
            unlabeled_data, _ = self.__loadfile(self.train_list[2][0])
            self.data = np.concatenate((self.data, unlabeled_data))
77
            self.labels = np.concatenate((self.labels, np.asarray([-1] * unlabeled_data.shape[0])))
Elad Hoffer's avatar
Elad Hoffer committed
78

79
        elif self.split == "unlabeled":
Elad Hoffer's avatar
Elad Hoffer committed
80
            self.data, _ = self.__loadfile(self.train_list[2][0])
81
            self.labels = np.asarray([-1] * self.data.shape[0])
Elad Hoffer's avatar
Elad Hoffer committed
82
        else:  # self.split == 'test':
83
            self.data, self.labels = self.__loadfile(self.test_list[0][0], self.test_list[1][0])
Elad Hoffer's avatar
Elad Hoffer committed
84

85
        class_file = os.path.join(self.root, self.base_folder, self.class_names_file)
Elad Hoffer's avatar
Elad Hoffer committed
86
87
88
89
        if os.path.isfile(class_file):
            with open(class_file) as f:
                self.classes = f.read().splitlines()

90
    def _verify_folds(self, folds: Optional[int]) -> Optional[int]:
91
92
93
94
95
        if folds is None:
            return folds
        elif isinstance(folds, int):
            if folds in range(10):
                return folds
96
            msg = "Value for argument folds should be in the range [0, 10), but got {}."
97
98
99
100
101
            raise ValueError(msg.format(folds))
        else:
            msg = "Expected type None or int for argument folds, but got type {}."
            raise ValueError(msg.format(type(folds)))

102
    def __getitem__(self, index: int) -> Tuple[Any, Any]:
103
104
105
106
107
108
109
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
110
        target: Optional[int]
Elad Hoffer's avatar
Elad Hoffer committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
        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

128
    def __len__(self) -> int:
Elad Hoffer's avatar
Elad Hoffer committed
129
130
        return self.data.shape[0]

131
    def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]:
Elad Hoffer's avatar
Elad Hoffer committed
132
133
        labels = None
        if labels_file:
134
135
            path_to_labels = os.path.join(self.root, self.base_folder, labels_file)
            with open(path_to_labels, "rb") as f:
Elad Hoffer's avatar
Elad Hoffer committed
136
137
138
                labels = np.fromfile(f, dtype=np.uint8) - 1  # 0-based

        path_to_data = os.path.join(self.root, self.base_folder, data_file)
139
        with open(path_to_data, "rb") as f:
Elad Hoffer's avatar
Elad Hoffer committed
140
141
142
143
144
145
            # 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
146

147
    def _check_integrity(self) -> bool:
148
149
        for filename, md5 in self.train_list + self.test_list:
            fpath = os.path.join(self.root, self.base_folder, filename)
Francisco Massa's avatar
Francisco Massa committed
150
151
152
153
            if not check_integrity(fpath, md5):
                return False
        return True

154
    def download(self) -> None:
Francisco Massa's avatar
Francisco Massa committed
155
        if self._check_integrity():
156
            print("Files already downloaded and verified")
Francisco Massa's avatar
Francisco Massa committed
157
            return
158
        download_and_extract_archive(self.url, self.root, filename=self.filename, md5=self.tgz_md5)
159
        self._check_integrity()
Francisco Massa's avatar
Francisco Massa committed
160

161
    def extra_repr(self) -> str:
162
        return "Split: {split}".format(**self.__dict__)
163

164
    def __load_folds(self, folds: Optional[int]) -> None:
165
        # loads one of the folds if specified
166
167
        if folds is None:
            return
168
        path_to_folds = os.path.join(self.root, self.base_folder, self.folds_list_file)
169
        with open(path_to_folds) as f:
170
            str_idx = f.read().splitlines()[folds]
171
            list_idx = np.fromstring(str_idx, dtype=np.int64, sep=" ")
Vasilis Vryniotis's avatar
Vasilis Vryniotis committed
172
173
174
            self.data = self.data[list_idx, :, :, :]
            if self.labels is not None:
                self.labels = self.labels[list_idx]