| 48 | |
| 49 | class PlainImageLoader(torchvision.datasets.VisionDataset): |
| 50 | def __init__(self, root, corruptions, scan_for_broken_images:bool = False): |
| 51 | super(PlainImageLoader, self).__init__(root) |
| 52 | self.corruptions = corruptions |
| 53 | self.all_image_filenames = glob(os.path.join(root, "*.jpg")) |
| 54 | self.all_image_filenames.extend(glob(os.path.join(root, "**", "*.jpg"))) # Include subdirectories. |
| 55 | self.all_image_filenames.extend(glob(os.path.join(root, "*.png"))) |
| 56 | self.all_image_filenames.extend(glob(os.path.join(root, "**", "*.png"))) |
| 57 | self.all_image_filenames.extend(glob(os.path.join(root, "*.gif"))) |
| 58 | self.all_image_filenames.extend(glob(os.path.join(root, "**", "*.gif"))) |
| 59 | if scan_for_broken_images: |
| 60 | # Filter images that can't get loaded. This is a little slow but saves some headache. |
| 61 | print("Scanning for broken images...") |
| 62 | to_remove = list() |
| 63 | for filename in tqdm(self.all_image_filenames): |
| 64 | try: |
| 65 | _ = Image.open(filename).convert("RGB") |
| 66 | except KeyboardInterrupt: |
| 67 | raise |
| 68 | except Exception as e: |
| 69 | print(f"Failed to read {filename}: {e}") |
| 70 | to_remove.append(filename) |
| 71 | for filename in to_remove: |
| 72 | self.all_image_filenames.remove(filename) |
| 73 | print(f"Training set has {len(self.all_image_filenames)} images.") |
| 74 | |
| 75 | def __getitem__(self, index: int) -> Any: |
| 76 | img_left = self.corruptions(Image.open(self.all_image_filenames[index]).convert("RGB")) |