| 47 | |
| 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")) |
| 77 | if random.choice([True, False]): |
| 78 | other_index = random.randint(0, len(self.all_image_filenames)-1) |
| 79 | img_right = self.corruptions(Image.open(self.all_image_filenames[other_index]).convert("RGB")) |
| 80 | label = -1.0 |
| 81 | if other_index == index: |
| 82 | label = 1.0 # In the slim chance we happen to pick exactly the same index at random... |
| 83 | else: |
| 84 | img_right = self.corruptions(Image.open(self.all_image_filenames[index]).convert("RGB")) |
| 85 | label = 1.0 |
| 86 | label = torch.tensor(label) |
| 87 | return img_left, img_right, label |
| 88 | |
| 89 | def __len__(self) -> int: |
| 90 | return len(self.all_image_filenames) |
| 91 | |
| 92 | |
| 93 | # Set up |