"""
# Baseline model evaluation
"""
"""
In order to make sure that model is learning correctly, it is useful to frequently run evaluation loops and keep track whether statistic of interest is being improved. In case of the problem of Landmark Retrieval, the statistic often used is Mean Average Precision (mAP). It is exactly the statistic used for model evaluation in this competiton. It is defined, in [contest evaluation protocol](https://www.kaggle.com/c/landmark-retrieval-2020/overview/evaluation), as:
"""
"""
![mAP](https://i.imgur.com/ukUp5cy.png)
"""
"""
There are a lot of great blogposts exaplaining this metric in-depth, for example [this one](https://towardsdatascience.com/breaking-down-mean-average-precision-map-ae462f623a52), but the simple thing to know is this: hidden between all these obscure math symbols is a recipe for coming up with a single number that summarizes our model in a way that if it provides embeddings (model output) that are useful for image retrieval, the mAP score gets bigger. In particular, for each query image, embeddings for all index images are compared with the embedding for the query image (using Euclidean distance between embedding vectors), and sorted in a way, such that index images with most similar embedding vectors are placed first. The metric used in this contest takes 100 most similar images and gets bigger, if ones corresponding to the landmark presented in query image are listed first.
"""
"""
Since the dataset used in the competition is very large (>1.5M images), it is not practical to use them all for evaluation - inference on such an enormous dataset takes a significant amount of time. Moreover, it is hard to pick a specific small subset of landmarks to use in evaluation - the dataset is very diverse and it is a challenge to come up with representative set of images.
"""
"""
In order to address this problem, smaller datasets are used in literature as well. One of those datasets are Revisited Oxford and Paris datasets, presented originally in [this paper](https://arxiv.org/abs/1803.11285), as a relabelling of popular [Oxford Buildings Dataset](https://www.robots.ox.ac.uk/~vgg/data/oxbuildings/) and [The Paris Dataset](https://www.robots.ox.ac.uk/~vgg/data/parisbuildings/). The relabeling was a considerate effort, which adressed noisiness present in original datasets. That means that the dataset can be treated with high confidence as a representative one. Moreover, the mAP metrics on these datasets are often reported in majority of recent Image Retrieval literature.
"""
"""
In this notebook, I present how to evaluate baseline model provided by Google on Revisited Oxford & Paris datasets. The evaluation loops can be easily re-used for your custom model trained for this competition. The code is based on https://github.com/filipradenovic/revisitop
"""
"""
### 1. Define utility functions
"""
import os
import pickle
import random

import numpy as np
import tensorflow as tf
from tqdm  import tqdm
from PIL import Image, ImageFile
from scipy.io import savemat, loadmat
import matplotlib.pyplot as plt
# Utility script - parsing dataset config .pkl files and mAP computation
from roxfordparis_tools import configdataset, compute_map
# Adapted from: https://github.com/filipradenovic/revisitop/blob/master/python/example_process_images.py

def pil_loader(path):
    # to avoid crashing for truncated (corrupted images)
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    # open path as file to avoid ResourceWarning 
    # (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')
# Adapted from: https://github.com/filipradenovic/revisitop/blob/master/python/example_process_images.py

def extract_features(test_dataset, cfg, model):
    """
    Generates file with serialized model outputs for each image from test_dataset.
    
    Arguments
    ---------
    test_dataset   : name of dataset of interest (roxford5k | rparis6k)
    cfg            : unserialized dataset config, containing annotation metadata
    model          : loaded Tensorflow baseline model object
    """
    
    print('>> Processing query images...', flush=True)
    Q = []
    for i in tqdm(np.arange(cfg['nq'])):
        qim = pil_loader(cfg['qim_fname'](cfg, i)).crop(cfg['gnd'][i]['bbx'])
        image_data = np.array(qim)
        image_tensor = tf.convert_to_tensor(image_data)
        Q.append(model(image_tensor)['global_descriptor'].numpy())
    Q = np.array(Q, dtype=np.float32)
    Q = Q.transpose()
    
    print('>> Processing index images...', flush=True)
    X = []
    for i in tqdm(np.arange(cfg['n'])):
        im = pil_loader(cfg['im_fname'](cfg, i))
        image_data = np.array(im)
        image_tensor = tf.convert_to_tensor(image_data)
        X.append(model(image_tensor)['global_descriptor'].numpy())
    X = np.array(X, dtype=np.float32)
    X = X.transpose()

    feature_dict = {'X': X, 'Q': Q}
    mat_save_path = "{}_delg_baseline.mat".format(test_dataset)
    print('>> Saving model outputs to: {}'.format(mat_save_path))
    savemat(mat_save_path, feature_dict)
# Adapted from: https://github.com/filipradenovic/revisitop/blob/master/python/example_evaluate.py

