File size: 1,619 Bytes
3b2b066 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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
|