{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "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",
    "\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.sa['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_cross(center, res):\n",
    "    plt.plot(\n",
    "        [\n",
    "            (res[1] - 1) / 2 + center[0] - 5,\n",
    "            (res[1] - 1) / 2 + center[0] + 5\n",
    "        ],\n",
    "        [\n",
    "            (res[0] - 1) / 2 + center[1],\n",
    "            (res[0] - 1) / 2 + center[1]\n",
    "        ],\n",
    "        color=[0, 1, 0])\n",
    "    plt.plot(\n",
    "        [\n",
    "            (res[1] - 1) / 2 + center[0],\n",
    "            (res[1] - 1) / 2 + center[0]\n",
    "        ],\n",
    "        [\n",
    "            (res[0] - 1) / 2 + center[1] - 5,\n",
    "            (res[0] - 1) / 2 + center[1] + 5\n",
    "        ],\n",
    "        color=[0, 1, 0])\n",
    "\n",
    "\n",
    "def plot_figures(left_images, right_images, left_center, right_center):\n",
    "    # Plot Fovea\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(121)\n",
    "    img.plot(left_images['layers_img'][0])\n",
    "    fovea_res = left_images['layers_img'][0].size()[-2:]\n",
    "    plot_cross((0, 0), fovea_res)\n",
    "    plt.subplot(122)\n",
    "    img.plot(right_images['layers_img'][0])\n",
    "    plot_cross((0, 0), fovea_res)\n",
    "\n",
    "    # Plot Mid\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(121)\n",
    "    img.plot(left_images['layers_img'][1])\n",
    "    plt.subplot(122)\n",
    "    img.plot(right_images['layers_img'][1])\n",
    "\n",
    "    # Plot Periph\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(121)\n",
    "    img.plot(left_images['layers_img'][2])\n",
    "    plt.subplot(122)\n",
    "    img.plot(right_images['layers_img'][2])\n",
    "\n",
    "    # Plot Blended\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.subplot(121)\n",
    "    img.plot(left_images['blended'])\n",
    "    full_res = left_images['blended'].size()[-2:]\n",
    "    plot_cross(left_center, full_res)\n",
    "    plt.subplot(122)\n",
    "    img.plot(right_images['blended'])\n",
    "    plot_cross(right_center, full_res)\n",
    "\n",
    "\n",
    "scenes = {\n",
    "    'classroom': 'classroom_all',\n",
    "    'stones': 'stones_all',\n",
    "    'barbershop': 'barbershop_all',\n",
    "    'lobby': 'lobby_all'\n",
    "}\n",
    "\n",
    "\n",
    "fov_list = [20, 45, 110]\n",
    "res_list = [(256, 256), (256, 256), (256, 230)]\n",
    "res_full = (1600, 1440)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Change working directory to  /home/dengnc/dvs/data/__new/barbershop_all\n",
      "Load net from fovea200@snerffast4-rgb_e6_fc512x4_d1.20-6.00_s64_~p.pth ...\n",
      "Load net from periph200@snerffast2-rgb_e6_fc256x4_d1.20-6.00_s32_~p.pth ...\n",
      "barbershop 0 Saved\n",
      "barbershop 1 Saved\n",
      "barbershop 2 Saved\n",
      "barbershop 3 Saved\n",
      "barbershop 4 Saved\n",
      "barbershop 0 Saved\n",
      "barbershop 1 Saved\n",
      "barbershop 2 Saved\n",
      "barbershop 3 Saved\n",
      "barbershop 4 Saved\n",
      "barbershop 0 Saved\n",
      "barbershop 1 Saved\n",
      "barbershop 2 Saved\n",
      "barbershop 3 Saved\n",
      "barbershop 4 Saved\n",
      "barbershop 0 Saved\n",
      "barbershop 1 Saved\n",
      "barbershop 2 Saved\n",
      "barbershop 3 Saved\n",
      "barbershop 4 Saved\n"
     ]
    }
   ],
   "source": [
    "params = {\n",
    "    'classroom': [\n",
    "        [(0, 0, 0, 0, 0), (22, -83), (16, -83)],\n",
    "        [(0, 0, 0, 0, 0), (60, 55), (55, 55)],\n",
    "    ],\n",
    "    'lobby': [\n",
    "        [(0, 0, 0, 0, 0), (114, 0), (110, 0)],\n",
    "        [(0, 0, 0, 0, 0), (21, 150), (12, 150)]\n",
    "    ],\n",
    "    'barbershop': [\n",
    "        #[(0, 0, 0, 0, 0), (106, -67), (90, -67)],\n",
    "        #[(0, 0, 0, 0, 0), (-114, 10), (-126, 10)],\n",
    "        [(0, 0, 0, 25, 20), (189, -45), (173, -45)],\n",
    "        [(0, 0, 0, 25, 20), (-148, 130), (-163, 130)],\n",
    "        [(0.15, 0.15, 0, 43, 2), (9, 0), (-9, 0)],\n",
    "        [(0.15, 0, 0.15, -13, -5), (6, 0), (-6, 0)],\n",
    "        [(-0.15, 0.15, 0.15, -53, -21), (3, 0), (-3, 0)]\n",
    "    ]\n",
    "}\n",
    "\n",
    "#for scene in ['classroom', 'lobby', 'barbershop']:\n",
    "for scene in ['barbershop']:\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,\n",
    "                                    nn.ModuleList([fovea_net, periph_net, periph_net]),\n",
    "                                    res_full, device=device.default())\n",
    "\n",
    "    for mono_periph in range(0,4):\n",
    "        for i, param in enumerate(params[scene]):\n",
    "            view = Trans(torch.tensor(param[0][:3], device=device.default()),\n",
    "                        torch.tensor(euler_to_matrix([-param[0][4], param[0][3], 0]),\n",
    "                                    device=device.default()).view(3, 3))\n",
    "            eye_offset = torch.tensor([0.03, 0, 0], device=device.default())\n",
    "            left_view = Trans(view.trans_point(-eye_offset), view.r)\n",
    "            right_view = Trans(view.trans_point(eye_offset), view.r)\n",
    "            left_images, right_images = renderer(view, param[1], param[2],\n",
    "                                                stereo_disparity=0.06,\n",
    "                                                using_mask=True,\n",
    "                                                mono_periph_mode=mono_periph,\n",
    "                                                ret_raw=False)\n",
    "            if True:\n",
    "                outputdir = '../__demo/stereo_m%d' % mono_periph if mono_periph else '../__demo/stereo'\n",
    "                misc.create_dir(outputdir)\n",
    "                img.save(torch.cat([\n",
    "                    left_images['blended'],\n",
    "                    right_images['blended']\n",
    "                ], dim=-1), '%s/%s_%d.png' % (outputdir, scene, i))\n",
    "                img.save(left_images['blended'], '%s/%s_%d_l.png' % (outputdir, scene, i))\n",
    "                img.save(right_images['blended'], '%s/%s_%d_r.png' % (outputdir, scene, i))\n",
    "                stereo_overlap = torch.cat([\n",
    "                    left_images['blended'][:, 0:1],\n",
    "                    right_images['blended'][:, 1:3]\n",
    "                ], dim=1)\n",
    "                img.save(stereo_overlap, '%s/%s_%d_stereo.png' % (outputdir, scene, i))\n",
    "                #misc.create_dir(outputdir + '/mid')\n",
    "                #img.save(left_images['layers_img'][1], '%s/mid/%s_%d_l.png' % (outputdir, scene, i))\n",
    "                #img.save(right_images['layers_img'][1], '%s/mid/%s_%d_r.png' % (outputdir, scene, i))\n",
    "                print(\"%s %d Saved\" % (scene, i))\n",
    "            else:\n",
    "                plot_figures(left_images, right_images, param[1], param[2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "82066b63b621a9e3d15e3b7c11ca76da6238eff3834294910d715044bd0561e5"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}