[Demo] 2. Image-based Localization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import joblib\n",
    "import quaternion as q\n",
    "import numpy as np\n",
    "from src.models.autoencoder.autoenc import Embedder\n",
    "import torch.nn.functional as F\n",
    "import imageio\n",
    "from IPython.display import clear_output, HTML, display\n",
    "from src.utils.render_utils import add_title, add_agent_view_on_w, get_angle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'\n",
    "embedder = Embedder(pretrained_ckpt='pretrained/autoenc_large.ckpt',\n",
    "                   img_res=128, w_size=128, coordinate_scale=32, w_ch=32, nerf_res=128, voxel_res=128)\n",
    "embedder = embedder.to(device).eval()\n",
    "\n",
    "from src.utils.image_rotator import ImageRotator\n",
    "from src.models.localization.models import UConv\n",
    "localizer = UConv(w_size=128,num_rot=36, w_ch=32, angle_ch=18)\n",
    "rotator = ImageRotator(36)\n",
    "sd = torch.load('pretrained/img_loc.ckpt', map_location='cpu')\n",
    "localizer.load_state_dict(sd)\n",
    "localizer = localizer.cuda().eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Setup Habitat Simulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-09 13:44:18,225 Initializing dataset PointNav-v1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 14 scenes found\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['GLOG_minloglevel'] = \"3\"\n",
    "os.environ['MAGNUM_LOG'] = \"quiet\"\n",
    "os.environ['HABITAT_SIM_LOG'] = \"quiet\"\n",
    "import habitat\n",
    "from habitat import get_config\n",
    "from habitat.sims import make_sim\n",
    "\n",
    "config = get_config()\n",
    "cfg = config\n",
    "cfg.defrost()\n",
    "habitat_api_path = os.path.join(os.path.dirname(habitat.__file__), '../')\n",
    "cfg.DATASET.SCENES_DIR = os.path.join(habitat_api_path, cfg.DATASET.SCENES_DIR)\n",
    "cfg.DATASET.DATA_PATH = os.path.join(habitat_api_path, cfg.DATASET.DATA_PATH.replace(\"habitat-test-scenes\",\"gibson\"))\n",
    "#cfg.SIMULATOR.SCENE_DATASET = 'data/scene_datasets/mp3d.scene_dataset_config.json'\n",
    "#cfg.SIMULATOR.SCENE_DATASET = os.path.join(habitat_api_path, cfg.SIMULATOR.SCENE_DATASET)\n",
    "cfg.DATASET.TYPE = \"PointNav-v1\"\n",
    "cfg.SIMULATOR.RGB_SENSOR.HEIGHT = 128\n",
    "cfg.SIMULATOR.RGB_SENSOR.WIDTH = 128\n",
    "cfg.SIMULATOR.DEPTH_SENSOR.HEIGHT = 128\n",
    "cfg.SIMULATOR.DEPTH_SENSOR.WIDTH = 128\n",
    "cfg.TASK.SENSORS = cfg.SIMULATOR.AGENT_0.SENSORS = ['RGB_SENSOR', 'DEPTH_SENSOR']\n",
    "cfg.freeze()\n",
    "\n",
    "from habitat.datasets import make_dataset\n",
    "dataset = make_dataset('PointNav-v1')\n",
    "cfg.defrost()\n",
    "cfg.DATASET.SPLIT = 'val'\n",
    "cfg.freeze()\n",
    "val_scenes = dataset.get_scenes_to_load(cfg.DATASET)\n",
    "all_scenes =  val_scenes\n",
    "print(f'Total {len(all_scenes)} scenes found')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cantwell\n"
     ]
    }
   ],
   "source": [
    "scene = np.random.choice(all_scenes)\n",
    "print(scene)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-09 13:44:25,524 initializing sim Sim-v0\n"
     ]
    }
   ],
   "source": [
    "try: sim.close()\n",
    "except: pass\n",
    "cfg.defrost()\n",
    "cfg.SIMULATOR.SCENE = os.path.join(cfg.DATASET.SCENES_DIR,'gibson_habitat/{}.glb'.format(scene))\n",
    "cfg.freeze()\n",
    "past_room = scene\n",
    "sim = make_sim(id_sim=cfg.SIMULATOR.TYPE, config=cfg.SIMULATOR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "NUM_GOALS = 5\n",
    "VIS_RES = 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Draw RNR-Map from random trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ad01493624/codes/venv_rnr/lib/python3.8/site-packages/torch/functional.py:478: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2894.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    }
   ],
   "source": [
    "from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower\n",
    "\n",
    "follower = ShortestPathFollower(sim, 0.5, return_one_hot=False)\n",
    "\n",
    "start_position = sim.sample_navigable_point()\n",
    "start_rotation = q.from_euler_angles([0., np.random.rand() * 2 * np.pi, 0.])\n",
    "goal_points = []\n",
    "for i in range(NUM_GOALS):\n",
    "    while True:\n",
    "        random_point = sim.sample_navigable_point()\n",
    "        if abs(random_point[1]-start_position[1]) > 0.5: continue\n",
    "        break\n",
    "    goal_points.append(random_point)\n",
    "\n",
    "\n",
    "w, w_mask = None, None\n",
    "\n",
    "K = torch.eye(3)\n",
    "K[0,0] = (embedder.img_res/2.) / np.tan(np.deg2rad(90.0) / 2)\n",
    "K[1,1] = -(embedder.img_res/2.) / np.tan(np.deg2rad(90.0) / 2)\n",
    "K = K.unsqueeze(0).to(device)\n",
    "\n",
    "orig_Rt = np.eye(4)\n",
    "orig_Rt[:3,3] = start_position\n",
    "orig_Rt[:3,:3] = q.as_rotation_matrix(start_rotation)\n",
    "orig_Rt = np.linalg.inv(orig_Rt)\n",
    "\n",
    "sim.set_agent_state(start_position, start_rotation)\n",
    "done = False\n",
    "goal_point = goal_points.pop(0)\n",
    "rgbs, depths, positions, rotations = [], [], [], []\n",
    "while not done:\n",
    "    state = sim.get_agent_state()\n",
    "    obs = sim.get_observations_at(state.position, state.rotation)\n",
    "    rgbs.append(obs['rgb'])\n",
    "    depths.append(obs['depth'])\n",
    "    positions.append(state.position)\n",
    "    rotations.append(state.rotation.components)\n",
    "\n",
    "\n",
    "    Rt_t = np.eye(4)\n",
    "    Rt_t[:3,3] = state.position\n",
    "    Rt_t[:3,:3] = q.as_rotation_matrix(state.rotation)\n",
    "    Rt_t = np.linalg.inv(Rt_t)\n",
    "    Rt_t = Rt_t @ np.linalg.inv(orig_Rt)\n",
    "\n",
    "    rgb_t = torch.from_numpy(obs['rgb']/255.).unsqueeze(0).permute(0,3,1,2).to(device)\n",
    "    depth_t = torch.from_numpy(obs['depth']).unsqueeze(0).permute(0,3,1,2).to(device)\n",
    "    Rt_t = torch.from_numpy(Rt_t).unsqueeze(0).float().to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = embedder.calculate_mask_func(depth_t*10.0, Rt_t, K)\n",
    "        sorted_indices, seq_unique_list, seq_unique_counts, _ = output\n",
    "        input_dict = {'rgb': rgb_t.unsqueeze(1),\n",
    "                      'depth': depth_t.unsqueeze(1),\n",
    "                    'sorted_indices': sorted_indices.unsqueeze(1),\n",
    "                    'seq_unique_counts': seq_unique_counts.unsqueeze(1),\n",
    "                      'seq_unique_list': seq_unique_list.unsqueeze(1)}\n",
    "        w, w_mask = embedder.embed_obs(input_dict, past_w=w, past_w_num_mask=w_mask)\n",
    "\n",
    "        w_im = w.mean(0).mean(0).detach().cpu().numpy()\n",
    "        w_im = ((w_im - w_im.min())/(w_im.max()-w_im.min()) * 255).astype(np.uint8)\n",
    "        w_im = cv2.applyColorMap(w_im, cv2.COLORMAP_VIRIDIS)[:,:,::-1]\n",
    "        last_w_im = w_im\n",
    "\n",
    "        rgb = add_title(obs['rgb'], 'observation')\n",
    "        w_im = add_agent_view_on_w(w_im, Rt_t, embedder.coordinate_scale, embedder.w_size, agent_size=4, view_size=15)\n",
    "        w_im = np.fliplr(w_im)\n",
    "        w_im = add_title(w_im, 'RNR-Map')\n",
    "\n",
    "    cv2.imshow(\"View\", np.concatenate([rgb, w_im],1)[:,:,::-1])\n",
    "    cv2.waitKey(1)\n",
    "\n",
    "    action = follower.get_next_action(goal_point)\n",
    "    if action is None or action == 0:\n",
    "        if len(goal_points) == 0:\n",
    "            break\n",
    "        goal_point = goal_points.pop(0)\n",
    "        action = follower.get_next_action(goal_point)\n",
    "    sim.step(action)\n",
    "\n",
    "data = {'rgb': rgbs, 'depth': depths, 'position': positions, 'rotation': rotations, 'map': w, 'orig_Rt': orig_Rt}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Localize seen images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "average error: 0.173m\n"
     ]
    }
   ],
   "source": [
    "origin = torch.eye(4).unsqueeze(0).to(device)\n",
    "images = []\n",
    "diffs = []\n",
    "\n",
    "coordinate_scale = embedder.coordinate_scale\n",
    "map_size = embedder.w_size\n",
    "patch_size = map_size//4\n",
    "angle_bin = 18\n",
    "for t in range(len(rgbs)):\n",
    "    # make target RNR-Map\n",
    "    rgb = torch.from_numpy(data['rgb'][t]).unsqueeze(0).permute(0,3,1,2).to(device)\n",
    "    depth = torch.from_numpy(data['depth'][t]).unsqueeze(0).permute(0,3,1,2).to(device)\n",
    "    sorted_indices, seq_unique_list, seq_unique_counts, pose_map = embedder.calculate_mask_func(depth * 10.0, origin, K)\n",
    "\n",
    "    sample_dict = {'sorted_indices': sorted_indices.unsqueeze(0),\n",
    "               'seq_unique_list': seq_unique_list.unsqueeze(0),\n",
    "               'seq_unique_counts': seq_unique_counts.unsqueeze(0),\n",
    "              'rgb': rgb.unsqueeze(1)/255., 'depth': depth.unsqueeze(1)}\n",
    "    for k,v in sample_dict.items():\n",
    "        sample_dict[k] = v.cuda()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        latent_target, _ = embedder.embed_obs(sample_dict)\n",
    "        latent_target = latent_target[:, :, map_size//2 - patch_size//2 : map_size//2 + patch_size//2,\n",
    "                                            map_size//2 - patch_size//2 : map_size//2 + patch_size//2]\n",
    "\n",
    "        # Localize\n",
    "        pred_heatmap, pred_angle = localizer(data['map'], latent_target, rotator)\n",
    "        seen_area = (data['map'].mean(dim=1) != 0)\n",
    "        bs, ws, hs = torch.where(seen_area == 0)\n",
    "        pred_heatmap[bs, :, ws, hs] = -99999\n",
    "        pred_heatmap[bs, :, :, -1] = -99999\n",
    "        pred_heatmap[bs, :, -1, :] = -99999\n",
    "        pred_heatmap_flattened = F.softmax(pred_heatmap.view(1, -1), dim=-1)\n",
    "        pred = pred_heatmap_flattened.view(map_size+1, map_size+1)\n",
    "\n",
    "    pred_max = pred_heatmap.view(1, -1).argmax(dim=1).item()\n",
    "    pred_h, pred_w = pred_max//pred_heatmap.shape[-1], pred_max%pred_heatmap.shape[-1]\n",
    "    pred_x = (pred_h-(map_size//2))/(map_size//2) * (coordinate_scale/2)\n",
    "    pred_y = (pred_w-(map_size//2))/(map_size//2) * (coordinate_scale/2)\n",
    "\n",
    "    pred_Rt = np.eye(4)\n",
    "    pred_Rt[:3,3] = np.array([pred_x, 0., pred_y])\n",
    "    pred_Rt[:3,:3] = q.as_rotation_matrix(q.from_euler_angles([0., 2*np.pi/angle_bin * pred_angle.argmax().item(), 0.0]))\n",
    "    pred_Rt = np.linalg.inv(pred_Rt)\n",
    "    pred_sim_Rt = np.linalg.inv(np.matmul(pred_Rt, data['orig_Rt']))\n",
    "\n",
    "    # Get observation from predicted pose\n",
    "    pred_position_ = pred_sim_Rt[:3,3]\n",
    "    pred_rotation = q.from_rotation_matrix(pred_sim_Rt[:3,:3])\n",
    "    pred_obs = sim.get_observations_at(pred_position_, pred_rotation)\n",
    "\n",
    "    # Calculate Localization error\n",
    "    Rtt = np.eye(4)\n",
    "    Rtt[:3,3] = data['position'][t]\n",
    "    Rtt[:3,:3] = q.as_rotation_matrix(q.from_float_array(data['rotation'][t]))\n",
    "    Rtt = np.linalg.inv(Rtt)@np.linalg.inv(data['orig_Rt'])\n",
    "    answer_x, _, answer_y = np.linalg.inv(Rtt)[:3,3]\n",
    "    answer_h = int(answer_x*(map_size/2)/(coordinate_scale/2.)+(map_size/2))\n",
    "    answer_w = int(answer_y*(map_size/2)/(coordinate_scale/2.)+(map_size/2))\n",
    "\n",
    "    \n",
    "    diff = np.linalg.norm(data['position'][t][:2]-pred_position_[:2])\n",
    "    diffs.append(diff)\n",
    "\n",
    "    # Visualization\n",
    "    map_im = last_w_im.copy()\n",
    "    map_im = add_agent_view_on_w(map_im, Rtt, embedder.coordinate_scale, embedder.w_size, agent_size=4, view_size=15, agent_color=(255,0,0), view_color=(255,0,0))\n",
    "    map_im = add_agent_view_on_w(map_im, pred_Rt, embedder.coordinate_scale, embedder.w_size, agent_size=4, view_size=15, agent_color=(0,0,255), view_color=(0,0,255))\n",
    "    map_im = cv2.resize(map_im, dsize=(VIS_RES, VIS_RES))\n",
    "\n",
    "    pred_im = pred[:-1,:-1].detach().cpu().numpy()\n",
    "    pred_im = (pred_im - pred_im.min())/(pred_im.max()-pred_im.min())\n",
    "    pred_im = cv2.resize((pred_im*255).astype(np.uint8), dsize=(VIS_RES, VIS_RES))\n",
    "    pred_im = cv2.applyColorMap(pred_im, cv2.COLORMAP_VIRIDIS)[:,:,::-1]\n",
    "    pred_im = cv2.addWeighted(cv2.resize(last_w_im, dsize=(VIS_RES, VIS_RES)), 0.3, pred_im, 0.7, 0.0)\n",
    "\n",
    "    map_im = add_title(map_im, 'RNR-Map')\n",
    "    pred_im = add_title(pred_im, 'Loc. Heatmap')\n",
    "\n",
    "    rgb = cv2.resize(data['rgb'][t], dsize=(VIS_RES, VIS_RES))\n",
    "    rgb = add_title(rgb, 'Query Img.')\n",
    "    pred_rgb = cv2.resize(pred_obs['rgb'], dsize=(VIS_RES, VIS_RES))\n",
    "    pred_rgb = add_title(pred_rgb, 'Localized')\n",
    "\n",
    "    view_im = np.concatenate([rgb, pred_rgb,map_im, pred_im],1)\n",
    "    images.append(view_im)\n",
    "    cv2.imshow(\"view\", view_im[:,:,::-1])\n",
    "    key = cv2.waitKey(1)\n",
    "    if key == ord(\"q\"): break\n",
    "\n",
    "print(f'average error: {np.mean(diffs):.3f}m' )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=demo/image-based-localization.gif>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "imageio.mimwrite(\"demo/image-based-localization.gif\", images, fps=15)\n",
    "display(HTML('<img src={}>'.format(\"demo/image-based-localization.gif\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Localize images from unseen view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "average error: 0.481m\n"
     ]
    }
   ],
   "source": [
    "origin = torch.eye(4).unsqueeze(0).to(device)\n",
    "images = []\n",
    "diffs = []\n",
    "\n",
    "coordinate_scale = embedder.coordinate_scale\n",
    "map_size = embedder.w_size\n",
    "patch_size = map_size//4\n",
    "angle_bin = 18\n",
    "\n",
    "noise_dist = 2\n",
    "noise_angle = np.pi/3\n",
    "for t in range(len(rgbs)):\n",
    "    while True:\n",
    "        query_position = data['position'][t] + np.array([np.random.rand()*noise_dist-noise_dist/2,0,\n",
    "                                                        np.random.rand()*noise_dist-noise_dist/2])\n",
    "        if not sim.is_navigable(query_position): continue\n",
    "        query_angle = get_angle(q.from_float_array(data['rotation'][t])) + np.random.rand()*noise_angle - noise_angle/2.\n",
    "        query_rotation = q.from_euler_angles([0, query_angle, 0])\n",
    "        query_obs = sim.get_observations_at(query_position, query_rotation)\n",
    "        break\n",
    "    \n",
    "    # make target RNR-Map\n",
    "    rgb = torch.from_numpy(query_obs['rgb']).unsqueeze(0).permute(0,3,1,2).to(device)\n",
    "    depth = torch.from_numpy(query_obs['depth']).unsqueeze(0).permute(0,3,1,2).to(device)\n",
    "    sorted_indices, seq_unique_list, seq_unique_counts, pose_map = embedder.calculate_mask_func(depth * 10.0, origin, K)\n",
    "\n",
    "    sample_dict = {'sorted_indices': sorted_indices.unsqueeze(0),\n",
    "               'seq_unique_list': seq_unique_list.unsqueeze(0),\n",
    "               'seq_unique_counts': seq_unique_counts.unsqueeze(0),\n",
    "              'rgb': rgb.unsqueeze(1)/255., 'depth': depth.unsqueeze(1)}\n",
    "    for k,v in sample_dict.items():\n",
    "        sample_dict[k] = v.cuda()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        latent_target, _ = embedder.embed_obs(sample_dict)\n",
    "        latent_target = latent_target[:, :, map_size//2 - patch_size//2 : map_size//2 + patch_size//2,\n",
    "                                            map_size//2 - patch_size//2 : map_size//2 + patch_size//2]\n",
    "\n",
    "        # Localize\n",
    "        pred_heatmap, pred_angle = localizer(data['map'], latent_target, rotator)\n",
    "        seen_area = (data['map'].mean(dim=1) != 0)\n",
    "        bs, ws, hs = torch.where(seen_area == 0)\n",
    "        pred_heatmap[bs, :, ws, hs] = -99999\n",
    "        pred_heatmap[bs, :, :, -1] = -99999\n",
    "        pred_heatmap[bs, :, -1, :] = -99999\n",
    "        pred_heatmap_flattened = F.softmax(pred_heatmap.view(1, -1), dim=-1)\n",
    "        pred = pred_heatmap_flattened.view(map_size+1, map_size+1)\n",
    "\n",
    "    pred_max = pred_heatmap.view(1, -1).argmax(dim=1).item()\n",
    "    pred_h, pred_w = pred_max//pred_heatmap.shape[-1], pred_max%pred_heatmap.shape[-1]\n",
    "    pred_x = (pred_h-(map_size//2))/(map_size//2) * (coordinate_scale/2)\n",
    "    pred_y = (pred_w-(map_size//2))/(map_size//2) * (coordinate_scale/2)\n",
    "\n",
    "    pred_Rt = np.eye(4)\n",
    "    pred_Rt[:3,3] = np.array([pred_x, 0., pred_y])\n",
    "    pred_Rt[:3,:3] = q.as_rotation_matrix(q.from_euler_angles([0., 2*np.pi/angle_bin * pred_angle.argmax().item(), 0.0]))\n",
    "    pred_Rt = np.linalg.inv(pred_Rt)\n",
    "    pred_sim_Rt = np.linalg.inv(np.matmul(pred_Rt, data['orig_Rt']))\n",
    "\n",
    "    # Get observation from predicted pose\n",
    "    pred_position_ = pred_sim_Rt[:3,3]\n",
    "    pred_rotation = q.from_rotation_matrix(pred_sim_Rt[:3,:3])\n",
    "    pred_obs = sim.get_observations_at(pred_position_, pred_rotation)\n",
    "\n",
    "    # Calculate Localization error\n",
    "    Rtt = np.eye(4)\n",
    "    Rtt[:3,3] = query_position\n",
    "    Rtt[:3,:3] = q.as_rotation_matrix(query_rotation)\n",
    "    Rtt = np.linalg.inv(Rtt)@np.linalg.inv(data['orig_Rt'])\n",
    "    answer_x, _, answer_y = np.linalg.inv(Rtt)[:3,3]\n",
    "    answer_h = int(answer_x*(map_size/2)/(coordinate_scale/2.)+(map_size/2))\n",
    "    answer_w = int(answer_y*(map_size/2)/(coordinate_scale/2.)+(map_size/2))\n",
    "\n",
    "    \n",
    "    diff = np.linalg.norm(query_position[:2]-pred_position_[:2])\n",
    "    diffs.append(diff)\n",
    "\n",
    "    # Visualization\n",
    "    map_im = last_w_im.copy()\n",
    "    map_im = add_agent_view_on_w(map_im, Rtt, embedder.coordinate_scale, embedder.w_size, agent_size=4, view_size=15, agent_color=(255,0,0), view_color=(255,0,0))\n",
    "    map_im = add_agent_view_on_w(map_im, pred_Rt, embedder.coordinate_scale, embedder.w_size, agent_size=4, view_size=15, agent_color=(0,0,255), view_color=(0,0,255))\n",
    "    map_im = cv2.resize(map_im, dsize=(VIS_RES, VIS_RES))\n",
    "\n",
    "    pred_im = pred[:-1,:-1].detach().cpu().numpy()\n",
    "    pred_im = (pred_im - pred_im.min())/(pred_im.max()-pred_im.min())\n",
    "    pred_im = cv2.resize((pred_im*255).astype(np.uint8), dsize=(VIS_RES, VIS_RES))\n",
    "    pred_im = cv2.applyColorMap(pred_im, cv2.COLORMAP_VIRIDIS)[:,:,::-1]\n",
    "    pred_im = cv2.addWeighted(cv2.resize(last_w_im, dsize=(VIS_RES, VIS_RES)), 0.3, pred_im, 0.7, 0.0)\n",
    "\n",
    "    map_im = add_title(map_im, 'RNR-Map')\n",
    "    pred_im = add_title(pred_im, 'Loc. Heatmap')\n",
    "\n",
    "    rgb = cv2.resize(query_obs['rgb'], dsize=(VIS_RES, VIS_RES))\n",
    "    rgb = add_title(rgb, 'Query Img.')\n",
    "    pred_rgb = cv2.resize(pred_obs['rgb'], dsize=(VIS_RES, VIS_RES))\n",
    "    pred_rgb = add_title(pred_rgb, 'Localized')\n",
    "\n",
    "    view_im = np.concatenate([rgb, pred_rgb,map_im, pred_im],1)\n",
    "    images.append(view_im)\n",
    "    cv2.imshow(\"view\", view_im[:,:,::-1])\n",
    "    key = cv2.waitKey(1)\n",
    "    if key == ord(\"q\"): break\n",
    "\n",
    "print(f'average error: {np.mean(diffs):.3f}m' )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=demo/image-based-localization-noise.gif>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "imageio.mimwrite(\"demo/image-based-localization-noise.gif\", images, fps=5)\n",
    "display(HTML('<img src={}>'.format(\"demo/image-based-localization-noise.gif\")))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rnr",
   "language": "python",
   "name": "rnr"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}