Args: train_or_test (str): either 'train' or 'test' shuffle (bool): shuffle the dataset
(self, train_or_test, shuffle=True, dir=None)
| 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] |
nothing calls this directly
no test coverage detected