# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github"
# <a href="https://colab.research.google.com/github/murthylab/sleap/blob/main/docs/notebooks/Model_evaluation.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="WxEsSjtzA_Af"
# # Model evaluation
#
# After you've trained several models, you may want to generate some accuracy metrics to compare between them. This notebook demonstrates how you'll do that given a trained model.
#
# Let's start by installing SLEAP and downloading the trained model.

# + id="5bNDjxe1BZXV"
# !pip uninstall -y opencv-python opencv-contrib-python > /dev/null 2>&1
# !pip install sleap > /dev/null 2>&1
# !apt install tree > /dev/null 2>&1
# !wget https://storage.googleapis.com/sleap-data/reference/flies13/td_fast.210505_012601.centered_instance.n%3D1800.zip > /dev/null 2>&1
# !unzip -o -d "td_fast.210505_012601.centered_instance.n=1800" "td_fast.210505_012601.centered_instance.n=1800.zip" > /dev/null 2>&1

# + [markdown] id="g6KIMZjwTEmu"
# A trained SLEAP model will be a folder containing files that specify metadata that is useful for evaluation and analysis. The exact set of files may depend on the configuration, but all models will come with:
#
# - `metrics.train.npz`: Metrics for the training split.
# - `metrics.val.npz`: Metrics for the validation split. This is what you'll want to use most of the time since it wasn't directly used for optimizing the model.
#
# **Note:** A test split will also be evaluated if it was provided during training and saved to `metrics.test.npz`.

# + colab={"base_uri": "https://localhost:8080/"} id="dTgeUbf8AT0P" outputId="4f03ee9a-5297-4d53-cb8e-20a78e7c7285"
# !tree td_fast.210505_012601.centered_instance.n=1800

# + [markdown] id="ywbriY3nUnHk"
# Additionally, the following files are included and may also be useful:
# - `best_model.h5`: The actual saved model and weights. This can be loaded with `tf.keras.model.load_model()` but it is recommended to use `sleap.load_model()` instead as it takes care of adding some additional inference-only procedures.
# - `training_config.json`: The configuration for the model training job, including metadata inferred during the training procedure. It can be loaded with `sleap.load_config()`.
# - `labels_gt.train.slp` and `labels_pr.train.slp`: These are SLEAP labels files containing the ground truth and predicted points for the training split. They do not contain the images, but can be used to retrieve the poses used.
# - `labels_gt.train.slp` and `labels_pr.train.slp`: These are SLEAP labels files containing the ground truth and predicted points for the validation split. They do not contain the images, but can be used to retrieve the poses used.

# + colab={"base_uri": "https://localhost:8080/"} id="_FLiavFrBhCK" outputId="fedb9d7b-6dcc-4048-d030-eba38a006086"
import sleap
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns

mpl.style.use("seaborn-deep")
sleap.versions()

# + [markdown] id="r0TLPw8aU6Rh"
# SLEAP metrics can be loaded using the `sleap.load_metrics()` API:

# + colab={"base_uri": "https://localhost:8080/"} id="2o8ncdDYUy94" outputId="6a8a49f0-1252-438b-a2c3-99160fa4bbc4"
help(sleap.load_metrics)

# + [markdown] id="wG1YbA1-SxGN"
# Loading the metrics for the validation split of the model we can see all of the available keys:

# + colab={"base_uri": "https://localhost:8080/"} id="hQHHjfUHMerr" outputId="4f1b2c2a-f86d-47c7-979c-ad70f7a14f61"
metrics = sleap.load_metrics("td_fast.210505_012601.centered_instance.n=1800", split="val")
print("\n".join(metrics.keys()))

# + [markdown] id="ROcOrWX2VB1E"
# To start, let's look at the summary of the **localization errors**:

# + colab={"base_uri": "https://localhost:8080/"} id="EchwGMNcNBoN" outputId="dae62b47-cddb-4605-91c2-c22c899839b6"
print("Error distance (50%):", metrics["dist.p50"])
print("Error distance (90%):", metrics["dist.p90"])
print("Error distance (95%):", metrics["dist.p95"])

# + [markdown] id="o3c91HExVIdx"
# These are the percentiles of the distribution of how far off the model was from the ground truth location.
#
# We can visualize the entire distribution like this:

# + colab={"base_uri": "https://localhost:8080/", "height": 467} id="vo3qyz5nNfgf" outputId="59d0c939-53a3-4580-cf0b-be85b58ad067"
plt.figure(figsize=(6, 3), dpi=150, facecolor="w")
sns.histplot(metrics["dist.dists"].flatten(), binrange=(0, 20), kde=True, kde_kws={"clip": (0, 20)}, stat="probability")
plt.xlabel("Localization error (px)");

