-
Notifications
You must be signed in to change notification settings - Fork 11
/
prepare_captions.py
268 lines (206 loc) · 9.52 KB
/
prepare_captions.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
"""Data preparation for training image captioning model
This script will do the followings:
1) Come up with a vocab list by pooling all training and val captions
2) Convert each word from captions to an integer based on the vocab list
3) Produce image-name-index mapping, that maps an image to an integer based on its name (e.g. COCO_train2014_000000417432.jpg -> 1)
4) Rename all images using the image-name-index mapping above
"""
import json
import os
import collections
import tensorflow as tf
import re
import h5py
import argparse
import sys
import numpy as np
import pandas as pd
FLAGS = None
BUFFER_TOKENS = ['<NULL>', '<START>', '<END>', '<UNK>']
def _parse_sentence(s):
s = s.replace('.', '')
s = s.replace(',', '')
s = s.replace('"', '')
s = s.replace("'", '')
s = s.lower()
s = re.sub("\s\s+", " ", s)
s = s.split(' ')
return s
def preprocess_json_files(path_to_dir):
"""Extract captions from each file and combine into lists, as well as image ids, and returned as dict"""
assert os.path.exists(path_to_dir), 'Path to directory of files does not exist!'
results = {}
for file in os.listdir(path_to_dir):
if 'captions_train2014' not in file and 'captions_val2014' not in file:
print("Skipping file {}".format(file))
continue
temp_path = os.path.join(path_to_dir, file)
with open(temp_path, 'r') as f:
data = json.load(f)
caps = data['annotations']
images = [item['image_id'] for item in caps]
urls = {}
for img in data['images']:
urls[img['id']] = img['flickr_url']
caps = [_parse_sentence(item['caption']) for item in caps]
results[file] = (caps, images, urls)
del data
# return dict of each file, having list of captions and image_ids
"""
results is a dict of two files (train and val), each of which has a caps list (results[file1][0]) and a images list (results[file1][1]), and urls dict
(results[file1][2]). cap list is a list of sentences(list of words), images list is a list of image ids(integers), and urls dict is a dict mapping each
image id to its url
"""
return results
def rename_images(dir, image_id_to_idx):
image_dict = pd.read_csv(image_id_to_idx) # cols: image_idx, image_id
image_dict = image_dict.set_index('image_id')
image_dict = image_dict['image_index'].to_dict()
for img_name in os.listdir(dir):
original_img_path = os.path.join(dir, img_name)
temp_num = int(re.split('\.|_', img_name)[-2])
temp_num = image_dict[temp_num] # convert image id to idx
new_img_path = os.path.join(dir, '{0}.jpg'.format(temp_num))
os.rename(original_img_path, new_img_path)
print("Renaming images for folder {} done. ".format(dir))
def main(_):
## get the vocaboluary
list_of_all_words = None
results = preprocess_json_files(FLAGS.file_dir)
for k, v in results.items():
if list_of_all_words is None:
list_of_all_words = results[k][0].copy()
else:
list_of_all_words += results[k][0]
list_of_all_words = [item for sublist in list_of_all_words for item in sublist]
counter = collections.Counter(list_of_all_words)
vocab = counter.most_common(FLAGS.total_vocab)
print("\nVocab generated! Most, median and least frequent words from the vocab are: \n{0}\n{1}\n{2}\n".format(vocab[0], vocab[int(FLAGS.total_vocab/2)], vocab[-1]))
## create word_to_idx, and idx_to_word
vocab = [i[0] for i in vocab]
word_to_idx = {}
idx_to_word = {}
# add in BUFFER_TOKENS
for i in range(len(BUFFER_TOKENS)):
idx_to_word[int(i)] = BUFFER_TOKENS[i]
word_to_idx[BUFFER_TOKENS[i]] = i
for i in range(len(vocab)):
word_to_idx[vocab[i]] = i + len(BUFFER_TOKENS)
idx_to_word[int(i + len(BUFFER_TOKENS))] = vocab[i]
word_dict = {}
word_dict['idx_to_word'] = idx_to_word
word_dict['word_to_idx'] = word_to_idx
with open(os.path.join(FLAGS.file_dir, 'coco2014_vocab.json'), 'w') as f:
json.dump(word_dict, f)
## convert sentences into encoding/integers
# pad all sentence to length of FLAGS.padding_len - 2
def _convert_sentence_to_numbers(s):
"""Convert a sentence s (a list of words) to list of numbers using word_to_idx"""
UNK_IDX = BUFFER_TOKENS.index('<UNK>')
NULL_IDX = BUFFER_TOKENS.index('<NULL>')
END_IDX = BUFFER_TOKENS.index('<END>')
s_encoded = [word_to_idx.get(w, UNK_IDX) for w in s]
s_encoded += [END_IDX]
s_encoded += [NULL_IDX] * (FLAGS.padding_len - 1 - len(s_encoded))
return s_encoded
h = h5py.File(os.path.join(FLAGS.file_dir,'coco2014_captions.h5'), 'w')
for k, _ in results.items():
results_to_save = {}
all_captions = results[k][0] # list of lists of words
all_images = results[k][1]
all_urls = results[k][2]
all_captions = [_convert_sentence_to_numbers(s) for s in all_captions] # list of numbers
valid_rows = [i for i in range(len(all_captions)) if len(all_captions[i]) == FLAGS.padding_len-1]
all_captions= [row for row in all_captions if len(row) == FLAGS.padding_len-1]
all_captions = np.array(all_captions)
all_images = np.array(all_images)
all_images = all_images[valid_rows]
assert all_images.shape[0] == all_captions.shape[0], "Processing error! all_captions and all_images diff in length."
# concatenate START and END tokens at two sides
START_TOKEN = BUFFER_TOKENS.index('<START>')
#END_TOKEN = BUFFER_TOKENS.index('<END>')
col_start = np.array([START_TOKEN] * all_images.shape[0]).reshape(-1, 1)
#col_end = np.array([END_TOKEN] * all_images.shape[0]).reshape(-1, 1)
all_captions = np.hstack([col_start, all_captions])
## create dicts that maps image ids to 0,...,total_images - image_idx_to_id, image_id_to_idx
image_ids = set(all_images)
image_idx = list(range(len(image_ids)))
image_id_to_idx = {}
image_idx_to_id = {}
for idx, id in enumerate(image_ids):
image_id_to_idx[id] = idx
image_idx_to_id[idx] = id
all_images_idx = np.array([image_id_to_idx.get(id) for id in all_images])
## save all the data
if 'train' in k:
h.create_dataset('train_captions', data=all_captions)
h.create_dataset('train_image_idx', data=all_images_idx)
df = pd.DataFrame.from_dict(image_id_to_idx, 'index')
df['image_id'] = df.index.values
df.columns = ['image_index', 'image_id']
df.to_csv(os.path.join(FLAGS.file_dir, 'train_image_id_to_idx.csv'), index = False)
## write urls file to local as train2014_urls.txt
with open(os.path.join(FLAGS.file_dir, 'train2014_urls.txt'), 'w') as f:
for idx in range(len(image_idx_to_id)):
this_url = all_urls[image_idx_to_id[idx]]
f.write(this_url + '\n')
elif 'val' in k:
h.create_dataset('val_captions', data=all_captions)
h.create_dataset('val_image_idx', data=all_images_idx)
df = pd.DataFrame.from_dict(image_id_to_idx, 'index')
df['image_id'] = df.index.values
df.columns = ['image_index', 'image_id']
df.to_csv(os.path.join(FLAGS.file_dir, 'val_image_id_to_idx.csv'), index = False)
## write urls file to local as val2014_urls.txt
with open(os.path.join(FLAGS.file_dir, 'val2014_urls.txt'), 'w') as f:
for idx in range(len(image_idx_to_id)):
this_url = all_urls[image_idx_to_id[idx]]
f.write(this_url + '\n')
else:
print("Strange file name found in dir: {0}, \nit does not belong to train nor val, so it is not able to save results!".format(k))
h.close()
print("Data generation done.\n Start renaming images in sequence ...")
if FLAGS.train_image_dir != '':
train_dict = os.path.join(FLAGS.file_dir, 'train_image_id_to_idx.csv')
rename_images(FLAGS.train_image_dir, train_dict)
if FLAGS.val_image_dir != '':
val_dict = os.path.join(FLAGS.file_dir, 'val_image_id_to_idx.csv')
rename_images(FLAGS.val_image_dir, val_dict)
print("all done. ")
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--file_dir',
type=str,
#default='C:\\Users\\WAWEIMIN\\Google Drive\\ShowAndTellWeimin\\coco_captioning\\original_captioning',
default= '/home/ubuntu/COCO/dataset/COCO_captioning/',
help="""\
Path to captions_train2014.json, captions_val2014.json\
"""
)
parser.add_argument(
'--total_vocab',
type=int,
default=1000,
help='Total number of vacobulary to use.'
)
parser.add_argument(
'--padding_len',
type=int,
default=17,
help='Total len of padding the sentence.'
)
parser.add_argument(
'--train_image_dir',
type=str,
default='/home/ubuntu/COCO/dataset/train2014',
help='Absolute path to training dir containing images that are to be renamed.'
)
parser.add_argument(
'--val_image_dir',
type=str,
default='/home/ubuntu/COCO/dataset/val2014',
help='Absolute path to val dir containing images that are to be renamed.'
)
FLAGS, unparsed = parser.parse_known_args()
tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)