{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3c00934d-0f14-4c1a-a563-94d627b377e5",
   "metadata": {},
   "source": [
    "## Notebook to visualize the viewer functionality\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8501e01d-f932-4a75-8259-2def67a546d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b15dfdc1-5054-40bc-9631-4c7426eecfd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_imports import *\n",
    "\n",
    "import time\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from pyrad.data.dataloader import TrainDataloader\n",
    "from pyrad.data.image_dataset import ImageDataset, PanopticImageDataset\n",
    "from pyrad.data.image_sampler import CacheImageSampler\n",
    "from pyrad.data.pixel_sampler import PixelSampler\n",
    "from pyrad.data.utils import DatasetInputs, get_dataset_inputs_from_dataset_config\n",
    "from pyrad.graphs.modules.ray_generator import RayGenerator\n",
    "from pyrad.graphs.modules.scene_colliders import SceneBoundsCollider, AABBBoxCollider\n",
    "from pyrad.cameras.cameras import get_camera\n",
    "from pyrad.cameras.camera_paths import InterpolatedCameraPath\n",
    "from pyrad.cameras.rays import RayBundle\n",
    "from pyrad.utils.io import get_absolute_path\n",
    "from pyrad.utils.plotly import get_line_segments_from_lines\n",
    "from pyrad.cameras.cameras import get_camera_model\n",
    "from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config\n",
    "from pyrad.engine.trainer import Trainer\n",
    "from pyrad.utils.config import Config, setup_config\n",
    "\n",
    "# the new import\n",
    "from pyrad.viewer.backend import vis_utils\n",
    "\n",
    "\n",
    "from hydra import compose, initialize\n",
    "from omegaconf import open_dict\n",
    "import pprint\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "\n",
    "import umsgpack\n",
    "from pyrad.viewer.backend.utils import get_chunks, get_intrinsics_matrix_and_camera_to_world_h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7fe7478d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with initialize(version_base=None, config_path=\"../configs\"):\n",
    "    config = compose(\n",
    "        config_name=\"graph_instant_ngp.yaml\",\n",
    "        overrides=[\n",
    "            \"trainer.resume_train.load_dir=/shared/ethanweber/pyrad/outputs/blender_lego/instant_ngp/2022-07-05_200513/checkpoints\",\n",
    "            \"trainer.resume_train.load_step=5000\",\n",
    "            \"viewer.render_image_height=50\",\n",
    "        ],\n",
    "    )\n",
    "    config = setup_config(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5d231e36-dd38-489b-a3a3-860cdf49e551",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-05 23:33:06.262342: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(config, local_rank=0, world_size=1)\n",
    "trainer.setup()\n",
    "trainer.draw_scene_in_viewer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "642b0a25-0e88-4eac-b661-a57a74fb10fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_inputs = get_dataset_inputs_from_dataset_config(**config.data.dataset_inputs_train, split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "938f090b-3b28-4fd6-82dc-d4b5c5917134",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ImageDataset\n",
    "image_dataset_train = instantiate_from_dict_config(config.data.image_dataset_train, **dataset_inputs.as_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fd04703a-aa50-45e4-95fd-4d881ff18d61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "estimated_time: 3.3333333333333335\n",
      "time_elapsed: 14.046947002410889\n"
     ]
    }
   ],
   "source": [
    "# move a camera around\n",
    "num_cameras = len(dataset_inputs.intrinsics)\n",
    "intrinsics = dataset_inputs.intrinsics\n",
    "camera_to_world = dataset_inputs.camera_to_world\n",
    "idx0, idx1 = random.sample(range(num_cameras), k=2)\n",
    "camera_a = get_camera(intrinsics[idx0], camera_to_world[idx0], None)\n",
    "camera_b = get_camera(intrinsics[idx1], camera_to_world[idx1], None)\n",
    "\n",
    "num_steps = 100\n",
    "fps = 30\n",
    "estimated_time = num_steps / fps\n",
    "print(\"estimated_time:\", estimated_time)\n",
    "\n",
    "camera_path = InterpolatedCameraPath(camera_a, camera_b)\n",
    "cameras = camera_path.get_path(steps=num_steps)\n",
    "\n",
    "start = time.time()\n",
    "for camera in cameras:\n",
    "    vis_utils.set_camera(trainer.vis, camera)\n",
    "\n",
    "    # render an image\n",
    "    outputs = trainer.render_image_in_viewer()\n",
    "\n",
    "    time.sleep(1 / fps)\n",
    "    # break\n",
    "time_elapsed = time.time() - start\n",
    "print(\"time_elapsed:\", time_elapsed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fbb89f42-bf37-4786-a865-64c4ef018005",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = trainer.render_image_in_viewer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91dfa73-ba0c-4412-a908-74ea51c37c58",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('mattport')",
   "language": "python",
   "name": "python3"
  },
  "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.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "f0938bb18d5c7de3ca7c6a73cb6be5f9b7603b44d115fc3e0726e01396d4d2be"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}