| 33 | |
| 34 | |
| 35 | class CsvDataset(Dataset): |
| 36 | def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t", tokenizer=None): |
| 37 | logging.debug(f'Loading csv data from {input_filename}.') |
| 38 | df = pd.read_csv(input_filename, sep=sep) |
| 39 | |
| 40 | self.images = df[img_key].tolist() |
| 41 | self.captions = df[caption_key].tolist() |
| 42 | self.transforms = transforms |
| 43 | logging.debug('Done loading data.') |
| 44 | |
| 45 | self.tokenize = tokenizer |
| 46 | |
| 47 | def __len__(self): |
| 48 | return len(self.captions) |
| 49 | |
| 50 | def __getitem__(self, idx): |
| 51 | images = self.transforms(Image.open(str(self.images[idx]))) |
| 52 | texts = self.tokenize([str(self.captions[idx])])[0] |
| 53 | return images, texts |
| 54 | |
| 55 | |
| 56 | class SharedEpoch: |