{
 "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",
    "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.gen_final import GenFinal\n",
    "\n",
    "\n",
    "def load_net(path):\n",
    "    config = SphericalViewSynConfig()\n",
    "    config.from_id(path[:-4])\n",
    "    config.sa['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) -> tuple[view.Trans, torch.Tensor]:\n",
    "    with open(data_desc_file, 'r', encoding='utf-8') as file:\n",
    "        lines = file.readlines()\n",
    "        n = len(lines) // 7\n",
    "        gazes = torch.empty(n * 2, 3)\n",
    "        views = torch.empty(n * 2, 4, 4)\n",
    "        view_idx = 0\n",
    "        for i in range(0, len(lines), 7):\n",
    "            gazes[view_idx * 2] = torch.tensor([\n",
    "                float(str) for str in lines[i + 1].split(',')\n",
    "            ])\n",
    "            gazes[view_idx * 2 + 1] = torch.tensor([\n",
    "                float(str) for str in lines[i + 2].split(',')\n",
    "            ])\n",
    "            views[view_idx * 2] = torch.tensor([\n",
    "                float(str) for str in lines[i + 3].split(',')\n",
    "            ]).view(4, 4)\n",
    "            views[view_idx * 2 + 1] = torch.tensor([\n",
    "                float(str) for str in lines[i + 4].split(',')\n",
    "            ]).view(4, 4)\n",
    "            view_idx += 1\n",
    "        gazes = gazes.to(device.default())\n",
    "        views = views.to(device.default())\n",
    "    return view.Trans(views[:, :3, 3], views[:, :3, :3]), gazes\n",
    "\n",
    "fov_list = [20, 45, 110]\n",
    "res_list = [(128, 128), (256, 256), (256, 230)]  # (192,256)]\n",
    "res_full = (1600, 1440)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(os.path.join(rootdir, 'data/__0_user_study/us_gas_all_in_one'))\n",
    "#os.chdir(os.path.join(rootdir, 'data/__0_user_study/us_mc_all_in_one'))\n",
    "#os.chdir(os.path.join(rootdir, 'data/bedroom_all_in_one'))\n",
    "#os.chdir(os.path.join(rootdir, 'data/lobby_all_in_one'))\n",
    "#os.chdir(os.path.join(rootdir, 'data/gallery_all_in_one'))\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, gazes = load_views('hmd.csv')\n",
    "print('Dataset loaded.')\n",
    "print('views:', views.size())\n",
    "\n",
    "gen = GenFinal(fov_list, res_list, res_full, fovea_net, periph_net,\n",
    "               device=device.default())\n",
    "gaze_centers = gen.full_cam.proj(gazes, center_as_origin=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for view_idx in range(gaze_centers.size(0) / 2):\n",
    "    left_center = (gaze_centers[view_idx * 2][0].item(),\n",
    "                   gaze_centers[view_idx * 2][1].item())\n",
    "    right_center = (gaze_centers[view_idx * 2 + 1][0].item(),\n",
    "                    gaze_centers[view_idx * 2 + 1][1].item())\n",
    "    left_view = views.get(view_idx * 2)\n",
    "    right_view = views.get(view_idx * 2 + 1)\n",
    "    mono_trans = view.Trans((left_view.t + right_view.t) / 2, left_view.r)\n",
    "    left_images = gen.gen(left_center, left_view, mono_trans=mono_trans)\n",
    "    right_images = gen.gen(right_center, right_view, mono_trans=mono_trans)\n",
    "    \n",
    "    os.makedirs('output/video_frames/hmd2', exist_ok=True)\n",
    "    img.save(torch.cat([left_images['blended'], right_images['blended']], -1),\n",
    "                          'output/video_frames/hmd2/view%04d.png' % view_idx)\n",
    "    print('Frame %d saved' % view_idx)\n"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}