"""
author: Timothy C. Arlen
date: 28 Feb 2018

Calculate Mean Average Precision (mAP) for a set of bounding boxes corresponding to specific
image Ids. Usage:

> python calculate_mean_ap.py

Will display a plot of precision vs recall curves at 10 distinct IoU thresholds as well as output
summary information regarding the average precision and mAP scores.

NOTE: Requires the files `ground_truth_boxes.json` and `predicted_boxes.json` which can be
downloaded fromt this gist.
"""

from __future__ import absolute_import, division, print_function

from copy import deepcopy
import json
import glob
import os
import time

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

sns.set_style('white')
sns.set_context('poster')

COLORS = [
    '#1f77b4', '#aec7e8', '#ff7f0e', '#ffbb78', '#2ca02c',
    '#98df8a', '#d62728', '#ff9896', '#9467bd', '#c5b0d5',
    '#8c564b', '#c49c94', '#e377c2', '#f7b6d2', '#7f7f7f',
    '#c7c7c7', '#bcbd22', '#dbdb8d', '#17becf', '#9edae5']


def calc_iou_individual(pred_box, gt_box):
    """Calculate IoU of single predicted and ground truth box

    Args:
        pred_box (list of floats): location of predicted object as
            [xmin, ymin, xmax, ymax]
        gt_box (list of floats): location of ground truth object as
            [xmin, ymin, xmax, ymax]

    Returns:
        float: value of the IoU for the two boxes.

    Raises:
        AssertionError: if the box is obviously malformed
    """
    x1_t, y1_t, x2_t, y2_t = gt_box
    x1_p, y1_p, x2_p, y2_p = pred_box

    if (x1_p > x2_p) or (y1_p > y2_p):
        raise AssertionError(
            "Prediction box is malformed? pred box: {}".format(pred_box))
    if (x1_t > x2_t) or (y1_t > y2_t):
        raise AssertionError(
            "Ground Truth box is malformed? true box: {}".format(gt_box))

    if (x2_t < x1_p or x2_p < x1_t or y2_t < y1_p or y2_p < y1_t):
        return 0.0

    far_x = np.min([x2_t, x2_p])
    near_x = np.max([x1_t, x1_p])
    far_y = np.min([y2_t, y2_p])
    near_y = np.max([y1_t, y1_p])

    inter_area = (far_x - near_x + 1) * (far_y - near_y + 1)
    true_box_area = (x2_t - x1_t + 1) * (y2_t - y1_t + 1)
    pred_box_area = (x2_p - x1_p + 1) * (y2_p - y1_p + 1)
    iou = inter_area / (true_box_area + pred_box_area - inter_area)
    return iou


def get_single_image_results(gt_boxes, pred_boxes, iou_thr):
    """Calculates number of true_pos, false_pos, false_neg from single batch of boxes.

    Args:
        gt_boxes (list of list of floats): list of locations of ground truth
            objects as [xmin, ymin, xmax, ymax]
        pred_boxes (dict): dict of dicts of 'boxes' (formatted like `gt_boxes`)
            and 'scores'
        iou_thr (float): value of IoU to consider as threshold for a
            true prediction.

    Returns:
        dict: true positives (int), false positives (int), false negatives (int)
    """

    all_pred_indices = range(len(pred_boxes))
    all_gt_indices = range(len(gt_boxes))
    if len(all_pred_indices) == 0:
        tp = 0
        fp = 0
        fn = len(gt_boxes)
        return {'true_pos': tp, 'false_pos': fp, 'false_neg': fn}
    if len(all_gt_indices) == 0:
        tp = 0
        fp = len(pred_boxes)
        fn = 0
        return {'true_pos': tp, 'false_pos': fp, 'false_neg': fn}

    gt_idx_thr = []
    pred_idx_thr = []
    ious = []
    for ipb, pred_box in enumerate(pred_boxes):
        for igb, gt_box in enumerate(gt_boxes):
            iou = calc_iou_individual(pred_box, gt_box)
            if iou > iou_thr:
                gt_idx_thr.append(igb)
                pred_idx_thr.append(ipb)
                ious.append(iou)

    args_desc = np.argsort(ious)[::-1]
    if len(args_desc) == 0:
        # No matches
        tp = 0
        fp = len(pred_boxes)
        fn = len(gt_boxes)
    else:
        gt_match_idx = []
        pred_match_idx = []
        for idx in args_desc:
            gt_idx = gt_idx_thr[idx]
            pr_idx = pred_idx_thr[idx]
            # If the boxes are unmatched, add them to matches
            if (gt_idx not in gt_match_idx) and (pr_idx not in pred_match_idx):
                gt_match_idx.append(gt_idx)
                pred_match_idx.append(pr_idx)
        tp = len(gt_match_idx)
        fp = len(pred_boxes) - len(pred_match_idx)
        fn = len(gt_boxes) - len(gt_match_idx)

    return {'true_pos': tp, 'false_pos': fp, 'false_neg': fn}


