Folder datasets which returns the index of the image (for memory_bank)
| 6 | |
| 7 | |
| 8 | class ImageFolderInstance(datasets.ImageFolder): |
| 9 | """Folder datasets which returns the index of the image (for memory_bank) |
| 10 | """ |
| 11 | def __init__(self, root, transform=None, target_transform=None, |
| 12 | two_crop=False, jigsaw_transform=None): |
| 13 | super(ImageFolderInstance, self).__init__(root, transform, target_transform) |
| 14 | self.two_crop = two_crop |
| 15 | self.jigsaw_transform = jigsaw_transform |
| 16 | self.use_jigsaw = (jigsaw_transform is not None) |
| 17 | self.num = self.__len__() |
| 18 | |
| 19 | def __getitem__(self, index): |
| 20 | """ |
| 21 | Args: |
| 22 | index (int): index |
| 23 | Returns: |
| 24 | tuple: (image, index, ...) |
| 25 | """ |
| 26 | path, target = self.imgs[index] |
| 27 | image = self.loader(path) |
| 28 | |
| 29 | # # image |
| 30 | if self.transform is not None: |
| 31 | img = self.transform(image) |
| 32 | if self.two_crop: |
| 33 | img2 = self.transform(image) |
| 34 | img = torch.cat([img, img2], dim=0) |
| 35 | else: |
| 36 | img = image |
| 37 | |
| 38 | # # jigsaw |
| 39 | if self.use_jigsaw: |
| 40 | jigsaw_image = self.jigsaw_transform(image) |
| 41 | |
| 42 | if self.use_jigsaw: |
| 43 | return img, index, jigsaw_image |
| 44 | else: |
| 45 | return img, index |
no outgoing calls
no test coverage detected