Spaces:
Running
Running
| import os | |
| import json | |
| import random | |
| from PIL import Image | |
| import torch | |
| from torch.utils.data import Dataset | |
| from data.utils import pre_question | |
| from torchvision.datasets.utils import download_url | |
| class vqa_dataset(Dataset): | |
| def __init__(self, transform, ann_root, vqa_root, vg_root, train_files=[], split="train"): | |
| self.split = split | |
| self.transform = transform | |
| self.vqa_root = vqa_root | |
| self.vg_root = vg_root | |
| if split=='train': | |
| urls = {'vqa_train':'https://storage.googleapis.com/sfr-vision-language-research/datasets/vqa_train.json', | |
| 'vqa_val':'https://storage.googleapis.com/sfr-vision-language-research/datasets/vqa_val.json', | |
| 'vg_qa':'https://storage.googleapis.com/sfr-vision-language-research/datasets/vg_qa.json'} | |
| self.annotation = [] | |
| for f in train_files: | |
| download_url(urls[f],ann_root) | |
| self.annotation += json.load(open(os.path.join(ann_root,'%s.json'%f),'r')) | |
| else: | |
| download_url('https://storage.googleapis.com/sfr-vision-language-research/datasets/vqa_test.json',ann_root) | |
| self.annotation = json.load(open(os.path.join(ann_root,'vqa_test.json'),'r')) | |
| download_url('https://storage.googleapis.com/sfr-vision-language-research/datasets/answer_list.json',ann_root) | |
| self.answer_list = json.load(open(os.path.join(ann_root,'answer_list.json'),'r')) | |
| def __len__(self): | |
| return len(self.annotation) | |
| def __getitem__(self, index): | |
| ann = self.annotation[index] | |
| if ann['dataset']=='vqa': | |
| image_path = os.path.join(self.vqa_root,ann['image']) | |
| elif ann['dataset']=='vg': | |
| image_path = os.path.join(self.vg_root,ann['image']) | |
| image = Image.open(image_path).convert('RGB') | |
| image = self.transform(image) | |
| if self.split == 'test': | |
| question = pre_question(ann['question']) | |
| question_id = ann['question_id'] | |
| return image, question, question_id | |
| elif self.split=='train': | |
| question = pre_question(ann['question']) | |
| if ann['dataset']=='vqa': | |
| answer_weight = {} | |
| for answer in ann['answer']: | |
| if answer in answer_weight.keys(): | |
| answer_weight[answer] += 1/len(ann['answer']) | |
| else: | |
| answer_weight[answer] = 1/len(ann['answer']) | |
| answers = list(answer_weight.keys()) | |
| weights = list(answer_weight.values()) | |
| elif ann['dataset']=='vg': | |
| answers = [ann['answer']] | |
| weights = [0.2] | |
| return image, question, answers, weights | |
| def vqa_collate_fn(batch): | |
| image_list, question_list, answer_list, weight_list, n = [], [], [], [], [] | |
| for image, question, answer, weights in batch: | |
| image_list.append(image) | |
| question_list.append(question) | |
| weight_list += weights | |
| answer_list += answer | |
| n.append(len(answer)) | |
| return torch.stack(image_list,dim=0), question_list, answer_list, torch.Tensor(weight_list), n |