def calc_precision_recall(img_results):
    """Calculates precision and recall from the set of images

    Args:
        img_results (dict): dictionary formatted like:
            {
                'img_id1': {'true_pos': int, 'false_pos': int, 'false_neg': int},
                'img_id2': ...
                ...
            }

    Returns:
        tuple: of floats of (precision, recall)
    """
    true_pos = 0; false_pos = 0; false_neg = 0
    for _, res in img_results.items():
        true_pos += res['true_pos']
        false_pos += res['false_pos']
        false_neg += res['false_neg']

    try:
        precision = true_pos/(true_pos + false_pos)
    except ZeroDivisionError:
        precision = 0.0
    try:
        recall = true_pos/(true_pos + false_neg)
    except ZeroDivisionError:
        recall = 0.0

    return (precision, recall)

def get_model_scores_map(pred_boxes):
    """Creates a dictionary of from model_scores to image ids.

    Args:
        pred_boxes (dict): dict of dicts of 'boxes' and 'scores'

    Returns:
        dict: keys are model_scores and values are image ids (usually filenames)

    """
    model_scores_map = {}
    for img_id, val in pred_boxes.items():
        for score in val['scores']:
            if score not in model_scores_map.keys():
                model_scores_map[score] = [img_id]
            else:
                model_scores_map[score].append(img_id)
    return model_scores_map

def get_avg_precision_at_iou(gt_boxes, pred_boxes, iou_thr=0.5):
    """Calculates average precision at given IoU threshold.

    Args:
        gt_boxes (list of list of floats): list of locations of ground truth
            objects as [xmin, ymin, xmax, ymax]
        pred_boxes (list of list of floats): list of locations of predicted
            objects as [xmin, ymin, xmax, ymax]
        iou_thr (float): value of IoU to consider as threshold for a
            true prediction.

    Returns:
        dict: avg precision as well as summary info about the PR curve

        Keys:
            'avg_prec' (float): average precision for this IoU threshold
            'precisions' (list of floats): precision value for the given
                model_threshold
            'recall' (list of floats): recall value for given
                model_threshold
            'models_thrs' (list of floats): model threshold value that
                precision and recall were computed for.
    """
    model_scores_map = get_model_scores_map(pred_boxes)
    sorted_model_scores = sorted(model_scores_map.keys())

    # Sort the predicted boxes in descending order (lowest scoring boxes first):
    for img_id in pred_boxes.keys():
        arg_sort = np.argsort(pred_boxes[img_id]['scores'])
        pred_boxes[img_id]['scores'] = np.array(pred_boxes[img_id]['scores'])[arg_sort].tolist()
        pred_boxes[img_id]['boxes'] = np.array(pred_boxes[img_id]['boxes'])[arg_sort].tolist()
        if 'class' in pred_boxes[img_id]: # TODO or re add the class field when splitting
            pred_boxes[img_id]['class'] = np.array(pred_boxes[img_id]['class'])[arg_sort].tolist()

    pred_boxes_pruned = deepcopy(pred_boxes)

    precisions = []
    recalls = []
    model_thrs = []
    img_results = {}
    # Loop over model score thresholds and calculate precision, recall
    for ithr, model_score_thr in enumerate(sorted_model_scores[:-1]):
        # On first iteration, define img_results for the first time:
        img_ids = gt_boxes.keys() if ithr == 0 else model_scores_map[model_score_thr]
        for img_id in img_ids:
            gt_boxes_img = gt_boxes[img_id]['boxes']
            box_scores = pred_boxes_pruned[img_id]['scores']
            start_idx = 0
            for score in box_scores:
                if score <= model_score_thr:
                    pred_boxes_pruned[img_id]
                    start_idx += 1
                else:
                    break

            # Remove boxes, scores of lower than threshold scores:
            pred_boxes_pruned[img_id]['scores'] = pred_boxes_pruned[img_id]['scores'][start_idx:]
            pred_boxes_pruned[img_id]['boxes'] = pred_boxes_pruned[img_id]['boxes'][start_idx:]
            if 'class' in pred_boxes[img_id]:  # TODO or re add the class field when splitting
                pred_boxes_pruned[img_id]['class'] = pred_boxes_pruned[img_id]['class'][start_idx:]

            # Recalculate image results for this image
            img_results[img_id] = get_single_image_results(
                gt_boxes_img, pred_boxes_pruned[img_id]['boxes'], iou_thr)

        prec, rec = calc_precision_recall(img_results)
        precisions.append(prec)
        recalls.append(rec)
        model_thrs.append(model_score_thr)

    precisions = np.array(precisions)
    recalls = np.array(recalls)
    prec_at_rec = []
    for recall_level in np.linspace(0.0, 1.0, 11):
        try:
            args = np.argwhere(recalls >= recall_level).flatten()
            prec = max(precisions[args])
        except ValueError:
            prec = 0.0
        prec_at_rec.append(prec)
    avg_prec = np.mean(prec_at_rec)

    return {
        'avg_prec': avg_prec,
        'precisions': precisions,
        'recalls': recalls,
        'model_thrs': model_thrs}


