{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import time\n",
    "\n",
    "from common import *\n",
    "from utils import img\n",
    "from utils import sphere\n",
    "from utils import device\n",
    "from utils import misc\n",
    "from utils.mem_profiler import *\n",
    "from data import Dataset, RaysLoader\n",
    "\n",
    "MemProfiler.enable = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validate Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DATA_DESC_FILE = f'{rootdir}/data/pabellon_fovea_r40x40_t0.3/train.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/gas_fovea_r80x60_t0.3_2021.01.26/train.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/nerf_fern/train.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/lobby_fovea_2021.01.18/train.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/street_fovea_r360x80_t1.0/train1.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/stones_fovea_r360x80_t1.0/train1.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/lobby_periph_r360x180_t1.0/train1.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/barbershop_periph_r360x180_t0.6/train_t0.3.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/classroom_all/nerf.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/barbershop_all/nerf_cvt.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__nerf/horns/images_4.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__new/city_fovea_r360x80_t5.0/train1.json'\n",
    "#DATA_DESC_FILE = f'{rootdir}/data/__captured/room/train.json'\n",
    "DATA_DESC_FILE = f'{rootdir}/data/__captured/bedroom/images4_train.json'\n",
    "\n",
    "MemProfiler.print_memory_stats('Start')\n",
    "\n",
    "dataset = Dataset(DATA_DESC_FILE)\n",
    "res = dataset.res\n",
    "data_loader = RaysLoader(dataset, res[0] * res[1], device=torch.device(\"cuda\"))\n",
    "\n",
    "MemProfiler.print_memory_stats('After dataset loaded')\n",
    "\n",
    "fig = plt.figure(figsize=(12, 6))\n",
    "i = 0\n",
    "for data in data_loader:\n",
    "    if i >= 4:\n",
    "        break\n",
    "    plt.subplot(2, 2, i + 1)\n",
    "    img.plot(data['color'].view(1, res[0], res[1], 3))\n",
    "    MemProfiler.print_memory_stats(f'After view {i} is plotted')\n",
    "    i += 1\n",
    "    time.sleep(1)\n",
    "\n",
    "#plt.show()\n",
    "MemProfiler.print_memory_stats(f'After all views are plotted')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "selector = torch.arange(res[0] * res[1]).reshape(res[0], res[1])\n",
    "selector = selector[::3, ::3].flatten()\n",
    "idx_range = [0, 4, 20, 24, 62, 100, 104, 120, 124]\n",
    "for r in torch.arange(11, 50, 5):\n",
    "    p = None\n",
    "    centers = None\n",
    "    pixels = None\n",
    "    idx = 0\n",
    "    MemProfiler.print_memory_stats(f'Before iter')\n",
    "    for data in data_loader:\n",
    "        if idx > max(idx_range):\n",
    "            break\n",
    "        if idx not in idx_range:\n",
    "            idx += 1\n",
    "            continue\n",
    "        colors = data['color'][selector]\n",
    "        rays_o = data['rays_o'][selector]\n",
    "        rays_d = data['rays_d'][selector]\n",
    "        r = torch.tensor([[r]], device=device.default())\n",
    "        p_ = misc.torch2np(sphere.ray_sphere_intersect(rays_o, rays_d, r)[0].view(-1, 3))\n",
    "        p = p_ if p is None else np.concatenate((p, p_), axis=0)\n",
    "        pixels_ = misc.torch2np(colors)\n",
    "        pixels = pixels_ if pixels is None else np.concatenate((pixels, pixels_), axis=0)\n",
    "        idx += 1\n",
    "    plt.figure(facecolor='white', figsize=(12, 12))\n",
    "    ax = plt.axes(projection='3d')\n",
    "    #ax = plt.subplot(1, 2, ri % 2 + 1, projection='3d')\n",
    "    plt.xlabel('x')\n",
    "    plt.ylabel('z')\n",
    "    plt.title('r = %f' % r)\n",
    "    ax.scatter([0], [0], [0], color=\"k\", s=10)\n",
    "    ax.scatter(p[:, 0], p[:, 2], p[:, 1], color=pixels, s=0.5)\n",
    "    ax.view_init(elev=0, azim=-90)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "65406b00395a48e1d89cf658ae895e7869e05878f5469716b06a752a3915211c"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('base')",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}