-
Notifications
You must be signed in to change notification settings - Fork 15
/
Copy pathdata_provider.py
155 lines (115 loc) · 5.58 KB
/
data_provider.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
import torch
import torch.utils.data as data
import pickle
from bigfile import BigFile
from textlib import TextTool, Vocabulary
from txt2vec import get_lang, BowVec, BowVecNSW, W2Vec, W2VecNSW, IndexVec
def collate_vision(data):
vis_feats, idxs, vis_ids = zip(*data)
vis_feats = torch.stack(vis_feats, 0)
return vis_feats, idxs, vis_ids
def collate_text(data):
data.sort(key=lambda x: len(TextTool.tokenize(x[0])), reverse=True)
captions, idxs, cap_ids = zip(*data)
return captions, idxs, cap_ids
def collate_pair(data):
data.sort(key=lambda x: len(TextTool.tokenize(x[1])), reverse=True)
vis_feats, captions, idxs, vis_ids, cap_ids = zip(*data)
vis_feats = torch.stack(vis_feats, 0)
return vis_feats, captions, idxs, vis_ids, cap_ids
class VisionDataset(data.Dataset):
def __init__(self, params):
self.vis_feat_file = BigFile(params['vis_feat']) if isinstance(params['vis_feat'], str) else params['vis_feat']
self.vis_ids = self.vis_feat_file.names if params.get('vis_ids', None) is None else params['vis_ids']
self.length = len(self.vis_ids)
def __getitem__(self, index):
vis_id = self.vis_ids[index]
vis_tensor = self.get_feat_by_id(vis_id)
return vis_tensor, index, vis_id
def get_feat_by_id(self, vis_id):
vis_tensor = torch.Tensor(self.vis_feat_file.read_one(vis_id))
return vis_tensor
def __len__(self):
return self.length
class TextDataset(data.Dataset):
def __init__(self, params):
capfile = params['capfile']
self.captions = {}
self.cap_ids = []
with open(capfile, 'r') as reader:
for line in reader.readlines():
cap_id, caption = line.strip().split(' ', 1)
self.captions[cap_id] = caption
self.cap_ids.append(cap_id)
self.length = len(self.cap_ids)
def __getitem__(self, index):
cap_id = self.cap_ids[index]
caption = self.get_caption_by_id(cap_id)
return caption, index, cap_id
def get_caption_by_id(self, cap_id):
caption = self.captions[cap_id]
return caption
def __len__(self):
return self.length
class PairDataset(data.Dataset):
def __init__(self, params):
self.visData = VisionDataset(params)
self.txtData = TextDataset(params)
self.cap_ids = self.txtData.cap_ids
self.length = len(self.cap_ids)
def __getitem__(self, index):
cap_id = self.cap_ids[index]
vis_id = self.get_visId_by_capId(cap_id)
caption = self.txtData.get_caption_by_id(cap_id)
vis_feat = self.visData.get_feat_by_id(vis_id)
return vis_feat, caption, index, vis_id, cap_id
def get_visId_by_capId(self, cap_id):
vis_id = cap_id.split('#', 1)[0]
return vis_id
def __len__(self):
return self.length
def vis_provider(params):
data_loader = torch.utils.data.DataLoader(dataset=VisionDataset(params),
batch_size=params.get('batch_size', 1),
shuffle=params.get('shuffle', False),
pin_memory=params.get('pin_memory', False),
num_workers=params.get('num_workers', 0),
collate_fn=collate_vision)
return data_loader
def txt_provider(params):
data_loader = torch.utils.data.DataLoader(dataset=TextDataset(params),
batch_size=params.get('batch_size', 1),
shuffle=params.get('shuffle', False),
pin_memory=params.get('pin_memory', False),
num_workers=params.get('num_workers', 0),
collate_fn=collate_text)
return data_loader
def pair_provider(params):
data_loader = torch.utils.data.DataLoader(dataset=PairDataset(params),
batch_size=params.get('batch_size', 1),
shuffle=params.get('shuffle', False),
pin_memory=params.get('pin_memory', False),
num_workers=params.get('num_workers', 0),
collate_fn=collate_pair)
return data_loader
if __name__ == '__main__':
import os
data_path = 'VisualSearch'
collection = 'tgif-msrvtt10k'
vid_feat = 'mean_resnext101_resnet152'
vid_feat_dir = os.path.join(data_path, collection, 'FeatureData', vid_feat)
vis_loader = vis_provider({'vis_feat': vid_feat_dir, 'batch_size':100, 'num_workers':2})
for i, (feat_vecs, idxs, vis_ids) in enumerate(vis_loader):
print i, feat_vecs.shape, len(idxs)
break
capfile = os.path.join(data_path, collection, 'TextData', '%s.caption.txt' % collection)
txt_loader = txt_provider({'capfile':capfile, 'batch_size':100, 'num_workers':2})
for i, (captions, idxs, cap_ids) in enumerate(txt_loader):
print i, captions, len(cap_ids)
print [len(cap) for cap in captions]
break
pair_loader = pair_provider({'vis_feat': vid_feat_dir, 'capfile': capfile, 'batch_size':100, 'num_workers':2, 'shuffle':True})
for i, (vis_feats, captions, idxs, vis_ids, cap_ids) in enumerate(pair_loader):
print i, vis_feats.shape, captions[:10], len(cap_ids)
print [len(cap) for cap in captions]
break