doevent commited on
Commit
eba757c
1 Parent(s): b4bfe5e

Upload data/vqa_dataset.py

Browse files
Files changed (1) hide show
  1. data/vqa_dataset.py +88 -0
data/vqa_dataset.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ import random
4
+ from PIL import Image
5
+
6
+ import torch
7
+ from torch.utils.data import Dataset
8
+ from data.utils import pre_question
9
+
10
+ from torchvision.datasets.utils import download_url
11
+
12
+ class vqa_dataset(Dataset):
13
+ def __init__(self, transform, ann_root, vqa_root, vg_root, train_files=[], split="train"):
14
+ self.split = split
15
+
16
+ self.transform = transform
17
+ self.vqa_root = vqa_root
18
+ self.vg_root = vg_root
19
+
20
+ if split=='train':
21
+ urls = {'vqa_train':'https://storage.googleapis.com/sfr-vision-language-research/datasets/vqa_train.json',
22
+ 'vqa_val':'https://storage.googleapis.com/sfr-vision-language-research/datasets/vqa_val.json',
23
+ 'vg_qa':'https://storage.googleapis.com/sfr-vision-language-research/datasets/vg_qa.json'}
24
+
25
+ self.annotation = []
26
+ for f in train_files:
27
+ download_url(urls[f],ann_root)
28
+ self.annotation += json.load(open(os.path.join(ann_root,'%s.json'%f),'r'))
29
+ else:
30
+ download_url('https://storage.googleapis.com/sfr-vision-language-research/datasets/vqa_test.json',ann_root)
31
+ self.annotation = json.load(open(os.path.join(ann_root,'vqa_test.json'),'r'))
32
+
33
+ download_url('https://storage.googleapis.com/sfr-vision-language-research/datasets/answer_list.json',ann_root)
34
+ self.answer_list = json.load(open(os.path.join(ann_root,'answer_list.json'),'r'))
35
+
36
+
37
+ def __len__(self):
38
+ return len(self.annotation)
39
+
40
+ def __getitem__(self, index):
41
+
42
+ ann = self.annotation[index]
43
+
44
+ if ann['dataset']=='vqa':
45
+ image_path = os.path.join(self.vqa_root,ann['image'])
46
+ elif ann['dataset']=='vg':
47
+ image_path = os.path.join(self.vg_root,ann['image'])
48
+
49
+ image = Image.open(image_path).convert('RGB')
50
+ image = self.transform(image)
51
+
52
+ if self.split == 'test':
53
+ question = pre_question(ann['question'])
54
+ question_id = ann['question_id']
55
+ return image, question, question_id
56
+
57
+
58
+ elif self.split=='train':
59
+
60
+ question = pre_question(ann['question'])
61
+
62
+ if ann['dataset']=='vqa':
63
+ answer_weight = {}
64
+ for answer in ann['answer']:
65
+ if answer in answer_weight.keys():
66
+ answer_weight[answer] += 1/len(ann['answer'])
67
+ else:
68
+ answer_weight[answer] = 1/len(ann['answer'])
69
+
70
+ answers = list(answer_weight.keys())
71
+ weights = list(answer_weight.values())
72
+
73
+ elif ann['dataset']=='vg':
74
+ answers = [ann['answer']]
75
+ weights = [0.2]
76
+
77
+ return image, question, answers, weights
78
+
79
+
80
+ def vqa_collate_fn(batch):
81
+ image_list, question_list, answer_list, weight_list, n = [], [], [], [], []
82
+ for image, question, answer, weights in batch:
83
+ image_list.append(image)
84
+ question_list.append(question)
85
+ weight_list += weights
86
+ answer_list += answer
87
+ n.append(len(answer))
88
+ return torch.stack(image_list,dim=0), question_list, answer_list, torch.Tensor(weight_list), n