| 24 | |
| 25 | |
| 26 | class ImageNetV2Dataset(Dataset): |
| 27 | def __init__(self, variant='matched-frequency', transform=None, location='.'): |
| 28 | self.dataset_root = pathlib.Path(f'{location}/ImageNetV2-{variant}/') |
| 29 | self.tar_root = pathlib.Path(f'{location}/ImageNetV2-{variant}.tar.gz') |
| 30 | self.fnames = list(self.dataset_root.glob('**/*.jpeg')) |
| 31 | self.transform = transform |
| 32 | assert variant in URLS, f'unknown V2 Variant: {variant}' |
| 33 | if not self.dataset_root.exists() or len(self.fnames) != V2_DATASET_SIZE: |
| 34 | if not self.tar_root.exists(): |
| 35 | print(f'Dataset {variant} not found on disk, downloading....') |
| 36 | response = requests.get(URLS[variant], stream=True) |
| 37 | total_size_in_bytes = int(response.headers.get('content-length', 0)) |
| 38 | block_size = 1024 # 1 Kibibyte |
| 39 | progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True) |
| 40 | with open(self.tar_root, 'wb') as f: |
| 41 | for data in response.iter_content(block_size): |
| 42 | progress_bar.update(len(data)) |
| 43 | f.write(data) |
| 44 | progress_bar.close() |
| 45 | if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes: |
| 46 | assert False, f'Downloading from {URLS[variant]} failed' |
| 47 | print('Extracting....') |
| 48 | tarfile.open(self.tar_root).extractall(f'{location}') |
| 49 | shutil.move(f'{location}/{FNAMES[variant]}', self.dataset_root) |
| 50 | self.fnames = list(self.dataset_root.glob('**/*.jpeg')) |
| 51 | |
| 52 | def __len__(self): |
| 53 | return len(self.fnames) |
| 54 | |
| 55 | def __getitem__(self, i): |
| 56 | img, label = Image.open(self.fnames[i]), int(self.fnames[i].parent.name) |
| 57 | if self.transform is not None: |
| 58 | img = self.transform(img) |
| 59 | return img, label |