| import json | |
| import os | |
| import random | |
| from torch.utils.data import Dataset | |
| from PIL import Image | |
| from PIL import ImageFile | |
| ImageFile.LOAD_TRUNCATED_IMAGES = True | |
| Image.MAX_IMAGE_PIXELS = None | |
| from data.utils import pre_caption | |
| import os,glob | |
| class pretrain_dataset(Dataset): | |
| def __init__(self, ann_file, laion_path, transform): | |
| self.ann_pretrain = [] | |
| for f in ann_file: | |
| print('loading '+f) | |
| ann = json.load(open(f,'r')) | |
| self.ann_pretrain += ann | |
| self.laion_path = laion_path | |
| if self.laion_path: | |
| self.laion_files = glob.glob(os.path.join(laion_path,'*.json')) | |
| print('loading '+self.laion_files[0]) | |
| with open(self.laion_files[0],'r') as f: | |
| self.ann_laion = json.load(f) | |
| self.annotation = self.ann_pretrain + self.ann_laion | |
| else: | |
| self.annotation = self.ann_pretrain | |
| self.transform = transform | |
| def reload_laion(self, epoch): | |
| n = epoch%len(self.laion_files) | |
| print('loading '+self.laion_files[n]) | |
| with open(self.laion_files[n],'r') as f: | |
| self.ann_laion = json.load(f) | |
| self.annotation = self.ann_pretrain + self.ann_laion | |
| def __len__(self): | |
| return len(self.annotation) | |
| def __getitem__(self, index): | |
| ann = self.annotation[index] | |
| image = Image.open(ann['image']).convert('RGB') | |
| image = self.transform(image) | |
| caption = pre_caption(ann['caption'],30) | |
| return image, caption |