notebooks/visualize_singleview_predictions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import torch\n",
    "from cosypose.config import LOCAL_DATA_DIR\n",
    "from cosypose.datasets.datasets_cfg import make_scene_dataset\n",
    "from cosypose.rendering.bullet_scene_renderer import BulletSceneRenderer\n",
    "from cosypose.visualization.singleview import make_singleview_prediction_plots, filter_predictions\n",
    "from cosypose.visualization.singleview import filter_predictions\n",
    "from bokeh.plotting import gridplot\n",
    "from bokeh.io import show, output_notebook; output_notebook()\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# result_id = 'tless-siso-n_views=1--684390594'\n",
    "# ds_name, urdf_ds_name = 'tless.primesense.test', 'tless.cad'\n",
    "# pred_key = 'pix2pose_detections/refiner/iteration=4'\n",
    "\n",
    "result_id = 'ycbv-n_views=1--5154971130'\n",
    "ds_name, urdf_ds_name = 'ycbv.test.keyframes', 'ycbv'\n",
    "pred_key = 'posecnn_init/refiner/iteration=2'\n",
    "\n",
    "\n",
    "results = LOCAL_DATA_DIR / 'results' / result_id / 'results.pth.tar'\n",
    "scene_ds = make_scene_dataset(ds_name)\n",
    "results = torch.load(results)['predictions']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results[pred_key].infos.loc[:, ['scene_id', 'view_id']].groupby('scene_id').first()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# scene_id, view_id = 12, 103\n",
    "\n",
    "# Replace this here, you can use the dataframe above to get examples of scene/view ids.\n",
    "scene_id, view_id = 48, 733\n",
    "\n",
    "this_preds = filter_predictions(results[pred_key], scene_id, view_id)\n",
    "renderer = BulletSceneRenderer(urdf_ds_name)\n",
    "figures = make_singleview_prediction_plots(scene_ds, renderer, this_preds)\n",
    "renderer.disconnect()\n",
    "# print(this_preds)\n",
    "\n",
    "show(figures['input_im'])\n",
    "show(figures['pred_overlay'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "robotics",
   "language": "python",
   "name": "robotics"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}