| 64 | |
| 65 | |
| 66 | class CLIPCapDataset(torch.utils.data.Dataset): |
| 67 | def __init__(self, data, prefix='A photo depicts'): |
| 68 | self.data = data |
| 69 | self.prefix = prefix |
| 70 | if self.prefix[-1] != ' ': |
| 71 | self.prefix += ' ' |
| 72 | |
| 73 | def __getitem__(self, idx): |
| 74 | c_data = self.data[idx] |
| 75 | c_data = clip.tokenize(self.prefix + c_data, truncate=True).squeeze() |
| 76 | return {'caption': c_data} |
| 77 | |
| 78 | def __len__(self): |
| 79 | return len(self.data) |
| 80 | |
| 81 | |
| 82 | class CLIPImageDataset(torch.utils.data.Dataset): |
no outgoing calls
no test coverage detected