File size: 1,274 Bytes
32b542e |
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 |
import json
from collections import defaultdict
original_json = json.load(open("mscoco_dataset/new_annotations/dataset_coco.json"))
subsets = ['train', 'val', 'test']
savepath = "mscoco_dataset/new_annotations"
import os
if not os.path.exists(savepath):
os.makedirs(savepath)
savename = {
'train': "captions_train113k.json",
'val': "captions_val5k.json",
'test': "captions_test5k.json",
}
imagefields = defaultdict(list)
annotationsfields = defaultdict(list)
for imagecaps in original_json['images']:
filepath = imagecaps['filepath']
filename = imagecaps['filename']
image_id = int(filename.split(".")[0].split('_')[-1])
split = imagecaps['split']
if split == 'restval':
split = 'train'
imagefields[split].append({
"file_name": filename,
"file_path": filepath,
"id": image_id
})
for sen in imagecaps['sentences']:
annotationsfields[split].append({
"image_id": image_id,
"id": sen["sentid"],
"caption": sen["raw"],
})
for subset in subsets:
data = {
"images": imagefields[subset],
"annotations": annotationsfields[subset]
}
json.dump(data, open(os.path.join(savepath, savename[subset]), "w"))
pass
|