| 80 | |
| 81 | |
| 82 | class CLIPImageDataset(torch.utils.data.Dataset): |
| 83 | def __init__(self, data): |
| 84 | self.data = data |
| 85 | # only 224x224 ViT-B/32 supported for now |
| 86 | self.preprocess = self._transform_test(224) |
| 87 | |
| 88 | def _transform_test(self, n_px): |
| 89 | return Compose([ |
| 90 | Resize(n_px, interpolation=Image.BICUBIC), |
| 91 | CenterCrop(n_px), |
| 92 | lambda image: image.convert("RGB"), |
| 93 | ToTensor(), |
| 94 | Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), |
| 95 | ]) |
| 96 | |
| 97 | def __getitem__(self, idx): |
| 98 | c_data = self.data[idx] |
| 99 | image = Image.open(c_data) |
| 100 | image = self.preprocess(image) |
| 101 | return {'image':image} |
| 102 | |
| 103 | def __len__(self): |
| 104 | return len(self.data) |
| 105 | |
| 106 | |
| 107 | def extract_all_captions(captions, model, device, batch_size=256, num_workers=8): |
no outgoing calls
no test coverage detected