|
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 |
