This repository has been archived by the owner on Nov 3, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
agents.py
363 lines (321 loc) · 13.6 KB
/
agents.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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from typing import Optional
from parlai.core.params import ParlaiParser
from parlai.core.opt import Opt
from parlai.core.teachers import FixedDialogTeacher
from parlai.core.image_featurizers import ImageLoader
from parlai.utils.io import PathManager
from .build_2014 import build as build_2014
from .build_2014 import buildImage as buildImage_2014
from .build_2017 import build as build_2017
from .build_2017 import buildImage as buildImage_2017
import os
import json
import random
"""
Agents for MSCOCO Image Captioning Task
There are two versions of the task - one comprising MSCOCO 2014 splits
(from the 2015 task competition), and one comprising MSCOCO 2017 splits
For the 2014 splits, we use the train, val, and test split of Karpathy et.
al, "Deep visual-semantic alignments for generating image descriptions"
(splits from here: https://cs.stanford.edu/people/karpathy/deepimagesent/).
This split has ~82k train images, 5k validation images, and 5k test images.
The val and test images are taken from the original validation set of ~40k.
For 2017, we use the splits from the official MSCOCO Image Captioning 2017
task.
"""
# There is no real dialog in this task, so for the purposes of display_data, we
# include a generic question that applies to all images.
QUESTION = "Describe the above picture in a sentence."
def load_candidates(datapath, datatype, version):
if not datatype.startswith('train'):
suffix = 'captions_{}{}.json'
suffix_val = suffix.format('val', version)
val_path = os.path.join(
datapath, 'COCO_{}_Caption'.format(version), 'annotations', suffix_val
)
val = json.load(open(val_path))['annotations']
val_caps = [x['caption'] for x in val]
if datatype.startswith('test'):
suffix_train = suffix.format('train', version)
train_path = os.path.join(
datapath, 'COCO_{}_Caption'.format(version), 'annotations', suffix_train
)
train = json.load(open(train_path))['annotations']
train_caps = [x['caption'] for x in train]
test_caps = train_caps + val_caps
return test_caps
else:
return val_caps
else:
return None
def _path(opt, version):
if version == '2014':
build_2014(opt)
buildImage_2014(opt)
elif version == '2017':
build_2017(opt)
buildImage_2017(opt)
else:
raise Exception('Unknown version for COCO Captions: %s' % version)
dt = opt['datatype'].split(':')[0]
if dt == 'train':
annotation_suffix = 'train{}'.format(version)
img_suffix = os.path.join(
'train{}'.format(version),
'COCO_train{}_'.format(version) if version == '2014' else '',
)
elif dt == 'valid' or (dt == 'test' and version == '2014'):
annotation_suffix = 'val{}'.format(version)
img_suffix = os.path.join(
'val{}'.format(version),
'COCO_val{}_'.format(version) if version == '2014' else '',
)
elif dt == 'test':
annotation_suffix = 'None'
img_suffix = os.path.join(
'test{}'.format(version),
'COCO_test{}_'.format(version) if version == '2014' else '',
)
else:
raise RuntimeError('Not valid datatype.')
if version == '2017':
test_info_path = os.path.join(
opt['datapath'],
'COCO_2017_Caption',
'annotations',
'image_info_test2017.json',
)
annotation_path = os.path.join(
opt['datapath'],
'COCO_2017_Caption',
'annotations',
'captions_' + annotation_suffix + '.json',
)
else:
test_info_path = None
annotation_path = os.path.join(
opt['datapath'], 'COCO_2014_Caption', 'dataset_coco.json'
)
image_path = os.path.join(
opt['datapath'], 'COCO-IMG-{}'.format(version), img_suffix
)
return test_info_path, annotation_path, image_path
class DefaultTeacher(FixedDialogTeacher):
"""
COCO default teacher that expects open-ended descriptions of images.
"""
def __init__(self, opt, shared=None, version='2017'):
super().__init__(opt, shared)
self.version = version
self.image_mode = opt.get('image_mode', 'no_image_model')
self.use_intro = opt.get('use_intro', False)
self.num_cands = opt.get('num_cands', -1)
self.include_rest_val = opt.get('include_rest_val', False)
test_info_path, annotation_path, self.image_path = _path(opt, version)
self.test_split = opt['test_split']
if shared:
# another instance was set up already, just reference its data
if 'annotation' in shared:
self.annotation = shared['annotation']
self.image_loader = shared['image_loader']
if 'cands' in shared:
self.cands = shared['cands']
else:
# need to set up data from scratch
self._setup_data(test_info_path, annotation_path, opt)
self.image_loader = ImageLoader(opt)
self.reset()
@classmethod
def add_cmdline_args(
cls, parser: ParlaiParser, partial_opt: Optional[Opt] = None
) -> ParlaiParser:
super().add_cmdline_args(parser, partial_opt)
agent = parser.add_argument_group('COCO Caption arguments')
agent.add_argument(
'--use_intro',
type='bool',
default=False,
help='Include an intro question with each image \
for readability (e.g. for coco_caption, \
Describe the above picture in a sentence.)',
)
agent.add_argument(
'--num_cands',
type=int,
default=150,
help='Number of candidates to use during \
evaluation, setting to -1 uses all.',
)
agent.add_argument(
'--include_rest_val',
type='bool',
default=False,
help='Include unused validation images in training',
)
agent.add_argument(
'--test-split',
type=int,
default=-1,
choices=[-1, 0, 1, 2, 3, 4],
help='Which 1k image split of dataset to use for candidates'
'if -1, use all 5k test images',
)
return parser
def reset(self):
super().reset() # call parent reset so other fields can be set up
self.example = None # set up caching fields
self.imageEpochDone = False
def num_examples(self):
if self.version == '2014' or not self.datatype.startswith('test'):
return len(self.annotation)
else:
# For 2017, we only have annotations for the train and val sets,
# so for the test set we need to determine how many images we have.
return len(self.test_info['images'])
def num_episodes(self):
return self.num_examples()
def submit_load_request(self, image_id, split=None):
if split == 'restval':
img_path = self.image_path.replace('train', 'val')
else:
img_path = self.image_path
img_path += '%012d.jpg' % (image_id)
self.data_loader.request_load(
self.receive_data, self.image_loader.load, (img_path,)
)
def get(self, episode_idx, entry_idx=0):
action = {'episode_done': True}
if self.use_intro:
action['text'] = QUESTION
if self.version == '2014':
ep = self.annotation[episode_idx]
action['labels'] = [s['raw'] for s in ep['sentences']]
action['image_id'] = ep['cocoid']
action['split'] = ep['split']
if not self.datatype.startswith('train'):
if self.num_cands > 0:
labels = action['labels']
cands_to_sample = [c for c in self.cands if c not in labels]
cands = (
random.Random(episode_idx).sample(
cands_to_sample, self.num_cands
)
) + labels
random.shuffle(cands)
action['label_candidates'] = cands
else:
action['label_candidates'] = self.cands
else:
if not self.datatype.startswith('test'):
# test set annotations are not available for this dataset
anno = self.annotation[episode_idx]
action['labels'] = [anno['caption']]
action['image_id'] = anno['image_id']
if not self.datatype.startswith('train'):
if self.num_cands == -1:
candidates = self.cands
else:
# Can only randomly select from validation set
candidates = random.Random(episode_idx).choices(
self.cands, k=self.num_cands
)
if anno['caption'] not in candidates:
candidates.pop(0)
else:
candidates.remove(anno['caption'])
candidate_labels = [anno['caption']]
candidate_labels += candidates
action['label_candidates'] = candidate_labels
else:
if self.num_cands == -1:
candidates = self.cands
else:
# Can only randomly select from validation set
candidates = random.Random(episode_idx).choices(
self.cands, k=self.num_cands
)
action['label_candidates'] = candidates
action['image_id'] = self.test_info['images'][episode_idx]['id']
return action
def next_example(self):
"""
Returns the next example from this dataset after starting to queue up the next
example.
"""
ready = None
# pull up the currently queued example
if self.example is not None:
if self.image_mode != 'no_image_model' and 'image_id' in self.example:
# move the image we loaded in the background into the example
image = self.data_queue.get()
self.example['image'] = image
ready = (self.example, self.imageEpochDone)
# get the next base example: super().next_example() calls self.get()
self.example, self.imageEpochDone = super().next_example()
if self.image_mode != 'no_image_model' and 'image_id' in self.example:
# load the next image in the background
image_id = self.example['image_id']
split = self.example.get('split', None)
self.submit_load_request(image_id, split)
# Try to return the previously cached example
if ready is None:
return self.next_example()
else:
return ready
def share(self):
shared = super().share()
if hasattr(self, 'annotation'):
shared['annotation'] = self.annotation
shared['image_loader'] = self.image_loader
if hasattr(self, 'cands'):
shared['cands'] = self.cands
return shared
def _setup_data(self, test_info_path, annotation_path, opt):
if self.version == '2014':
with PathManager.open(annotation_path) as data_file:
raw_data = json.load(data_file)['images']
if 'train' in self.datatype:
self.annotation = [d for d in raw_data if d['split'] == 'train']
if self.include_rest_val:
self.annotation += [d for d in raw_data if d['split'] == 'restval']
elif 'valid' in self.datatype:
self.annotation = [d for d in raw_data if d['split'] == 'val']
self.cands = [
l
for d in self.annotation
for l in [s['raw'] for s in d['sentences']]
]
else:
self.annotation = [d for d in raw_data if d['split'] == 'test']
if self.test_split != -1:
start = self.test_split * 1000
end = (self.test_split + 1) * 1000
self.annotation = self.annotation[start:end]
self.cands = [
l
for d in self.annotation
for l in [s['raw'] for s in d['sentences']]
]
else:
if not self.datatype.startswith('test'):
print('loading: ' + annotation_path)
with PathManager.open(annotation_path) as data_file:
self.annotation = json.load(data_file)['annotations']
else:
print('loading: ' + test_info_path)
with PathManager.open(test_info_path) as data_file:
self.test_info = json.load(data_file)
if not self.datatype.startswith('train'):
self.cands = load_candidates(
opt['datapath'], opt['datatype'], self.version
)
class V2014Teacher(DefaultTeacher):
def __init__(self, opt, shared=None):
super(V2014Teacher, self).__init__(opt, shared, '2014')
class V2017Teacher(DefaultTeacher):
def __init__(self, opt, shared=None):
super(V2017Teacher, self).__init__(opt, shared, '2017')