| import os | |
| import json | |
| import torch | |
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| class BaseDataset(Dataset): | |
| def __init__(self, args, tokenizer, split, transform=None): | |
| self.image_dir = args.image_dir | |
| self.ann_path = args.ann_path | |
| self.max_seq_length = args.max_seq_length | |
| self.split = split | |
| self.tokenizer = tokenizer | |
| self.transform = transform | |
| self.ann = json.loads(open(self.ann_path, 'r', encoding="utf_8_sig").read()) | |
| self.examples = self.ann[self.split] | |
| for i in range(len(self.examples)): | |
| self.examples[i]['ids'] = tokenizer(self.examples[i]['finding'])[:self.max_seq_length] | |
| self.examples[i]['mask'] = [1] * len(self.examples[i]['ids']) | |
| def __len__(self): | |
| return len(self.examples) | |
| class MyDataset(BaseDataset): | |
| def __getitem__(self, idx): | |
| example = self.examples[idx] | |
| image_id = example['uid'] | |
| image_path = example['image_path'] | |
| image_1 = Image.open(os.path.join(self.image_dir, image_path[0])).convert('RGB') | |
| image_2 = Image.open(os.path.join(self.image_dir, image_path[1])).convert('RGB') | |
| if self.transform is not None: | |
| image_1 = self.transform(image_1) | |
| image_2 = self.transform(image_2) | |
| image = torch.stack((image_1, image_2), 0) | |
| report_ids = example['ids'] | |
| report_masks = example['mask'] | |
| mesh_label = example['labels'] | |
| seq_length = len(report_ids) | |
| sample = (image_id, image, report_ids, report_masks, seq_length, mesh_label) | |
| return sample | |