Produces [image, label] in MNIST dataset, image is 28x28 in the range [0,1], label is an int.
| 61 | |
| 62 | |
| 63 | class Mnist(RNGDataFlow): |
| 64 | """ |
| 65 | Produces [image, label] in MNIST dataset, |
| 66 | image is 28x28 in the range [0,1], label is an int. |
| 67 | """ |
| 68 | |
| 69 | _DIR_NAME = 'mnist_data' |
| 70 | _SOURCE_URL = 'http://yann.lecun.com/exdb/mnist/' |
| 71 | |
| 72 | def __init__(self, train_or_test, shuffle=True, dir=None): |
| 73 | """ |
| 74 | Args: |
| 75 | train_or_test (str): either 'train' or 'test' |
| 76 | shuffle (bool): shuffle the dataset |
| 77 | """ |
| 78 | if dir is None: |
| 79 | dir = get_dataset_path(self._DIR_NAME) |
| 80 | assert train_or_test in ['train', 'test'] |
| 81 | self.train_or_test = train_or_test |
| 82 | self.shuffle = shuffle |
| 83 | |
| 84 | def get_images_and_labels(image_file, label_file): |
| 85 | f = maybe_download(self._SOURCE_URL + image_file, dir) |
| 86 | images = extract_images(f) |
| 87 | f = maybe_download(self._SOURCE_URL + label_file, dir) |
| 88 | labels = extract_labels(f) |
| 89 | assert images.shape[0] == labels.shape[0] |
| 90 | return images, labels |
| 91 | |
| 92 | if self.train_or_test == 'train': |
| 93 | self.images, self.labels = get_images_and_labels( |
| 94 | 'train-images-idx3-ubyte.gz', |
| 95 | 'train-labels-idx1-ubyte.gz') |
| 96 | else: |
| 97 | self.images, self.labels = get_images_and_labels( |
| 98 | 't10k-images-idx3-ubyte.gz', |
| 99 | 't10k-labels-idx1-ubyte.gz') |
| 100 | |
| 101 | def __len__(self): |
| 102 | return self.images.shape[0] |
| 103 | |
| 104 | def __iter__(self): |
| 105 | idxs = list(range(self.__len__())) |
| 106 | if self.shuffle: |
| 107 | self.rng.shuffle(idxs) |
| 108 | for k in idxs: |
| 109 | img = self.images[k].reshape((28, 28)) |
| 110 | label = self.labels[k] |
| 111 | yield [img, label] |
| 112 | |
| 113 | |
| 114 | class FashionMnist(Mnist): |
no outgoing calls
no test coverage detected
searching dependent graphs…