def split_boxes_by_class(gt_boxes, pred_boxes, classification: str):
    gt_boxes_class = {}
    pred_boxes_class = {}

    for gt_image in gt_boxes:
        gt_boxes_class[gt_image] = {}
        gt_boxes_class[gt_image]['boxes'] = []

        for gt_object_box, gt_object_class in zip(gt_boxes[gt_image]['boxes'], gt_boxes[gt_image]['class']):
            if gt_object_class == classification:
                gt_boxes_class[gt_image]['boxes'].append(gt_object_box)

    for pred_image in pred_boxes:
        pred_boxes_class[pred_image] = {}
        pred_boxes_class[pred_image]['boxes'] = []
        pred_boxes_class[pred_image]['scores'] = []

        for pred_object_box, pred_object_score, pred_object_class in \
                zip(pred_boxes[pred_image]['boxes'], pred_boxes[pred_image]['scores'], pred_boxes[pred_image]['class']):
            if pred_object_class == classification:
                pred_boxes_class[pred_image]['boxes'].append(pred_object_box)
                pred_boxes_class[pred_image]['scores'].append(pred_object_score)

    # print(classification)
    # print(gt_boxes_class)
    # print(pred_boxes_class)

    return gt_boxes_class, pred_boxes_class

# TODO this method can be rewritten such that boxes can be split by class within it
def get_avg_precision_at_iou_per_class(gt_boxes, pred_boxes, classification, iou_thr=0.5):
    """Calculates average precision at given IoU threshold.

    Args:
        gt_boxes (list of list of floats): list of locations of ground truth
            objects as [xmin, ymin, xmax, ymax]
        pred_boxes (list of list of floats): list of locations of predicted
            objects as [xmin, ymin, xmax, ymax]
        classification (str): the classification of object for which the AP
            should be calculated
        iou_thr (float): value of IoU to consider as threshold for a
            true prediction.

    Returns:
        dict: avg precision as well as summary info about the PR curve

        Keys:
            'avg_prec' (float): average precision for this IoU threshold
            'precisions' (list of floats): precision value for the given
                model_threshold
            'recall' (list of floats): recall value for given
                model_threshold
            'models_thrs' (list of floats): model threshold value that
                precision and recall were computed for.
    """
    model_scores_map = get_model_scores_map(pred_boxes)
    sorted_model_scores = sorted(model_scores_map.keys())

    # Sort the predicted boxes in descending order (lowest scoring boxes first):
    for img_id in pred_boxes.keys():
        arg_sort = np.argsort(pred_boxes[img_id]['scores'])
        pred_boxes[img_id]['scores'] = np.array(pred_boxes[img_id]['scores'])[arg_sort].tolist()
        pred_boxes[img_id]['boxes'] = np.array(pred_boxes[img_id]['boxes'])[arg_sort].tolist()
        pred_boxes[img_id]['class'] = np.array(pred_boxes[img_id]['class'])[arg_sort].tolist()

    pred_boxes_pruned = deepcopy(pred_boxes)

    precisions = []
    recalls = []
    model_thrs = []
    img_results = {}
    # Loop over model score thresholds and calculate precision, recall
    for ithr, model_score_thr in enumerate(sorted_model_scores[:-1]):
        # On first iteration, define img_results for the first time:
        img_ids = gt_boxes.keys() if ithr == 0 else model_scores_map[model_score_thr]
        for img_id in img_ids:
            gt_boxes_img = gt_boxes[img_id]['boxes']
            box_scores = pred_boxes_pruned[img_id]['scores']
            start_idx = 0
            for score in box_scores:
                if score <= model_score_thr:
                    pred_boxes_pruned[img_id]
                    start_idx += 1
                else:
                    break

            # Remove boxes, scores of lower than threshold scores:
            pred_boxes_pruned[img_id]['scores'] = pred_boxes_pruned[img_id]['scores'][start_idx:]
            pred_boxes_pruned[img_id]['boxes'] = pred_boxes_pruned[img_id]['boxes'][start_idx:]
            pred_boxes_pruned[img_id]['class'] = pred_boxes_pruned[img_id]['class'][start_idx:]

            # Recalculate image results for this image
            img_results[img_id] = get_single_image_results(
                gt_boxes_img, pred_boxes_pruned[img_id]['boxes'], iou_thr)

        prec, rec = calc_precision_recall(img_results)
        precisions.append(prec)
        recalls.append(rec)
        model_thrs.append(model_score_thr)

    precisions = np.array(precisions)
    recalls = np.array(recalls)
    prec_at_rec = []
    for recall_level in np.linspace(0.0, 1.0, 11):
        try:
            args = np.argwhere(recalls >= recall_level).flatten()
            prec = max(precisions[args])
        except ValueError:
            prec = 0.0
        prec_at_rec.append(prec)
    avg_prec = np.mean(prec_at_rec)

    return {
        'avg_prec': avg_prec,
        'precisions': precisions,
        'recalls': recalls,
        'model_thrs': model_thrs}


