CIFAR10 dataset for visualization. Splits: train_image_data, train_label_data, test_image_data, test_label_data, image_data, label_data
| 794 | |
| 795 | |
| 796 | class CIFAR10(Dataset): |
| 797 | """ |
| 798 | CIFAR10 dataset for visualization. |
| 799 | |
| 800 | Splits: |
| 801 | train_image_data, train_label_data, test_image_data, test_label_data, image_data, label_data |
| 802 | """ |
| 803 | def __init__(self): |
| 804 | super(CIFAR10, self).__init__( |
| 805 | "cifar10", |
| 806 | urls={ |
| 807 | "train_image_data": "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz", |
| 808 | "train_label_data": "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz", |
| 809 | "test_image_data": "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz", |
| 810 | "test_label_data": "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz", |
| 811 | "image_data": [], # depends on `train_image_data` & `test_image_data` |
| 812 | "label_data": [] # depends on `train_label_data` & `test_label_data` |
| 813 | }, |
| 814 | ) |
| 815 | |
| 816 | def load_images(self, *batch_files): |
| 817 | images = [] |
| 818 | for batch_file in batch_files: |
| 819 | batch = np.fromfile(batch_file, dtype=np.uint8) |
| 820 | batch = batch.reshape(-1, 32*32*3 + 1) |
| 821 | images.append(batch[:, 1:]) |
| 822 | return np.concatenate(images) |
| 823 | |
| 824 | def load_labels(self, meta_file, *batch_files): |
| 825 | classes = [] |
| 826 | with open(meta_file, "r") as fin: |
| 827 | for line in fin: |
| 828 | line = line.strip() |
| 829 | if line: |
| 830 | classes.append(line) |
| 831 | classes = np.asarray(classes) |
| 832 | labels = [] |
| 833 | for batch_file in batch_files: |
| 834 | batch = np.fromfile(batch_file, dtype=np.uint8) |
| 835 | batch = batch.reshape(-1, 32*32*3 + 1) |
| 836 | labels.append(batch[:, 0]) |
| 837 | return classes[np.concatenate(labels)] |
| 838 | |
| 839 | def train_image_data_preprocess(self, raw_path, save_file): |
| 840 | batch_files = glob.glob(os.path.join(raw_path, "cifar-10-batches-bin/data_batch_*.bin")) |
| 841 | return self.load_images(*batch_files) |
| 842 | |
| 843 | def train_label_data_preprocess(self, raw_path, save_file): |
| 844 | meta_file = os.path.join(raw_path, "cifar-10-batches-bin/batches.meta.txt") |
| 845 | batch_files = glob.glob(os.path.join(raw_path, "cifar-10-batches-bin/data_batch_*.bin")) |
| 846 | return self.load_labels(meta_file, *batch_files) |
| 847 | |
| 848 | def test_image_data_preprocess(self, raw_path, save_file): |
| 849 | batch_file = os.path.join(raw_path, "cifar-10-batches-bin/test_batch.bin") |
| 850 | return self.load_images(batch_file) |
| 851 | |
| 852 | def test_label_data_preprocess(self, raw_path, save_file): |
| 853 | meta_file = os.path.join(raw_path, "cifar-10-batches-bin/batches.meta.txt") |