|
import json |
|
import collections |
|
import logging |
|
import sys |
|
|
|
if len(sys.argv) != 4: |
|
print("Provide .tsv file name, images dir, output file name. e.g. python coco.py coco_captions_train2017.json /mnt/disks/data-1/flickr8k/coco_train.json coco_dataset_train.json") |
|
exit(1) |
|
|
|
annotation_file = sys.argv[1] |
|
images_dir = sys.argv[2] |
|
output_file = sys.argv[3] |
|
|
|
logging.info("Processing COCO dataset") |
|
|
|
with open(annotation_file, "r") as f: |
|
annotations = json.load(f)["annotations"] |
|
|
|
image_path_to_caption = collections.defaultdict(list) |
|
for element in annotations: |
|
caption = f"{element['caption'].lower().rstrip('.')}" |
|
image_path = images_dir + "/%012d.jpg" % (element["image_id"]) |
|
image_path_to_caption[image_path].append(caption) |
|
|
|
lines = [] |
|
for image_path, captions in image_path_to_caption.items(): |
|
lines.append(json.dumps({"image_path": image_path, "captions": captions})) |
|
|
|
train_lines = lines[:-10_001] |
|
valid_lines = lines[-10_001:] |
|
|
|
with open(output_file+"_train.json", "w") as f: |
|
f.write("\n".join(train_lines)) |
|
|
|
with open(output_file+"_val.json", "w") as f: |
|
f.write("\n".join(valid_lines)) |
|
|
|
logging.info(f"Processing COCO dataset done. {len(lines)} images processed.") |
|
|
|
|
|
|
