| 39 | |
| 40 | |
| 41 | class CLIPImageDataset(torch.utils.data.Dataset): |
| 42 | def __init__(self, data): |
| 43 | self.data = data |
| 44 | # only 224x224 ViT-B/32 supported for now |
| 45 | self.preprocess = self._transform_test(224) |
| 46 | |
| 47 | def _transform_test(self, n_px): |
| 48 | return Compose([ |
| 49 | Resize(n_px, interpolation=Image.BICUBIC), |
| 50 | CenterCrop(n_px), |
| 51 | Convert, |
| 52 | ToTensor(), |
| 53 | Normalize((0.48145466, 0.4578275, 0.40821073), |
| 54 | (0.26862954, 0.26130258, 0.27577711)), |
| 55 | ]) |
| 56 | |
| 57 | def __getitem__(self, idx): |
| 58 | c_data = self.data[idx] |
| 59 | image = Image.open(c_data) |
| 60 | image = self.preprocess(image) |
| 61 | return {'image': image} |
| 62 | |
| 63 | def __len__(self): |
| 64 | return len(self.data) |
| 65 | |
| 66 | |
| 67 | class DINOImageDataset(torch.utils.data.Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected