compute_scores.py

"""
Computer all metrics (BMCR) of all splits
It is used to as the inputs of the compute_dataslice function
and as pre-computed cider scores at training time

"""

import os
import sys
import json
import argparse
import string
import itertools
import numpy as np
from collections import OrderedDict

#sys.path.append("cider")
#from pyciderevalcap.cider.cider import Cider
#from pyciderevalcap.ciderD.ciderD import CiderD

sys.path.append('coco-caption')
#from pycocotools.coco import COCO
#from pycocoevalcap.eval import COCOEvalCap

from pycocoevalcap.bleu.bleu import Bleu
from pycocoevalcap.meteor.meteor import Meteor
from pycocoevalcap.rouge.rouge import Rouge
from pycocoevalcap.cider.cider import Cider

import logging
from datetime import datetime

import utils
logger = logging.getLogger(__name__)
import pickle

if __name__ == "__main__":
    logging.basicConfig(
        level=logging.DEBUG, format='%(asctime)s:%(levelname)s: %(message)s')
    parser = argparse.ArgumentParser()

    #parser.add_argument('pred_file', type=str, help='')
    parser.add_argument('cocofmt_file', type=str, help='')
    parser.add_argument('output_pkl', type=str, help='')
    parser.add_argument(
        '--seq_per_img',
        type=int,
        default=20,
        help='Number of caption per image/video')
    parser.add_argument(
        '--remove_in_ref',
        default=False,
        action='store_true',
        help='Remove current caption in the ref set')

    args = parser.parse_args()
    logger.info('Input arguments: %s', args)

    start = datetime.now()

    #logger.info('Loading prediction: %s', args.pred_file)
    #preds = json.load(open(args.pred_file))['predictions']
    #preds = {p['image_id']: [p['caption']] for p in preds}
    #scorer = CiderD(df=args.cached_tokens_file)

    logger.info('Setting up scorers...')
    scorers = [(Bleu(), "Bleu_4"), (Meteor(), "METEOR"), (Rouge(), "ROUGE_L"),
               (Cider(), "CIDEr")]

    logger.info('loading gt refs: %s', args.cocofmt_file)
    gt_refs = utils.load_gt_refs(args.cocofmt_file)

    logger.info('Computing score...')
    #score, scores = computer_score(preds, gt_refs, scorer)
    videos = sorted(gt_refs.keys())

    gt_scores = {}
    for scorer, method in scorers:
        gt_scores[method] = np.zeros((len(gt_refs), args.seq_per_img))

    for i in range(args.seq_per_img):
        logger.info('taking caption: %d', i)
        preds_i = {v: [gt_refs[v][i]] for v in videos}

        # removing the refs at i
        if args.remove_in_ref:
            gt_refs_i = {v: gt_refs[v][:i] + gt_refs[v][i + 1:] for v in videos}
        else:
            gt_refs_i = gt_refs

        for scorer, method in scorers:
            score_i, scores_i = scorer.compute_score(gt_refs_i, preds_i)

            if method == 'Bleu_4':
                score_i = score_i[-1]
                scores_i = scores_i[-1]

            # happens for BLeu and METEOR
            if type(scores_i) == list:
                scores_i = np.array(scores_i)

            gt_scores[method][:, i] = scores_i
            logger.info('%s: %f', method, score_i)

    pickle.dump(
        gt_scores,
        open(args.output_pkl, 'wb'),
        protocol=pickle.HIGHEST_PROTOCOL)

    logger.info('Time: %s', datetime.now() - start)