eval_retrieval.py

import pickle as pkl
from os.path import join
import numpy as np
from numpy.linalg import norm
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE, Isomap
from multiprocessing import Pool
import matplotlib.pyplot as plt
from sacred import Experiment
import logging
import sys

from configuration import CONFIG
from src.datasets.a2d2 import id_to_trainid, trainid_to_name
from src.log_utils import log_config

ex = Experiment('eval_retrieval')

log = logging.getLogger()
log.handlers = []

log_format = logging.Formatter('%(asctime)s || %(name)s - [%(levelname)s] - %(message)s')

streamhandler = logging.StreamHandler(sys.stdout)
streamhandler.setFormatter(log_format)
log.addHandler(streamhandler)

log.setLevel('INFO')
ex.logger = log


def lp_dist(point, all_points, d=2):
    return ((all_points - point) ** d).sum(1) ** (1.0 / d)


def cos_dist(point, all_points):
    return 1 - ((point * all_points).sum(1) / (norm(point) * norm(all_points, axis=1)))


def ap_wrapper(arguments):
    return ap(*arguments)


def ap(gt, retrieval_list):
    result = 0
    relevant = retrieval_list == gt
    for k in range(retrieval_list.shape[0]):
        result += (relevant[:(k + 1)].mean() * relevant[k])
    return result / relevant.sum()


def meanaverageprecision(query_indices,
                         gt,
                         embeddings,
                         distance_metric='euclid',
                         gt_annotation=None,
                         void=-1,
                         n_jobs=4):
    if gt_annotation is not None:
        gt[gt_annotation == 0] = void
    if distance_metric == 'euclid':
        args = [(gt[q], gt[np.argsort(lp_dist(embeddings[q], embeddings))][1:]) for q in query_indices]
    elif distance_metric == 'cos':
        args = [(gt[q], gt[np.argsort(cos_dist(embeddings[q], embeddings))][1:]) for q in query_indices]
    with Pool(n_jobs) as p:
        average_precisions = list(p.imap(ap_wrapper, args))

    return sum(average_precisions) / len(average_precisions)


@ex.config
def config():
    args = dict(
        embeddings_file=join(CONFIG.metaseg_io_path, 'embeddings_128_128_densenet201.p'),
        embedding_size=2,
        overwrite_embeddings=False,
        plot_dir=None,
        method='TSNE',
        distance_metric='euclid',
        n_jobs=10,
        min_query_count=2,
    )

    tsne = dict(
        perplexity=30,
        learning_rate=200.0,
        early_exaggeration=12.0,
        verbose=3,
    )


@ex.automain
def main(args, tsne, _run, _log):
    log_config(_run, _log)
    with open(args['embeddings_file'], 'rb') as f:
        data = pkl.load(f)

    gt = np.array(data['gt']).squeeze()
    _log.debug('Number of segments: {}'.format(gt.shape[0]))

    gt = np.vectorize(id_to_trainid.get)(gt)

    if (data['nn_embeddings'].shape[1] != args['embedding_size'] if 'nn_embeddings' in data.keys() else True)\
            or args['overwrite_embeddings']:
        embeddings = np.stack(data['embeddings'])

        # _log.info('Standardizing embeddings...')
        # embeddings = (embeddings - embeddings.mean()) / embeddings.std()

        if args['embedding_size'] < embeddings.shape[1] if args['embedding_size'] is not None else False:
            _log.info('Computing embeddings for nearest neighbor search...')
            if args['method'] == 'TSNE':
                _log.info('Using t-SNE with method \'{}\' and dimensionality {}'.format(
                    'barnes_hut' if args['embedding_size'] < 4 else 'exact',
                    args['embedding_size']
                ))
                embeddings = PCA(
                    n_components=50 if args['embedding_size'] < 50 else 100
                ).fit_transform(embeddings)
                embeddings = TSNE(
                    n_components=args['embedding_size'],
                    n_jobs=args['n_jobs'],
                    method='barnes_hut' if args['embedding_size'] < 4 else 'exact',
                    **tsne
                ).fit_transform(embeddings)
            elif args['method'] == 'Isomap':
                _log.info('Using Isomap method.')
                embeddings = PCA(
                    n_components=50 if args['embedding_size'] < 50 else 100
                ).fit_transform(embeddings)
                embeddings = Isomap(
                    n_components=args['embedding_size'],
                    n_jobs=args['n_jobs'],
                ).fit_transform(embeddings)
            elif args['method'] == 'PCA':
                _log.info('Using PCA method.')
                embeddings = PCA(n_components=args['embedding_size']).fit_transform(embeddings)

            data['nn_embeddings'] = embeddings
            _log.debug('Saving computed manifold to embeddings file.')
            with open(args['embeddings_file'], 'wb') as f:
                pkl.dump(data, f)
        else:
            _log.info('Leaving data as it is.')
    else:
        embeddings = data['nn_embeddings']
        _log.info('Using precomputed embeddings ({} dimensions) for nearest neighbor search...'.format(
            embeddings.shape[1]))

    embeddings = embeddings[gt != 255]
    gt = gt[gt != 255]

    if 'annotated' in data:
        annotated_gt = data['annotated']

    results = {}
    n_queries = {}
    # sel_classes = [12, 22, 3, 34]
    sel_classes = list(range(37))
    for cl in sel_classes:
        query_list = np.argwhere(gt == cl).flatten()
        if 'annotated' in data and query_list.size >= args['min_query_count']:
            query_list = np.array([q for q in query_list if annotated_gt[q] != 0])
        n_queries[cl] = len(query_list) if len(query_list) >= args['min_query_count'] else 0
        if query_list.size >= args['min_query_count']:
            results[cl] = meanaverageprecision(
                query_list,
                gt,
                embeddings,
                distance_metric=args['distance_metric'],
                gt_annotation=annotated_gt if 'annotated' in data else None,
                n_jobs=args['n_jobs'],
            )
            _log.info('{:>{width}s} ({:>4d}): {:>7.2%}'.format(trainid_to_name[cl],
                                                               len(query_list),
                                                               results[cl],
                                                               width=max(
                                                                   [len(str(v)) for v in trainid_to_name.values()])))

    _log.info('Average: {:.2%}'.format(sum(results.values()) / len(results.values())))
    _log.info('Weighted Average: {:.2%}'.format(sum([v * n_queries[k] for k, v in results.items()])
                                                / sum(n_queries.values())))

    if args['plot_dir'] is not None:
        _log.info('Start plotting...')
        fig = plt.figure('mAP values in % for retrieval in the embedding space')
        ax = fig.add_subplot(111)
        rects = ax.bar(x=np.arange(len(results) + 2),
                       height=([v * 100 for k, v in results.items()]
                               + [sum(results.values()) / len(results.values()) * 100]
                               + [sum([v * n_queries[k] for k, v in results.items()]) / sum(n_queries.values()) * 100]))
        ax.set_xticks(np.arange(len(results) + 2))
        ax.set_xticklabels(labels=[trainid_to_name[k] for k in results.keys()] + ['Average'] + ['Weighted Average'])
        for rect in rects:
            height = rect.get_height()
            ax.annotate('{:.1f}'.format(height),
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom')
        # ax.title.set_text('Retrieval results in the embedding space')
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.yaxis.grid(True)
        ax.xaxis.set_tick_params(rotation=50)
        ax.set_ylabel('mAP in %')
        ax.set_axisbelow(True)
        plt.savefig(join(args['plot_dir'], 'map_plot.eps'), dpi=300, bbox_inches='tight')
        _log.info('Saved plot of mAP results to \'{}\''.format(join(args['plot_dir'], 'map_plot.eps')))