{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Set CUDA:0 as current device.\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "rootdir = os.path.abspath(sys.path[0] + '/../')\n",
    "sys.path.append(rootdir)\n",
    "torch.cuda.set_device(0)\n",
    "print(\"Set CUDA:%d as current device.\" % torch.cuda.current_device())\n",
    "torch.autograd.set_grad_enabled(False)\n",
    "\n",
    "from data.spherical_view_syn import *\n",
    "from configs.spherical_view_syn import SphericalViewSynConfig\n",
    "from utils import netio\n",
    "from utils import img\n",
    "from utils import device\n",
    "from utils.view import *\n",
    "from components.fnr import FoveatedNeuralRenderer\n",
    "\n",
    "\n",
    "def load_net(path):\n",
    "    config = SphericalViewSynConfig()\n",
    "    config.from_id(os.path.splitext(path)[0])\n",
    "    config.SAMPLE_PARAMS['perturb_sample'] = False\n",
    "    net = config.create_net().to(device.default())\n",
    "    netio.load(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) -> 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.default()).view(-1, 3)\n",
    "        view_rots = torch.tensor(\n",
    "            data_desc['view_rots'], device=device.default()).view(-1, 3, 3)\n",
    "        return Trans(view_centers, view_rots)\n",
    "\n",
    "\n",
    "def plot_images(images):\n",
    "    plt.figure(figsize=(12, 4))\n",
    "    plt.subplot(131)\n",
    "    img.plot(images['layers_img'][0])\n",
    "    plt.subplot(132)\n",
    "    img.plot(images['layers_img'][1])\n",
    "    plt.subplot(133)\n",
    "    img.plot(images['layers_img'][2])\n",
    "    plt.figure(figsize=(12, 12))\n",
    "    img.plot(images['blended'])\n",
    "\n",
    "\n",
    "scenes = {\n",
    "    'classroom': 'classroom_all',\n",
    "    'stones': 'stones_all',\n",
    "    'barbershop': 'barbershop_all',\n",
    "    'lobby': 'lobby_all'\n",
    "}\n",
    "\n",
    "fov_list = [20, 45, 110]\n",
    "res_list = [(256, 256), (256, 256), (256, 230)]\n",
    "res_full = (1600, 1440)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Change working directory to  /home/dengnc/dvs/data/__new/lobby_all\n",
      "Load net from fovea@snerffast4-rgb_e6_fc512x4_d2.00-50.00_s64_~p.pth ...\n",
      "Load net from periph@snerffast4-rgb_e6_fc256x4_d2.00-50.00_s64_~p.pth ...\n"
     ]
    }
   ],
   "source": [
    "scene = 'lobby'\n",
    "os.chdir(f'{rootdir}/data/__new/{scenes[scene]}')\n",
    "print('Change working directory to ', os.getcwd())\n",
    "\n",
    "fovea_net = load_net(find_file('fovea'))\n",
    "periph_net = load_net(find_file('periph'))\n",
    "renderer = FoveatedNeuralRenderer(fov_list, res_list, nn.ModuleList([fovea_net, periph_net, periph_net]),\n",
    "                                  res_full, device=device.default())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dengnc/miniconda3/lib/python3.8/site-packages/torch/nn/functional.py:3828: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "params = {\n",
    "    'classroom': [\n",
    "        [0, 0, 0,   0, 0,   0, 0],\n",
    "        [0, 0, 0,   -53, 0,   0, 0],\n",
    "        [0, 0, 0,   20, -20,   0, 0]\n",
    "    ],\n",
    "    'stones': [\n",
    "        [0, 0, 0, 0, 10, -300, -50],\n",
    "        [0, 0, 0, 0, 10, 200, -50]\n",
    "    ],\n",
    "    'barbershop': [\n",
    "        [0, 0, 0, 20, 0, -300, 50],\n",
    "        [0, 0, 0, -140, -30, 150, -250],\n",
    "        [0, 0, 0, -60, -30, 75, -125],\n",
    "    ],\n",
    "    'lobby': [\n",
    "        #[0, 0, 0, 0, 0, 75, 0],\n",
    "        #[0, 0, 0, 0, 0, 5, 150],\n",
    "        [0, 0, 0, -120, 0, 75, 50],\n",
    "    ]\n",
    "}\n",
    "\n",
    "for i, param in enumerate(params[scene]):\n",
    "    view = Trans(torch.tensor(param[:3], device=device.default()),\n",
    "                 torch.tensor(euler_to_matrix([-param[4], param[3], 0]), device=device.default()).view(3, 3))\n",
    "    images = renderer(view, param[-2:])\n",
    "    if False:\n",
    "        outputdir = '../__demo/mono/'\n",
    "        misc.create_dir(outputdir)\n",
    "        img.save(images['layers_img'][0], f'{outputdir}{scene}_{i}_fovea.png')\n",
    "        img.save(images['layers_img'][1], f'{outputdir}{scene}_{i}_mid.png')\n",
    "        img.save(images['layers_img'][2], f'{outputdir}{scene}_{i}_periph.png')\n",
    "        img.save(images['blended'], f'{outputdir}{scene}_{i}_blended.png')\n",
    "    else:\n",
    "        images = plot_images(images)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Load Dataset\n",
    "views = load_views('train.json')\n",
    "print('Dataset loaded.')\n",
    "print('views:', views.size())\n",
    "for view_idx in range(views.size()[0]):\n",
    "    center = (0, 0)\n",
    "    test_view = views.get(view_idx)\n",
    "    render_view(test_view, center)\n",
    "    '''\n",
    "    images = gen(center, test_view)\n",
    "    outputdir = '../__2_demo/layer_blend/'\n",
    "    misc.create_dir(outputdir)\n",
    "    for key in images:\n",
    "        img.save(images[key], outputdir + '%s_view%04d_%s.png' % (scene, view_idx, key))\n",
    "    '''\n",
    "    '''\n",
    "    images = gen(\n",
    "        center, test_view,\n",
    "        mono_trans=Trans(test_view.trans_point(\n",
    "            torch.tensor([0.03, 0, 0], device=device.default())\n",
    "        ), test_view.r))\n",
    "    outputdir = '../__2_demo/output_mono/ref_as_right_eye/'\n",
    "    misc.create_dir(outputdir)\n",
    "    for key in images:\n",
    "    key = 'blended'\n",
    "    img.save(images[key], outputdir + '%s_view%04d_%s.png' % (scene, view_idx, key))\n",
    "    '''\n",
    "    '''\n",
    "    left_images = gen(center,\n",
    "                      Trans(\n",
    "                          test_view.trans_point(\n",
    "                              torch.tensor([-0.03, 0, 0], device=device.default())\n",
    "                          ),\n",
    "                          test_view.r),\n",
    "                      mono_trans=test_view)\n",
    "    right_images = gen(center, Trans(\n",
    "        test_view.trans_point(\n",
    "            torch.tensor([0.03, 0, 0], device=device.default())\n",
    "        ), test_view.r), mono_trans=test_view)\n",
    "    outputdir = '../__2_demo/mono_periph/stereo/'\n",
    "    misc.create_dir(outputdir)\n",
    "    key = 'blended'\n",
    "    img.save(left_images[key], outputdir + '%s_view%04d_%s_l.png' % (scene, view_idx, key))\n",
    "    img.save(right_images[key], outputdir + '%s_view%04d_%s_r.png' % (scene, view_idx, key))\n",
    "    '''\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('base': conda)",
   "name": "python385jvsc74a57bd082066b63b621a9e3d15e3b7c11ca76da6238eff3834294910d715044bd0561e5"
  },
  "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.5"
  },
  "metadata": {
   "interpreter": {
    "hash": "82066b63b621a9e3d15e3b7c11ca76da6238eff3834294910d715044bd0561e5"
   }
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}