{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Set CUDA:2 as current device.\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.transforms.functional as trans_f\n",
    "\n",
    "sys.path.append(os.path.abspath(sys.path[0] + '/../../'))\n",
    "__package__ = \"deep_view_syn.notebook\"\n",
    "torch.cuda.set_device(2)\n",
    "print(\"Set CUDA:%d as current device.\" % torch.cuda.current_device())\n",
    "\n",
    "from ..data.spherical_view_syn import *\n",
    "from ..configs.spherical_view_syn import SphericalViewSynConfig\n",
    "from ..my import netio\n",
    "from ..my import util\n",
    "from ..my import device\n",
    "from ..my import view\n",
    "from ..my.gen_final import GenFinal\n",
    "from ..my.simple_perf import SimplePerf\n",
    "\n",
    "\n",
    "def load_net(path):\n",
    "    config = SphericalViewSynConfig()\n",
    "    config.from_id(path[:-4])\n",
    "    config.SAMPLE_PARAMS['perturb_sample'] = False\n",
    "    config.print()\n",
    "    net = config.create_net().to(device.GetDevice())\n",
    "    netio.LoadNet(path, net)\n",
    "    return net\n",
    "\n",
    "\n",
    "def find_file(prefix):\n",
    "    for path in os.listdir():\n",
    "        if path.startswith(prefix):\n",
    "            return path\n",
    "    return None\n",
    "\n",
    "\n",
    "def load_views(data_desc_file) -> view.Trans:\n",
    "    with open(data_desc_file, 'r', encoding='utf-8') as file:\n",
    "        data_desc = json.loads(file.read())\n",
    "        view_centers = torch.tensor(\n",
    "            data_desc['view_centers'], device=device.GetDevice()).view(-1, 3)\n",
    "        view_rots = torch.tensor(\n",
    "            data_desc['view_rots'], device=device.GetDevice()).view(-1, 3, 3)\n",
    "        return view.Trans(view_centers, view_rots)\n",
    "\n",
    "\n",
    "def plot_figures(images, center):\n",
    "    plt.figure(figsize=(8, 4))\n",
    "    plt.subplot(121)\n",
    "    util.PlotImageTensor(images['fovea_raw'])\n",
    "    plt.subplot(122)\n",
    "    util.PlotImageTensor(images['fovea'])\n",
    "\n",
    "    plt.figure(figsize=(8, 4))\n",
    "    plt.subplot(121)\n",
    "    util.PlotImageTensor(images['mid_raw'])\n",
    "    plt.subplot(122)\n",
    "    util.PlotImageTensor(images['mid'])\n",
    "\n",
    "    plt.figure(figsize=(8, 4))\n",
    "    plt.subplot(121)\n",
    "    util.PlotImageTensor(images['periph_raw'])\n",
    "    plt.subplot(122)\n",
    "    util.PlotImageTensor(images['periph'])\n",
    "\n",
    "    # Plot Blended\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(121)\n",
    "    util.PlotImageTensor(images['blended_raw'])\n",
    "    plt.subplot(122)\n",
    "    util.PlotImageTensor(images['blended'])\n",
    "    plt.plot([(res_full[1] - 1) / 2 + center[0] - 5, (res_full[1] - 1) / 2 + center[0] + 5],\n",
    "                [(res_full[0] - 1) / 2 + center[1],\n",
    "                (res_full[0] - 1) / 2 + center[1]],\n",
    "                color=[0, 1, 0])\n",
    "    plt.plot([(res_full[1] - 1) / 2 + center[0], (res_full[1] - 1) / 2 + center[0]],\n",
    "                [(res_full[0] - 1) / 2 + center[1] - 5,\n",
    "                (res_full[0] - 1) / 2 + center[1] + 5],\n",
    "                color=[0, 1, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Change working directory to  /home/dengnc/deep_view_syn/data/__0_user_study/us_gas_all_in_one\n",
      "==== Config fovea ====\n",
      "Net type:  nmsl\n",
      "Encode dim:  10\n",
      "Optimizer decay:  0\n",
      "Normalize:  False\n",
      "Direction as input:  False\n",
      "Full-connected network parameters: {'nf': 128, 'n_layers': 4, 'skips': []}\n",
      "Sample parameters {'spherical': True, 'depth_range': (1.0, 50.0), 'n_samples': 32, 'perturb_sample': False, 'lindisp': True, 'inverse_r': True}\n",
      "==========================\n",
      "Load net from fovea@nmsl-rgb_e10_fc128x4_d1-50_s32.pth ...\n",
      "==== Config periph ====\n",
      "Net type:  nnmsl\n",
      "Encode dim:  10\n",
      "Optimizer decay:  0\n",
      "Normalize:  False\n",
      "Direction as input:  False\n",
      "Full-connected network parameters: {'nf': 64, 'n_layers': 4, 'skips': []}\n",
      "Sample parameters {'spherical': True, 'depth_range': (1.0, 50.0), 'n_samples': 16, 'perturb_sample': False, 'lindisp': True, 'inverse_r': True}\n",
      "==========================\n",
      "Load net from periph@nnmsl-rgb_e10_fc64x4_d1-50_s16.pth ...\n",
      "Dataset loaded.\n",
      "views: [110]\n"
     ]
    }
   ],
   "source": [
    "os.chdir(sys.path[0] + '/../data/__0_user_study/us_gas_all_in_one')\n",
    "#os.chdir(sys.path[0] + '/../data/__0_user_study/us_mc_all_in_one')\n",
    "#os.chdir(sys.path[0] + '/../data/lobby_all_in_one')\n",
    "print('Change working directory to ', os.getcwd())\n",
    "torch.autograd.set_grad_enabled(False)\n",
    "\n",
    "fovea_net = load_net(find_file('fovea'))\n",
    "periph_net = load_net(find_file('periph'))\n",
    "\n",
    "# Load Dataset\n",
    "views = load_views('nerf_views.json')\n",
    "print('Dataset loaded.')\n",
    "\n",
    "print('views:', views.size())\n",
    "#print('ref views:', ref_dataset.samples)\n",
    "\n",
    "fov_list = [20, 45, 110]\n",
    "res_list = [(128, 128), (256, 256), (256, 230)]  # (192,256)]\n",
    "res_full = (1600, 1440)\n",
    "gen = GenFinal(fov_list, res_list, res_full, fovea_net, periph_net,\n",
    "               device=device.GetDevice())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GetRays: 580.4ms\n",
      "Sample: 2.1ms\n",
      "Encode: 2.2ms\n",
      "Rays: tensor([[-0.1711,  0.1711,  0.9702],\n",
      "        [-0.1685,  0.1711,  0.9707],\n",
      "        [-0.1659,  0.1713,  0.9712],\n",
      "        ...,\n",
      "        [ 0.1633, -0.1687,  0.9722],\n",
      "        [ 0.1660, -0.1687,  0.9717],\n",
      "        [ 0.1686, -0.1686,  0.9712]], device='cuda:2')\n",
      "Spherical coords: tensor([[[1.0000, 1.7454, 1.3988],\n",
      "         [0.9684, 1.7454, 1.3988],\n",
      "         [0.9368, 1.7454, 1.3988],\n",
      "         ...,\n",
      "         [0.0832, 1.7454, 1.3988],\n",
      "         [0.0516, 1.7454, 1.3988],\n",
      "         [0.0200, 1.7454, 1.3988]],\n",
      "\n",
      "        [[1.0000, 1.7426, 1.3988],\n",
      "         [0.9684, 1.7426, 1.3988],\n",
      "         [0.9368, 1.7426, 1.3988],\n",
      "         ...,\n",
      "         [0.0832, 1.7426, 1.3988],\n",
      "         [0.0516, 1.7426, 1.3988],\n",
      "         [0.0200, 1.7426, 1.3988]],\n",
      "\n",
      "        [[1.0000, 1.7400, 1.3987],\n",
      "         [0.9684, 1.7400, 1.3987],\n",
      "         [0.9368, 1.7400, 1.3987],\n",
      "         ...,\n",
      "         [0.0832, 1.7400, 1.3987],\n",
      "         [0.0516, 1.7400, 1.3987],\n",
      "         [0.0200, 1.7400, 1.3987]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[1.0000, 1.4043, 1.7403],\n",
      "         [0.9684, 1.4043, 1.7403],\n",
      "         [0.9368, 1.4043, 1.7403],\n",
      "         ...,\n",
      "         [0.0832, 1.4043, 1.7403],\n",
      "         [0.0516, 1.4043, 1.7403],\n",
      "         [0.0200, 1.4043, 1.7403]],\n",
      "\n",
      "        [[1.0000, 1.4016, 1.7403],\n",
      "         [0.9684, 1.4016, 1.7403],\n",
      "         [0.9368, 1.4016, 1.7403],\n",
      "         ...,\n",
      "         [0.0832, 1.4016, 1.7403],\n",
      "         [0.0516, 1.4016, 1.7403],\n",
      "         [0.0200, 1.4016, 1.7403]],\n",
      "\n",
      "        [[1.0000, 1.3989, 1.7402],\n",
      "         [0.9684, 1.3989, 1.7402],\n",
      "         [0.9368, 1.3989, 1.7402],\n",
      "         ...,\n",
      "         [0.0832, 1.3989, 1.7402],\n",
      "         [0.0516, 1.3989, 1.7402],\n",
      "         [0.0200, 1.3989, 1.7402]]], device='cuda:2')\n",
      "Depths: tensor([[ 1.0001,  1.0327,  1.0675,  ..., 12.0161, 19.3760, 50.0026],\n",
      "        [ 1.0000,  1.0327,  1.0675,  ..., 12.0159, 19.3757, 50.0017],\n",
      "        [ 1.0000,  1.0326,  1.0675,  ..., 12.0151, 19.3744, 49.9984],\n",
      "        ...,\n",
      "        [ 0.9999,  1.0325,  1.0674,  ..., 12.0140, 19.3726, 49.9938],\n",
      "        [ 0.9999,  1.0326,  1.0674,  ..., 12.0144, 19.3732, 49.9954],\n",
      "        [ 1.0000,  1.0326,  1.0675,  ..., 12.0152, 19.3745, 49.9987]],\n",
      "       device='cuda:2')\n",
      "Encoded: tensor([[[ 1.0000,  1.7454,  1.3988,  ..., -0.9968,  0.1395,  0.9952],\n",
      "         [ 0.9684,  1.7454,  1.3988,  ...,  0.8486,  0.1395,  0.9952],\n",
      "         [ 0.9368,  1.7454,  1.3988,  ..., -0.5103,  0.1395,  0.9952],\n",
      "         ...,\n",
      "         [ 0.0832,  1.7454,  1.3988,  ...,  0.1988,  0.1395,  0.9952],\n",
      "         [ 0.0516,  1.7454,  1.3988,  ...,  0.2742,  0.1395,  0.9952],\n",
      "         [ 0.0200,  1.7454,  1.3988,  ..., -0.6857,  0.1395,  0.9952]],\n",
      "\n",
      "        [[ 1.0000,  1.7426,  1.3988,  ..., -0.9968,  0.9999,  0.9953],\n",
      "         [ 0.9684,  1.7426,  1.3988,  ...,  0.8486,  0.9999,  0.9953],\n",
      "         [ 0.9368,  1.7426,  1.3988,  ..., -0.5103,  0.9999,  0.9953],\n",
      "         ...,\n",
      "         [ 0.0832,  1.7426,  1.3988,  ...,  0.1988,  0.9999,  0.9953],\n",
      "         [ 0.0516,  1.7426,  1.3988,  ...,  0.2742,  0.9999,  0.9953],\n",
      "         [ 0.0200,  1.7426,  1.3988,  ..., -0.6857,  0.9999,  0.9953]],\n",
      "\n",
      "        [[ 1.0000,  1.7400,  1.3987,  ..., -0.9968,  0.2253,  0.9881],\n",
      "         [ 0.9684,  1.7400,  1.3987,  ...,  0.8486,  0.2253,  0.9881],\n",
      "         [ 0.9368,  1.7400,  1.3987,  ..., -0.5103,  0.2253,  0.9881],\n",
      "         ...,\n",
      "         [ 0.0832,  1.7400,  1.3987,  ...,  0.1988,  0.2253,  0.9881],\n",
      "         [ 0.0516,  1.7400,  1.3987,  ...,  0.2742,  0.2253,  0.9881],\n",
      "         [ 0.0200,  1.7400,  1.3987,  ..., -0.6857,  0.2253,  0.9881]],\n",
      "\n",
      "        ...,\n",
      "\n",
      "        [[ 1.0000,  1.4043,  1.7403,  ..., -0.9968, -0.9210,  0.3760],\n",
      "         [ 0.9684,  1.4043,  1.7403,  ...,  0.8486, -0.9210,  0.3760],\n",
      "         [ 0.9368,  1.4043,  1.7403,  ..., -0.5103, -0.9210,  0.3760],\n",
      "         ...,\n",
      "         [ 0.0832,  1.4043,  1.7403,  ...,  0.1988, -0.9210,  0.3760],\n",
      "         [ 0.0516,  1.4043,  1.7403,  ...,  0.2742, -0.9210,  0.3760],\n",
      "         [ 0.0200,  1.4043,  1.7403,  ..., -0.6857, -0.9210,  0.3760]],\n",
      "\n",
      "        [[ 1.0000,  1.4016,  1.7403,  ..., -0.9968,  0.2445,  0.3786],\n",
      "         [ 0.9684,  1.4016,  1.7403,  ...,  0.8486,  0.2445,  0.3786],\n",
      "         [ 0.9368,  1.4016,  1.7403,  ..., -0.5103,  0.2445,  0.3786],\n",
      "         ...,\n",
      "         [ 0.0832,  1.4016,  1.7403,  ...,  0.1988,  0.2445,  0.3786],\n",
      "         [ 0.0516,  1.4016,  1.7403,  ...,  0.2742,  0.2445,  0.3786],\n",
      "         [ 0.0200,  1.4016,  1.7403,  ..., -0.6857,  0.2445,  0.3786]],\n",
      "\n",
      "        [[ 1.0000,  1.3989,  1.7402,  ..., -0.9968,  0.9995,  0.3247],\n",
      "         [ 0.9684,  1.3989,  1.7402,  ...,  0.8486,  0.9995,  0.3247],\n",
      "         [ 0.9368,  1.3989,  1.7402,  ..., -0.5103,  0.9995,  0.3247],\n",
      "         ...,\n",
      "         [ 0.0832,  1.3989,  1.7402,  ...,  0.1988,  0.9995,  0.3247],\n",
      "         [ 0.0516,  1.3989,  1.7402,  ...,  0.2742,  0.9995,  0.3247],\n",
      "         [ 0.0200,  1.3989,  1.7402,  ..., -0.6857,  0.9995,  0.3247]]],\n",
      "       device='cuda:2')\n"
     ]
    }
   ],
   "source": [
    "test_view = view.Trans(\n",
    "    torch.tensor([[0.0, 0.0, 0.0]], device=device.GetDevice()),\n",
    "    torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]], device=device.GetDevice())\n",
    ")\n",
    "perf = SimplePerf(True, True)\n",
    "rays_o, rays_d = gen.layer_cams[0].get_global_rays(test_view, True)\n",
    "perf.Checkpoint(\"GetRays\")\n",
    "rays_o = rays_o.view(-1, 3)\n",
    "rays_d = rays_d.view(-1, 3)\n",
    "coords, pts, depths = fovea_net.sampler(rays_o, rays_d)\n",
    "perf.Checkpoint(\"Sample\")\n",
    "encoded = fovea_net.input_encoder(coords)\n",
    "perf.Checkpoint(\"Encode\")\n",
    "print(\"Rays:\", rays_d)\n",
    "print(\"Spherical coords:\", coords)\n",
    "print(\"Depths:\", depths)\n",
    "print(\"Encoded:\", encoded)\n",
    "#plot_figures(images, center)\n",
    "\n",
    "#util.CreateDirIfNeed('output/teasers')\n",
    "#for key in images:\n",
    "#    util.WriteImageTensor(\n",
    "#        images[key], 'output/teasers/view%04d_%s.png' % (view_idx, key))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.9"
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}