-
Notifications
You must be signed in to change notification settings - Fork 42
/
flickr30ke2odvg.py
180 lines (160 loc) · 6.39 KB
/
flickr30ke2odvg.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
import xml.etree.ElementTree as ET
import jsonlines
import random
from tqdm import tqdm
import argparse
import os
import glob
def get_sentence_data(fn):
"""
Parses a sentence file from the Flickr30K Entities dataset
input:
fn - full file path to the sentence file to parse
output:
a list of dictionaries for each sentence with the following fields:
sentence - the original sentence
phrases - a list of dictionaries for each phrase with the
following fields:
phrase - the text of the annotated phrase
first_word_index - the position of the first word of
the phrase in the sentence
phrase_id - an identifier for this phrase
phrase_type - a list of the coarse categories this
phrase belongs to
"""
with open(fn, 'r') as f:
sentences = f.read().split('\n')
annotations = []
for sentence in sentences:
if not sentence:
continue
first_word = []
phrases = []
phrase_id = []
phrase_type = []
words = []
current_phrase = []
add_to_phrase = False
for token in sentence.split():
if add_to_phrase:
if token[-1] == ']':
add_to_phrase = False
token = token[:-1]
current_phrase.append(token)
phrases.append(' '.join(current_phrase))
current_phrase = []
else:
current_phrase.append(token)
words.append(token)
else:
if token[0] == '[':
add_to_phrase = True
first_word.append(len(words))
parts = token.split('/')
phrase_id.append(parts[1][3:])
phrase_type.append(parts[2:])
else:
words.append(token)
sentence_data = {'sentence' : ' '.join(words), 'phrases' : []}
for index, phrase, p_id, p_type in zip(first_word, phrases, phrase_id, phrase_type):
sentence_data['phrases'].append({'first_word_index' : index,
'phrase' : phrase,
'phrase_id' : p_id,
'phrase_type' : p_type})
annotations.append(sentence_data)
return annotations
def get_annotations(fn):
"""
Parses the xml files in the Flickr30K Entities dataset
input:
fn - full file path to the annotations file to parse
output:
dictionary with the following fields:
scene - list of identifiers which were annotated as
pertaining to the whole scene
nobox - list of identifiers which were annotated as
not being visible in the image
boxes - a dictionary where the fields are identifiers
and the values are its list of boxes in the
[xmin ymin xmax ymax] format
"""
tree = ET.parse(fn)
root = tree.getroot()
filename = root.findall('filename')[0].text
size_container = root.findall('size')[0]
anno_info = {'filename': filename, 'boxes' : {}, 'scene' : [], 'nobox' : []}
for size_element in size_container:
anno_info[size_element.tag] = int(size_element.text)
for object_container in root.findall('object'):
for names in object_container.findall('name'):
box_id = names.text
box_container = object_container.findall('bndbox')
if len(box_container) > 0:
if box_id not in anno_info['boxes']:
anno_info['boxes'][box_id] = []
xmin = int(box_container[0].findall('xmin')[0].text) - 1
ymin = int(box_container[0].findall('ymin')[0].text) - 1
xmax = int(box_container[0].findall('xmax')[0].text) - 1
ymax = int(box_container[0].findall('ymax')[0].text) - 1
anno_info['boxes'][box_id].append([xmin, ymin, xmax, ymax])
else:
nobndbox = int(object_container.findall('nobndbox')[0].text)
if nobndbox > 0:
anno_info['nobox'].append(box_id)
scene = int(object_container.findall('scene')[0].text)
if scene > 0:
anno_info['scene'].append(box_id)
return anno_info
def gen_record(sd, an):
filename = an["filename"]
caption = sd["sentence"]
regions = []
for ph in sd["phrases"]:
if ph["phrase_id"] in an["boxes"]:
for box in an["boxes"][ph["phrase_id"]]:
regions.append(
{
"phrase": ph["phrase"],
"bbox": box
}
)
if len(regions) < 1:
print("no phrase regions")
return None
return {
"filename": filename,
"height": an["height"],
"width": an["width"],
"grounding":{
"caption": caption,
"regions": regions
}
}
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="flickr30k entities to ODVG List.")
parser.add_argument("--root", type=str, default="", help="Source anno root")
parser.add_argument("--output_file", type=str, default="flickr30k_entities_odvg.jsonl")
parser.add_argument("--osoi", action="store_true", default=False)
args = parser.parse_args()
print(args)
odvg_anno = []
sentence_list = os.path.join(args.root, "Sentences")
annotation_list = os.path.join(args.root, "Annotations")
sentence_list = sorted(glob.glob(sentence_list + "/*"))
annotation_list = sorted(glob.glob(annotation_list + "/*"))
len_anno = len(annotation_list)
for idx in tqdm(range(len_anno)):
sds = get_sentence_data(sentence_list[idx])
an = get_annotations(annotation_list[idx])
if args.osoi:
sd = sds[random.randint(0, len(sds)-1)]
x = gen_record(sd, an)
if x:
odvg_anno.append(x)
else:
for sd in sds:
x = gen_record(sd, an)
if x:
odvg_anno.append(x)
with jsonlines.open(args.output_file, mode="w") as fwriter:
fwriter.write_all(odvg_anno)