# + [markdown] id="LbrwS1iAVStb"
# This metric is intuitive, but it does not incorporate other sources of error like those stemming from poor instance detection and grouping, or missing points.
#
# The Object Keypoint Similarity (OKS) is a more holistic metric that takes factors such as landmark visibility, animal size, and the difficulty in locating keypoints (all are assumed to be "easy" for our calculations). You can read more about this and other pose estimation metrics in: https://arxiv.org/abs/1707.05388
#
# First let's plot the distribution of OKS scores:

# + colab={"base_uri": "https://localhost:8080/", "height": 449} id="8eQ1BHeiOjOx" outputId="820c53c2-7a9e-4430-e8c2-1bccd3254b3d"
plt.figure(figsize=(6, 3), dpi=150, facecolor="w")
sns.histplot(metrics["oks_voc.match_scores"].flatten(), binrange=(0, 1), kde=True, kde_kws={"clip": (0, 1)}, stat="probability")
plt.xlabel("Object Keypoint Similarity");

# + [markdown] id="yj84sWSLWGTl"
# Since these range from 0 to 1, it seems like we're doing pretty well!
#
# Another way to summarize this is through precision-recall curves, which evaluate how well the model does at different thresholds of OKS scores. The higher the threshold, the more stringent our criteria for classifying a prediction as correct.
#
# Here we plot this at different thresholds:

# + colab={"base_uri": "https://localhost:8080/", "height": 562} id="4WEhNB93O7Ns" outputId="56001b09-2f5e-4e4f-cea5-2f66d6d9a2fd"
plt.figure(figsize=(4, 4), dpi=150, facecolor="w")
for precision, thresh in zip(metrics["oks_voc.precisions"][::2], metrics["oks_voc.match_score_thresholds"][::2]):
    plt.plot(metrics["oks_voc.recall_thresholds"], precision, "-", label=f"OKS @ {thresh:.2f}")
plt.xlabel("Recall")
plt.ylabel("Precision")
plt.legend(loc="lower left");

# + [markdown] id="DNgkSGy2WkVW"
# An easy way to summarize this analysis is to take the average over all of these thresholds to compute the **mean Average Precision (mAP)** and **mean Average Recall (mAR)** which are widely used in the pose estimation literature.
#
# Here are those values saved out:

# + colab={"base_uri": "https://localhost:8080/"} id="G3_YhQiFNdu6" outputId="fa4f34ad-00df-48d5-d734-a8cc501e6996"
print("mAP:", metrics["oks_voc.mAP"])
print("mAR:", metrics["oks_voc.mAR"])

# + [markdown] id="bfJ5R2HHXENK"
# Great, but what if we have some new labels or want to evaluate the model with an updated set of labels for better comparisons with newer models?
#
# For this, we'll need to generate new predictions.
#
# First, let's download a new SLEAP labels package file (`.pkg.slp`). This is important since it denotes that the labels contain the images as well which we need for predicting.

# + id="YHCLd3pkRhGT"
# !wget https://storage.googleapis.com/sleap-data/datasets/wt_gold.13pt/tracking_split2/test.pkg.slp > /dev/null 2>&1

# + [markdown] id="9GIPBudQXsDZ"
# Next we can simply load the model, the ground truth (GT) labels, and generate the predictions (~1 min on CPU):

# + id="OMXHY-7YRyTB"
predictor = sleap.load_model("td_fast.210505_012601.centered_instance.n=1800")
labels_gt = sleap.load_file("test.pkg.slp")
labels_pr = predictor.predict(labels_gt)

# + [markdown] id="vUVPXmZSXx-i"
# Generating another set of metrics can then be calculated with the pair of GT and predicted labels:

# + colab={"base_uri": "https://localhost:8080/"} id="0sn0R8tGS4LU" outputId="1de841a5-b885-4b90-f901-ad2913437185"
metrics = sleap.nn.evals.evaluate(labels_gt, labels_pr)

print("Error distance (50%):", metrics["dist.p50"])
print("Error distance (90%):", metrics["dist.p90"])
print("Error distance (95%):", metrics["dist.p95"])
print("mAP:", metrics["oks_voc.mAP"])
print("mAR:", metrics["oks_voc.mAR"])

StarcoderdataJupyterScriptsDedupFiltered