def plot_pr_curve(precisions, recalls, category='Objects',
                  label=None, color=None, ax=None, linestyle='solid'):
    """Simple plotting helper function"""

    if ax is None:
        plt.figure(figsize=(10,8))
        ax = plt.gca()

    if color is None:
        color = COLORS[0]
    # ax.scatter(recalls, precisions, label=label, s=20, color=color)
    # ax.scatter(recalls, precisions, label='_nolegend_', s=20, color=color)
    ax.plot(recalls, precisions, label=label, color=color, linestyle=linestyle)
    ax.set_xlabel('recall')
    ax.set_ylabel('precision')
    ax.set_title('Precision-Recall curve for {}'.format(category))
    ax.set_xlim([0.0,1.3])
    ax.set_ylim([0.0,1.2])
    return ax


if __name__ == "__main__":

    with open('ground_truth_boxes.json') as infile:
        gt_boxes = json.load(infile)

    with open('predicted_boxes.json') as infile:
        pred_boxes = json.load(infile)

    # Split all groundtruths and detections into subsets by the object detected
    gt_boxes_source, pred_boxes_source = split_boxes_by_class(gt_boxes, pred_boxes, 'source')
    gt_boxes_sidelobe, pred_boxes_sidelobe = split_boxes_by_class(gt_boxes, pred_boxes, 'sidelobe')
    gt_boxes_galaxy, pred_boxes_galaxy = split_boxes_by_class(gt_boxes, pred_boxes, 'galaxy')

    # Run per class mAP calculation for one IoU threshold, 0.5
    iou_thr = 0.5
    # calculate source mAP
    data = get_avg_precision_at_iou(gt_boxes_source, pred_boxes_source, iou_thr=iou_thr)
    print('Source avg precision: {:.4f}'.format(data['avg_prec']))
    # calculate sidelobe mAP
    data = get_avg_precision_at_iou(gt_boxes_sidelobe, pred_boxes_sidelobe, iou_thr=iou_thr)
    print('Sidelobe avg precision: {:.4f}'.format(data['avg_prec']))
    # calculate galaxy mAP
    data = get_avg_precision_at_iou(gt_boxes_galaxy, pred_boxes_galaxy, iou_thr=iou_thr)
    print('Galaxy avg precision: {:.4f}'.format(data['avg_prec']))

    # Runs it for one IoU threshold
    iou_thr = 0.5
    start_time = time.time()
    data = get_avg_precision_at_iou(gt_boxes, pred_boxes, iou_thr=iou_thr)
    end_time = time.time()
    print('Single IoU calculation took {:.4f} secs'.format(end_time - start_time))
    print('avg precision: {:.4f}'.format(data['avg_prec']))

    start_time = time.time()
    ax = None
    avg_precs = []
    iou_thrs = []
    for idx, iou_thr in enumerate(np.linspace(0.5, 0.95, 10)):
        data = get_avg_precision_at_iou(gt_boxes, pred_boxes, iou_thr=iou_thr)
        avg_precs.append(data['avg_prec'])
        iou_thrs.append(iou_thr)

        precisions = data['precisions']
        recalls = data['recalls']
        ax = plot_pr_curve(
            precisions, recalls, label='{:.2f}'.format(iou_thr), color=COLORS[idx*2], ax=ax)

    # prettify for printing:
    avg_precs = [float('{:.4f}'.format(ap)) for ap in avg_precs]
    iou_thrs = [float('{:.4f}'.format(thr)) for thr in iou_thrs]
    print('map: {:.2f}'.format(100*np.mean(avg_precs)))
    print('avg precs: ', avg_precs)
    print('iou_thrs:  ', iou_thrs)
    plt.legend(loc='upper right', title='IOU Thr', frameon=True)
    for xval in np.linspace(0.0, 1.0, 11):
        plt.vlines(xval, 0.0, 1.1, color='gray', alpha=0.3, linestyles='dashed')
    end_time = time.time()
    print('\nPlotting and calculating mAP takes {:.4f} secs'.format(end_time - start_time))
    plt.show()

    # # Plotting per Class:
    # # Source
    # start_time = time.time()
    # ax = None
    # avg_precs = []
    # iou_thrs = []
    # for idx, iou_thr in enumerate(np.linspace(0.5, 0.95, 10)):
    #     data = get_avg_precision_at_iou(gt_boxes_source, pred_boxes_source, iou_thr=iou_thr)
    #     avg_precs.append(data['avg_prec'])
    #     iou_thrs.append(iou_thr)
    #
    #     precisions = data['precisions']
    #     recalls = data['recalls']
    #     ax = plot_pr_curve(
    #         precisions, recalls, label='{:.2f}'.format(iou_thr), color=COLORS[idx * 2], ax=ax)
    #
    # # prettify for printing:
    # avg_precs = [float('{:.4f}'.format(ap)) for ap in avg_precs]
    # iou_thrs = [float('{:.4f}'.format(thr)) for thr in iou_thrs]
    # print('map: {:.2f}'.format(100 * np.mean(avg_precs)))
    # print('avg precs: ', avg_precs)
    # print('iou_thrs:  ', iou_thrs)
    # plt.title('Precision-Recall curve for Source')
    # plt.legend(loc='upper right', title='IOU Thr', frameon=True)
    # for xval in np.linspace(0.0, 1.0, 11):
    #     plt.vlines(xval, 0.0, 1.1, color='gray', alpha=0.3, linestyles='dashed')
    # end_time = time.time()
    # print('\nPlotting and calculating mAP takes {:.4f} secs'.format(end_time - start_time))
    # plt.show()
    #
    # # Sidelobe
    # start_time = time.time()
    # ax = None
    # avg_precs = []
    # iou_thrs = []
    # for idx, iou_thr in enumerate(np.linspace(0.5, 0.95, 10)):
    #     data = get_avg_precision_at_iou(gt_boxes_sidelobe, pred_boxes_sidelobe, iou_thr=iou_thr)
    #     avg_precs.append(data['avg_prec'])
    #     iou_thrs.append(iou_thr)
    #
    #     precisions = data['precisions']
    #     recalls = data['recalls']
    #     ax = plot_pr_curve(
    #         precisions, recalls, label='{:.2f}'.format(iou_thr), color=COLORS[idx * 2], ax=ax)
    #
    # # prettify for printing:
    # avg_precs = [float('{:.4f}'.format(ap)) for ap in avg_precs]
    # iou_thrs = [float('{:.4f}'.format(thr)) for thr in iou_thrs]
    # print('map: {:.2f}'.format(100 * np.mean(avg_precs)))
    # print('avg precs: ', avg_precs)
    # print('iou_thrs:  ', iou_thrs)
    # plt.title('Precision-Recall curve for Sidelobe')
    # plt.legend(loc='upper right', title='IOU Thr', frameon=True)
    # for xval in np.linspace(0.0, 1.0, 11):
    #     plt.vlines(xval, 0.0, 1.1, color='gray', alpha=0.3, linestyles='dashed')
    # end_time = time.time()
    # print('\nPlotting and calculating mAP takes {:.4f} secs'.format(end_time - start_time))
    # plt.show()
    #
    # # Galaxy
    # start_time = time.time()
    # ax = None
    # avg_precs = []
    # iou_thrs = []
    # for idx, iou_thr in enumerate(np.linspace(0.5, 0.95, 10)):
    #     data = get_avg_precision_at_iou(gt_boxes_galaxy, pred_boxes_galaxy, iou_thr=iou_thr)
    #     avg_precs.append(data['avg_prec'])
    #     iou_thrs.append(iou_thr)
    #
    #     precisions = data['precisions']
    #     recalls = data['recalls']
    #     ax = plot_pr_curve(
    #         precisions, recalls, label='{:.2f}'.format(iou_thr), color=COLORS[idx * 2], ax=ax)
    #
    # # prettify for printing:
    # avg_precs = [float('{:.4f}'.format(ap)) for ap in avg_precs]
    # iou_thrs = [float('{:.4f}'.format(thr)) for thr in iou_thrs]
    # print('map: {:.2f}'.format(100 * np.mean(avg_precs)))
    # print('avg precs: ', avg_precs)
    # print('iou_thrs:  ', iou_thrs)
    # plt.title('Precision-Recall curve for Galaxy')
    # plt.legend(loc='upper right', title='IOU Thr', frameon=True)
    # for xval in np.linspace(0.0, 1.0, 11):
    #     plt.vlines(xval, 0.0, 1.1, color='gray', alpha=0.3, linestyles='dashed')
    # end_time = time.time()
    # print('\nPlotting and calculating mAP takes {:.4f} secs'.format(end_time - start_time))
    # plt.show()


    # Super-Imposed per Class
    start_time = time.time()
    ax = None
    avg_precs = []
    iou_thrs = []
    # for idx, iou_thr in enumerate(np.linspace(0.5, 0.95, 10)):
    #     data = get_avg_precision_at_iou(gt_boxes_galaxy, pred_boxes_galaxy, iou_thr=iou_thr)
    #     avg_precs.append(data['avg_prec'])
    #     iou_thrs.append(iou_thr)
    #
    #     precisions = data['precisions']
    #     recalls = data['recalls']
    #     ax = plot_pr_curve(
    #         precisions, recalls, label='{:.2f}'.format(iou_thr), color=COLORS[idx * 2], ax=ax)
    classes = ['source', 'sidelobe', 'galaxy']
    for i, classification in enumerate(classes):
        for iou_thr in np.linspace(0.5, 0.6, 2):
            if classification == 'source':
                data = get_avg_precision_at_iou(gt_boxes_source, pred_boxes_source, iou_thr=iou_thr)
                colour = 'blue'
            elif classification == 'sidelobe':
                data = get_avg_precision_at_iou(gt_boxes_sidelobe, pred_boxes_sidelobe, iou_thr=iou_thr)
                colour = 'red'
            elif classification == 'galaxy':
                data = get_avg_precision_at_iou(gt_boxes_galaxy, pred_boxes_galaxy, iou_thr=iou_thr)
                colour = 'green'
            avg_precs.append(data['avg_prec'])
            iou_thrs.append(iou_thr)

            precisions = data['precisions']
            recalls = data['recalls']

            if iou_thr == 0.5:
                line_style = 'solid'
            elif iou_thr == 0.6:
                line_style = 'dashed'
            # ax = plot_pr_curve(precisions, recalls, label=classification+' @ {:.2f}'.format(iou_thr),
            #                    color=COLORS[i*len(classes) + j], ax=ax)
            ax = plot_pr_curve(precisions, recalls, label=classification + ' @ {:.2f}'.format(iou_thr),
                               color=colour, ax=ax, linestyle=line_style)


    # prettify for printing:
    avg_precs = [float('{:.4f}'.format(ap)) for ap in avg_precs]
    iou_thrs = [float('{:.4f}'.format(thr)) for thr in iou_thrs]
    print('map: {:.2f}'.format(100 * np.mean(avg_precs)))
    print('avg precs: ', avg_precs)
    print('iou_thrs:  ', iou_thrs)
    plt.title('Precision-Recall Curves per Class')
    plt.legend(loc='lower right', title='Class @ IOU Thr', frameon=True)
    for xval in np.linspace(0.0, 1.0, 11):
        plt.vlines(xval, 0.0, 1.1, color='gray', alpha=0.3, linestyles='dashed')
    end_time = time.time()
    print('\nPlotting and calculating mAP takes {:.4f} secs'.format(end_time - start_time))
    plt.show()
