HarryLee/eCommerceImageCaptioning

1import json2import sys3import os.path as op4 5from pycocotools.coco import COCO6from pycocoevalcap.eval import COCOEvalCap7 8 9def evaluate_on_coco_caption(res_file, label_file, outfile=None):10    """11    res_file: txt file, each row is [image_key, json format list of captions].12             Each caption is a dict, with fields "caption", "conf".13    label_file: JSON file of ground truth captions in COCO format.14    """15    coco = COCO(label_file)16    cocoRes = coco.loadRes(res_file)17    cocoEval = COCOEvalCap(coco, cocoRes)18 19    # evaluate on a subset of images by setting20    # cocoEval.params['image_id'] = cocoRes.getImgIds()21    # please remove this line when evaluating the full validation set22    cocoEval.params['image_id'] = cocoRes.getImgIds()23 24    # evaluate results25    # SPICE will take a few minutes the first time, but speeds up due to caching26    cocoEval.evaluate()27    result = cocoEval.eval28    if not outfile:29        print(result)30    else:31        with open(outfile, 'w') as fp:32            json.dump(result, fp, indent=4)33    return result34 35 36if __name__ == "__main__":37    if len(sys.argv) == 3:38        evaluate_on_coco_caption(sys.argv[1], sys.argv[2])39    elif len(sys.argv) == 4:40        evaluate_on_coco_caption(sys.argv[1], sys.argv[2], sys.argv[3])41    else:42        raise NotImplementedError