def run_evaluation(test_dataset, cfg, features_dir):
    ks = [1, 5, 10]
    gnd = cfg['gnd']
    features = loadmat(os.path.join(features_dir, '{}_delg_baseline.mat'.format(test_dataset)))

    Q = features['Q']
    X = features['X']
    sim = np.dot(X.T, Q)
    ranks = np.argsort(-sim, axis=0)

    # search for easy
    gnd_t = []
    for i in range(len(gnd)):
        g = {}
        g['ok'] = np.concatenate([gnd[i]['easy']])
        g['junk'] = np.concatenate([gnd[i]['junk'], gnd[i]['hard']])
        gnd_t.append(g)
    mapE, apsE, mprE, prsE = compute_map(ranks, gnd_t, ks)

    # search for easy & hard
    gnd_t = []
    for i in range(len(gnd)):
        g = {}
        g['ok'] = np.concatenate([gnd[i]['easy'], gnd[i]['hard']])
        g['junk'] = np.concatenate([gnd[i]['junk']])
        gnd_t.append(g)
    mapM, apsM, mprM, prsM = compute_map(ranks, gnd_t, ks)

    # search for hard
    gnd_t = []
    for i in range(len(gnd)):
        g = {}
        g['ok'] = np.concatenate([gnd[i]['hard']])
        g['junk'] = np.concatenate([gnd[i]['junk'], gnd[i]['easy']])
        gnd_t.append(g)
    mapH, apsH, mprH, prsH = compute_map(ranks, gnd_t, ks)

    print('>> {}: mAP E: {}, M: {}, H: {}'.format(test_dataset, np.around(mapE*100, decimals=2), np.around(mapM*100, decimals=2), np.around(mapH*100, decimals=2)))
    print('>> {}: mP@k{} E: {}, M: {}, H: {}'.format(test_dataset, np.array(ks), np.around(mprE*100, decimals=2), np.around(mprM*100, decimals=2), np.around(mprH*100, decimals=2)))
"""
### 2. Feature extraction
"""
data_root = '/kaggle/input/roxfordparis'
model_root = '/kaggle/input/baseline-landmark-retrieval-model/baseline_landmark_retrieval_model'
model = tf.saved_model.load(model_root).signatures['serving_default']
"""
This notebook uses cached feature extractions from `delg-baseline-roxfordparis-output` Kaggle dataset. Those features can be generated in `/kaggle/working` directory as well, by changing `should_extract_features` flag value to `True`. Be vary that feature extraction takes around ~2h with GPU accelerator on.
"""
# Extract features for roxford5k & rparis6k datasets

should_extract_features = False

datasets = ['roxford5k', 'rparis6k']
for test_dataset in datasets:
    if should_extract_features:
        print('Processing dataset: {}'.format(test_dataset), flush=True)
        cfg = configdataset(test_dataset, data_root)
        extract_features(test_dataset, cfg, model)
"""
### 3. Model evaluation
"""
# Evaluate on roxford5k & rparis6k datasets

datasets = ['roxford5k', 'rparis6k']
for test_dataset in datasets:
    print('Evaluating on dataset: {}'.format(test_dataset), flush=True)
    cfg = configdataset(test_dataset, data_root)
    run_evaluation(test_dataset, cfg, '/kaggle/input/delg-baseline-roxfordparis-output')
    print()
"""
The final mAP metrics are:
- roxford5k - Easy: 90.92, Medium: 76.23, Hard: 55.54
- rparis5k - Easy: 94.06, Medium: 87.25, Hard: 74.2

It's interesting to noticed that these metrics are higher than ones presented in [original DELG paper](https://arxiv.org/pdf/2001.05027.pdf) the baseline model is based on (check rows corresponding to DELG in columns marked as `ROxf` and `RPar`). That suggests that the baseline model is more powerful than the one presented in the paper, which means it is a considerate challenge to come up with model that improves baseline score on the Leaderboard.
"""
"""
![DELG figure](https://i.imgur.com/IMQZei7.png)
"""
"""
### BONUS: Testing model with concrete input
"""
"""
The following code cells perform image retrieval of 5 most similar images using baseline model.
"""
def get_random_image(test_dataset, cfg):
    query_image_id = random.choice(range(cfg['nq']))
    im = pil_loader(cfg['im_fname'](cfg, query_image_id))
    return query_image_id, im
def run_inference(pil_image, model):
    image_data = np.array(pil_image)
    image_tensor = tf.convert_to_tensor(image_data)
    return model(image_tensor)['global_descriptor'].numpy()
test_dataset = 'roxford5k' # (roxford5k | rparis6k)
cfg = configdataset(test_dataset, data_root)

im_id, im = get_random_image(test_dataset, cfg)
im_feats = run_inference(im, model)
features_dir = '/kaggle/input/delg-baseline-roxfordparis-output'

features = loadmat(os.path.join(features_dir, '{}_delg_baseline.mat'.format(test_dataset)))
X = features['X']
sim = np.dot(X.T, im_feats)
ranks = np.argsort(-sim, axis=0)
plt.imshow(np.array(im))
plt.title('Query Image')
plt.axis('off')
fig=plt.figure(figsize=(16, 12))
columns = 5
rows = 1
for i in range(1, columns*rows +1):
    img = np.array(pil_loader(cfg['im_fname'](cfg, ranks[i-1])))
    fig.add_subplot(rows, columns, i)
    plt.imshow(img)
    plt.axis('off')
print('Five most similar images from query dataset:')