{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "import os\n",
    "os.chdir(\"..\")\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from pathlib import Path\n",
    "from src.utils.plotting import make_matching_figure, error_colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_prediction_and_evaluation_plot(root_dir, pe, path=None, source='ScanNet'):\n",
    "    img0 = cv2.imread(str(root_dir / pe['pair_names'][0]), cv2.IMREAD_GRAYSCALE)\n",
    "    img1 = cv2.imread(str(root_dir / pe['pair_names'][1]), cv2.IMREAD_GRAYSCALE)\n",
    "    if source == 'ScanNet':\n",
    "        img0 = cv2.resize(img0, (640, 480))\n",
    "        img1 = cv2.resize(img1, (640, 480))\n",
    "\n",
    "    thr = 5e-4\n",
    "    mkpts0 = pe['mkpts0_f']\n",
    "    mkpts1 = pe['mkpts1_f']\n",
    "    color = error_colormap(pe['epi_errs'], thr, alpha=0.3)\n",
    "\n",
    "    text = [\n",
    "        f\"LoFTR\",\n",
    "        f\"#Matches: {len(mkpts0)}\",\n",
    "        f\"$\\\\Delta$R:{pe['R_errs']:.2f}°,  $\\\\Delta$t:{pe['t_errs']:.2f}°\",\n",
    "    ]\n",
    "\n",
    "    if path:\n",
    "        make_matching_figure(img0, img1, mkpts0, mkpts1, color, text=text, path=path)\n",
    "    else:\n",
    "        return make_matching_figure(img0, img1, mkpts0, mkpts1, color, text=text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization on ScanNet\n",
    "- Prediction and Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = Path(\"data/scannet/test\")  # Scannet\n",
    "npy_path = \"dump/loftr_ds_indoor/LoFTR_pred_eval.npy\"\n",
    "dumps = np.load(npy_path, allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 420384 characters]",
      "text/plain": [
       "<Figure size 750x450 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = make_prediction_and_evaluation_plot(root_dir, dumps[2], source='ScanNet')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization on MegaDepth\n",
    "- Prediction and Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = Path(\"data/megadepth/test\")  # MegaDepth\n",
    "npy_path = \"dump/loftr_ds_outdoor/LoFTR_pred_eval.npy\"\n",
    "dumps = np.load(npy_path, allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 382132 characters]",
      "text/plain": [
       "<Figure size 750x450 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = make_prediction_and_evaluation_plot(root_dir, dumps[51], source='MegaDepth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.8 64-bit ('loftr': conda)",
   "name": "python388jvsc74a57bd0e2d1507a0fcefcbd70c2e8d5c2edae879585b2c8df0be6cdbe280bf251175c7f"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8-final"
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}