| | import os |
| | import json |
| | import torch |
| | from torch.utils.data import Dataset |
| | from torchvision.datasets.utils import download_url |
| | from PIL import Image, ImageFile |
| | ImageFile.LOAD_TRUNCATED_IMAGES = True |
| | from glob import glob |
| | from data.utils import pre_caption |
| |
|
| |
|
| | class facecaption_train(Dataset): |
| | def __init__(self, transform, image_root, ann_root, max_words=65, prompt=''): |
| | ''' |
| | image_root (string): Root directory of images (e.g. coco/images/) |
| | ann_root (string): directory to store the annotation file |
| | ''' |
| | all_json = sorted(glob(os.path.join(ann_root, '*.json'))) |
| | self.annotation = [] |
| | for json_path in all_json[0:1]: |
| | |
| | with open(json_path, 'r') as json_file: |
| | data = json.load(json_file) |
| | self.annotation.extend(data) |
| |
|
| | self.transform = transform |
| | self.image_root = image_root |
| | self.max_words = max_words |
| | self.prompt = prompt |
| | |
| | self.img_ids = {} |
| | n = 0 |
| | for ann in self.annotation: |
| | img_id = ann['image_id'] |
| | if img_id not in self.img_ids.keys(): |
| | self.img_ids[img_id] = n |
| | n += 1 |
| | |
| | def __len__(self): |
| | return len(self.annotation) |
| | |
| | def __getitem__(self, index): |
| | |
| | ann = self.annotation[index] |
| | image_path = os.path.join(self.image_root, ann['image']) |
| | image = Image.open(image_path).convert('RGB') |
| | image = self.transform(image) |
| | |
| | |
| | caption = self.prompt + pre_caption(*ann['caption'], self.max_words) |
| | image_id = self.img_ids[ann['image_id']] |
| | return image, caption, image_id |
| | |
| |
|
| | class facecaption_test(Dataset): |
| | def __init__(self, transform, image_root, ann_root, max_words=65): |
| | ''' |
| | image_root (string): Root directory of images (e.g. coco/images/) |
| | ann_root (string): directory to store the annotation file |
| | ''' |
| | all_json = sorted(glob(os.path.join(ann_root, '*.json'))) |
| | self.annotation = [] |
| | for json_path in all_json[-1:]: |
| | with open(json_path, 'r') as json_file: |
| | data = json.load(json_file) |
| | self.annotation.extend(data) |
| | self.annotation = self.annotation[:5000] |
| |
|
| | self.transform = transform |
| | self.image_root = image_root |
| | |
| | self.text = [] |
| | self.image = [] |
| | self.txt2img = {} |
| | self.img2txt = {} |
| |
|
| | txt_id = 0 |
| | for img_id, ann in enumerate(self.annotation): |
| | self.image.append(ann['image']) |
| | self.img2txt[img_id] = [] |
| | for i, caption in enumerate(ann['caption']): |
| | self.text.append(pre_caption(caption, max_words)) |
| | self.img2txt[img_id].append(txt_id) |
| | self.txt2img[txt_id] = img_id |
| | txt_id += 1 |
| | |
| | def __len__(self): |
| | return len(self.annotation) |
| | |
| | def __getitem__(self, index): |
| | |
| | ann = self.annotation[index] |
| | |
| | image_path = os.path.join(self.image_root, ann['image']) |
| | image = Image.open(image_path).convert('RGB') |
| | image = self.transform(image) |
| | return image, index |
