| 65 | |
| 66 | |
| 67 | class DINOImageDataset(torch.utils.data.Dataset): |
| 68 | def __init__(self, data): |
| 69 | self.data = data |
| 70 | # only 224x224 ViT-B/32 supported for now |
| 71 | self.preprocess = self._transform_test(224) |
| 72 | |
| 73 | def _transform_test(self, n_px): |
| 74 | return Compose([ |
| 75 | Resize(256, interpolation=Image.BICUBIC), |
| 76 | CenterCrop(n_px), |
| 77 | Convert, |
| 78 | ToTensor(), |
| 79 | Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), |
| 80 | ]) |
| 81 | |
| 82 | def __getitem__(self, idx): |
| 83 | c_data = self.data[idx] |
| 84 | image = Image.open(c_data) |
| 85 | image = self.preprocess(image) |
| 86 | return {'image': image} |
| 87 | |
| 88 | def __len__(self): |
| 89 | return len(self.data) |
| 90 | |
| 91 | |
| 92 | def extract_all_captions(captions, model, device, batch_size=256, num_workers=8, append=False): |
nothing calls this directly
no outgoing calls
no test coverage detected