Spaces:
Running
on
T4
Running
on
T4
File size: 6,546 Bytes
a277bb8 |
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) |