| 66 | |
| 67 | |
| 68 | class ImageDataset(Dataset): |
| 69 | def __init__(self, resolution, image_paths, classes=None, shard=0, num_shards=1): |
| 70 | super().__init__() |
| 71 | self.resolution = resolution |
| 72 | self.local_images = image_paths[shard:][::num_shards] |
| 73 | self.local_classes = None if classes is None else classes[shard:][::num_shards] |
| 74 | |
| 75 | def __len__(self): |
| 76 | return len(self.local_images) |
| 77 | |
| 78 | def __getitem__(self, idx): |
| 79 | path = self.local_images[idx] |
| 80 | with bf.BlobFile(path, "rb") as f: |
| 81 | pil_image = Image.open(f) |
| 82 | pil_image.load() |
| 83 | |
| 84 | # We are not on a new enough PIL to support the `reducing_gap` |
| 85 | # argument, which uses BOX downsampling at powers of two first. |
| 86 | # Thus, we do it by hand to improve downsample quality. |
| 87 | while min(*pil_image.size) >= 2 * self.resolution: |
| 88 | pil_image = pil_image.resize( |
| 89 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX |
| 90 | ) |
| 91 | |
| 92 | scale = self.resolution / min(*pil_image.size) |
| 93 | pil_image = pil_image.resize( |
| 94 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC |
| 95 | ) |
| 96 | |
| 97 | arr = np.array(pil_image.convert("RGB")) |
| 98 | crop_y = (arr.shape[0] - self.resolution) // 2 |
| 99 | crop_x = (arr.shape[1] - self.resolution) // 2 |
| 100 | arr = arr[crop_y : crop_y + self.resolution, crop_x : crop_x + self.resolution] |
| 101 | arr = arr.astype(np.float32) / 127.5 - 1 |
| 102 | |
| 103 | out_dict = {} |
| 104 | if self.local_classes is not None: |
| 105 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) |
| 106 | return np.transpose(arr, [2, 0, 1]), out_dict |