{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import torch\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(2)\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 misc\n",
    "from utils import img\n",
    "from utils import device\n",
    "from utils import view\n",
    "from components.fnr import FoveatedNeuralRenderer\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.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) -> 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.default()).view(-1, 3)\n",
    "        view_rots = torch.tensor(\n",
    "            data_desc['view_rots'], device=device.default()).view(-1, 3, 3)\n",
    "        return view.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_fovea(left_images, right_images, left_center, right_center):\n",
    "    plt.figure(figsize=(8, 4))\n",
    "    plt.subplot(121)\n",
    "    img.plot(left_images['fovea'])\n",
    "    fovea_res = left_images['fovea'].size()[-2:]\n",
    "    plot_cross((0, 0), fovea_res)\n",
    "    plt.subplot(122)\n",
    "    img.plot(right_images['fovea'])\n",
    "    plot_cross((0, 0), fovea_res)\n",
    "\n",
    "\n",
    "scenes = {\n",
    "    'gas': '__0_user_study/us_gas_all_in_one',\n",
    "    'mc': '__0_user_study/us_mc_all_in_one',\n",
    "    'bedroom': 'bedroom_all_in_one',\n",
    "    'gallery': 'gallery_all_in_one',\n",
    "    'lobby': 'lobby_all_in_one'\n",
    "}\n",
    "\n",
    "fov_list = [20, 45, 110]\n",
    "res_list = [(128, 128), (256, 256), (256, 230)]\n",
    "res_full = (1600, 1440)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "centers = {\n",
    "    'gas': [\n",
    "        [(3.5, 0), (-3.5, 0)],\n",
    "        [(1.5, 0), (-1.5, 0)]\n",
    "    ],\n",
    "    'mc': [\n",
    "        [(2, 0), (-2, 0)],\n",
    "        [(2, 0), (-2, 0)]\n",
    "    ],\n",
    "    'bedroom': [\n",
    "        [(5, 0), (-5, 0)],\n",
    "        [(6, 0), (-6, 0)],\n",
    "        [(5, 0), (-5, 0)]\n",
    "    ],\n",
    "    'gallery': [\n",
    "        [(2.5, 0), (-2.5, 0)],\n",
    "        [(11.5, 0), (-11.5, 0)]\n",
    "    ]\n",
    "}\n",
    "scene = 'bedroom'\n",
    "os.chdir(os.path.join(rootdir, f'data/{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",
    "\n",
    "# Load Dataset\n",
    "views = load_views('demo.json')\n",
    "print('Dataset loaded.')\n",
    "print('views:', views.size())\n",
    "gen = GenFinal(fov_list, res_list, res_full, fovea_net, periph_net,\n",
    "               device=device.default())\n",
    "\n",
    "for view_idx in range(views.size()[0]):\n",
    "    test_view = views.get(view_idx)\n",
    "    left_images = gen(centers[scene][view_idx][0], view.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",
    "    right_images = gen(centers[scene][view_idx][1], view.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",
    "    #plot_fovea(left_images, right_images, centers[scene][view_idx][0],\n",
    "    #           centers[scene][view_idx][1])\n",
    "    outputdir = '../__2_demo/mono_periph/stereo/'\n",
    "    misc.create_dir(outputdir)\n",
    "    # for key in images:\n",
    "    key = 'blended'\n",
    "    img.save(left_images[key], '%s%s_view%04d_%s_l.png' % (outputdir, scene, view_idx, key))\n",
    "    img.save(right_images[key], '%s%s_view%04d_%s_r.png' % (outputdir, scene, view_idx, key))\n",
    "    stereo_overlap = torch.cat([left_images['blended'][:, 0:1], right_images['blended'][:, 1:3]], dim=1)\n",
    "    img.save(stereo_overlap, '%s%s_view%04d_%s_stereo.png' % (outputdir, scene, view_idx, key))\n",
    "\n",
    "    left_images = gen(centers[scene][view_idx][0], view.Trans(\n",
    "        test_view.trans_point(\n",
    "            torch.tensor([-0.03, 0, 0], device=device.default())\n",
    "        ), test_view.r))\n",
    "    right_images = gen(centers[scene][view_idx][1], view.Trans(\n",
    "        test_view.trans_point(\n",
    "            torch.tensor([0.03, 0, 0], device=device.default())\n",
    "        ), test_view.r))\n",
    "    #plot_fovea(left_images, right_images, centers[scene][view_idx][0],\n",
    "    #           centers[scene][view_idx][1])\n",
    "    outputdir = '../__2_demo/stereo/'\n",
    "    misc.create_dir(outputdir)\n",
    "    # for key in images:\n",
    "    key = 'blended'\n",
    "    img.save(left_images[key], '%s%s_view%04d_%s_l.png' % (outputdir, scene, view_idx, key))\n",
    "    img.save(right_images[key], '%s%s_view%04d_%s_r.png' % (outputdir, scene, view_idx, key))\n",
    "    stereo_overlap = torch.cat([left_images['blended'][:, 0:1], right_images['blended'][:, 1:3]], dim=1)\n",
    "    img.save(stereo_overlap, '%s%s_view%04d_%s_stereo.png' % (outputdir, scene, view_idx, key))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.9 64-bit ('pytorch': conda)",
   "name": "python379jvsc74a57bd0660ca2a75467d3af74a68fcc6f40bc78ab96b99ff17d2f100b5ca821fbb183f2"
  },
  "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
}