/
data_loader_ks.py
179 lines (153 loc) · 6.89 KB
/
data_loader_ks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
"""
Author: Omkar Damle
Date: 28th Oct 2018
Dataloader class
Points to note:
1. The annotations are returned as an array of indices according to the vocabulary dictionary
2. Each annotation starts with the start symbol and ends with the end symbol. There can be padding after the end symbol in order
to make a batch
3. Each image feature has the shape - (36,2048)
reference:
https://github.com/hengyuan-hu/bottom-up-attention-vqa/blob/master/dataset.py
https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/03-advanced/image_captioning/data_loader.py
"""
import torch
import torchvision.transforms as transforms
import torch.utils.data as data
import os
import pickle
import numpy as np
import nltk
from PIL import Image
from data_helpers.vocab import Vocabulary
from pycocotools.coco import COCO
import h5py
import argparse
import pickle
class CocoDataset(data.Dataset):
"""COCO Custom Dataset compatible with torch.utils.data.DataLoader."""
def __init__(self, data_root, vocab, data_type, debug=False):
"""Set the path for images, captions and vocabulary wrapper.
Args:
data_root: root dir.
vocab: vocabulary wrapper.
transform: image transformer.
"""
self.data_type = data_type
self.imgid2idx = pickle.load(open(os.path.join(data_root, "karpathy_splits", data_type + "_ks_imgid2idx.pkl"), 'rb'))
print("Len of pkl file: " + str(len(self.imgid2idx)))
ids = None
self.debug = debug
if debug:
file_name = os.path.join(data_root, "karpathy_splits" ,data_type + "36_ks_mini.hdf5")
with open(os.path.join(data_root, "karpathy_splits", data_type + "_karpathy_mini.pkl"), 'rb') as f:
self.coco = pickle.load(f)
ids = list(self.coco.keys())
else:
if data_type == 'not_used':
file_name = os.path.join(data_root, "karpathy_splits" ,data_type + "36_ks_med.hdf5")
with open(os.path.join(data_root, "karpathy_splits", data_type + "_karpathy_medium.pkl"), 'rb') as f:
self.coco = pickle.load(f)
else:
file_name = os.path.join(data_root,"karpathy_splits",data_type + "36_ks.hdf5")
with open(os.path.join(data_root, "karpathy_splits", data_type + "_karpathy.pkl"), 'rb') as f:
self.coco = pickle.load(f)
#ks_json_path = os.path.join(data_root,"karpathy_splits","cocotalk.json")
#ks_h5_path = os.path.join(data_root,"karpathy_splits","cocotalk.h5")
#self.coco = COCO(annotation_path)
ids = list(self.coco.keys())
data_h5 = h5py.File(file_name,'r')
self.train_features = np.array(data_h5.get('image_features'))
self.data_root = data_root
self.ids = list(ids)
self.vocab = vocab
print('Initialization done for: ' + data_type)
print('Number of annotations: ' + str(len(self.ids)))
def __getitem__(self, index):
"""Returns one data pair (image and caption)."""
#print('Inside getitem, Retrieving for index: ' + str(index))
coco = None
vocab = self.vocab
ann_id = self.ids[index]
caption = None
img_id = None
coco = self.coco
caption = coco[ann_id]['caption']
img_id = coco[ann_id]['image_id']
#path = coco.loadImgs(img_id)[0]['file_name']
index = self.imgid2idx[img_id]
features = self.train_features[index]
features = torch.Tensor(features)
# Convert caption (string) to word ids.
#print('converting captions to word ids')
tokens = nltk.tokenize.word_tokenize(str(caption).lower())
caption = []
vocab_len = len(vocab)
caption.append(vocab('<start>'))
caption.extend([vocab(token) for token in tokens])
caption.append(vocab('<end>'))
#print(len(caption))
target = torch.Tensor(caption)
return img_id, features, target
def __len__(self):
return len(self.ids)
def indexto1hot(vocab_len, index):
#print("index type: ")
if isinstance(index,int) == False:
n = len(index)
#print("making a 1hot encoding of shape: " + str(n) + "," + str(vocab_len) )
one_hot = np.zeros([n,vocab_len])
#can this be optimized?
for i in range(n):
one_hot[i,index[i]]=1
return one_hot
else:
one_hot = np.zeros([vocab_len])
one_hot[index] = 1
return one_hot
def collate_fn(data):
"""Creates mini-batch tensors from the list of tuples (image, caption).
We should build custom collate_fn rather than using default collate_fn,
because merging caption (including padding) is not supported in default.
Args:
data: list of tuple (image, caption).
- feature: torch tensor of shape (36,2048).
- caption: torch tensor of shape (?); variable length.
Returns:
features: torch tensor of shape (batch_size, 36, 2048).
targets: torch tensor of shape (batch_size, padded_length).
lengths: list; valid length for each padded caption.
"""
#print("Length of list: " + str(len(data)))
# Sort a data list by caption length (descending order).
data.sort(key=lambda x: len(x[2]), reverse=True)
image_ids, features, captions = zip(*data)
# Merge images (from tuple of 3D tensor to 4D tensor).
features = torch.stack(features, 0)
# Merge captions (from tuple of 1D tensor to 2D tensor).
lengths = [len(cap) for cap in captions]
#vocab_len = len(captions[0][0])
#print("Vocab len: " + str(vocab_len) + "\n")
#print("Type: " + type(captions[0]))
#print("\n")
targets = torch.zeros(len(captions), max(lengths)).long()
for i, cap in enumerate(captions):
end = lengths[i]
targets[i, :end] = cap[:end]
return image_ids, features, targets, lengths
def get_loader(data_root, vocab, batch_size, data_type, shuffle, num_workers, debug=False):
"""Returns torch.utils.data.DataLoader for custom coco dataset."""
# COCO caption dataset
coco = CocoDataset(data_root=data_root,
vocab=vocab, data_type = data_type, debug=debug)
# Data loader for COCO dataset
# This will return (features, captions, lengths) for each iteration.
# features: a tensor of shape (batch_size, 36, 2048).
# captions: a tensor of shape (batch_size, padded_length).
# lengths: a list indicating valid length for each caption. length is (batch_size).
data_loader = torch.utils.data.DataLoader(dataset=coco,
batch_size=batch_size,
shuffle=shuffle,
num_workers=num_workers,
collate_fn = collate_fn)
return data_loader