{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3c00934d-0f14-4c1a-a563-94d627b377e5",
   "metadata": {},
   "source": [
    "## Notebook to visualize a dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15dfdc1-5054-40bc-9631-4c7426eecfd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_imports import *\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from pyrad.nerf.dataset.image_dataset import ImageDataset, PanopticImageDataset\n",
    "from pyrad.nerf.image_sampler import CacheImageSampler\n",
    "from pyrad.nerf.pixel_sampler import PixelSampler\n",
    "from pyrad.nerf.dataset.utils import get_dataset_inputs\n",
    "from pyrad.nerf.ray_generator import RayGenerator\n",
    "from pyrad.nerf.colliders import SceneBoundsCollider, AABBBoxCollider\n",
    "from pyrad.structures.rays import RayBundle\n",
    "from pyrad.utils.io import get_absolute_path\n",
    "from pyrad.viewer.plotly import get_line_segments_from_lines\n",
    "from pyrad.structures.cameras import get_camera_model\n",
    "from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config\n",
    "\n",
    "\n",
    "from hydra import compose, initialize\n",
    "from omegaconf import open_dict\n",
    "import pprint\n",
    "from tqdm import tqdm\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fe7478d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with initialize(version_base=None, config_path=\"../configs\"):\n",
    "    config = compose(config_name=\"default_setup.yaml\")\n",
    "dataset_inputs = get_dataset_inputs(**config.data.dataset, split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001bdf3b-60a0-4bb2-9f88-ae36f2036c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\"\n",
    "\n",
    "train_image_dataset = instantiate_from_dict_config(\n",
    "    config.data.image_dataset,\n",
    "    image_filenames=dataset_inputs.image_filenames,\n",
    "    downscale_factor=dataset_inputs.downscale_factor,\n",
    "    semantics=dataset_inputs.semantics,\n",
    "    alpha_color=dataset_inputs.alpha_color,\n",
    ")  # ImageDataset\n",
    "train_image_sampler = CacheImageSampler(\n",
    "    train_image_dataset,\n",
    "    num_samples_to_collate=len(train_image_dataset)\n",
    "    if config.data.image_sampler.num_images_to_sample_from == 0\n",
    "    else config.data.image_sampler.num_images_to_sample_from,\n",
    "    num_times_to_repeat_images=config.data.image_sampler.num_times_to_repeat_images,\n",
    "    device=device if config.data.image_sampler.move_to_graph_device else \"cpu\",\n",
    ")  # ImageSampler\n",
    "train_pixel_sampler = PixelSampler(\n",
    "    num_rays_per_batch=config.data.pixel_sampler.num_rays_per_batch, keep_full_image=True\n",
    ")  # PixelSampler\n",
    "\n",
    "ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)\n",
    "\n",
    "iter_train_image_sampler = iter(train_image_sampler)\n",
    "num_batches = 10\n",
    "for _ in tqdm(range(num_batches)):\n",
    "    image_batch = next(iter_train_image_sampler)\n",
    "    batch = train_pixel_sampler.sample(image_batch)\n",
    "    ray_bundle = ray_generator.forward(batch[\"indices\"])\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b85061-a1b0-458b-9d00-96cde6f37dca",
   "metadata": {},
   "outputs": [],
   "source": [
    "media.show_image(train_image_dataset.get_image(10), height=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b90f4a78-4b8b-4615-a127-7d7579874765",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_batch(batch):\n",
    "    # set the color of the sampled rays\n",
    "    print(batch.keys())\n",
    "    c, y, x = [i.flatten() for i in torch.split(batch[\"local_indices\"], 1, dim=-1)]\n",
    "    batch[\"image\"][c, y, x] = 0.0\n",
    "\n",
    "    # batch[\"image\"] is num_images, h, w, 3\n",
    "    images = torch.split(batch[\"image\"], 1, dim=0)\n",
    "    image_list = [image[0] for image in images]\n",
    "    image = torch.cat(image_list, dim=1)  # cat along the width dimension\n",
    "\n",
    "    # the black pixels are rays\n",
    "    media.show_image((image * 255).to(torch.uint8))\n",
    "\n",
    "\n",
    "def sample_and_show_batch():\n",
    "    image_batch = next(iter_train_image_sampler)\n",
    "    batch = train_pixel_sampler.sample(image_batch)\n",
    "    show_batch(batch)\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09e625f7-fa1f-4b3f-acd4-bddfcdf7cd9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = sample_and_show_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd8cf48a-823e-4395-99ca-eda4c2026440",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray_indices = batch[\"indices\"]\n",
    "ray_bundle = ray_generator(ray_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "646107c0-f979-408a-941f-98b2ad3e07f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "skip = 1\n",
    "size = 8\n",
    "data = []\n",
    "data += [\n",
    "    go.Scatter3d(\n",
    "        x=ray_generator.camera_to_world[::skip, 0, 3],\n",
    "        y=ray_generator.camera_to_world[::skip, 1, 3],\n",
    "        z=ray_generator.camera_to_world[::skip, 2, 3],\n",
    "        mode=\"markers\",\n",
    "        name=\"origins\",\n",
    "        marker=dict(color=\"rgba(0, 0, 0, 1)\", size=size),\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc627f05",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_bounds_collider = AABBBoxCollider(dataset_inputs.scene_bounds)\n",
    "intersected_ray_bundle = scene_bounds_collider(ray_bundle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c73ba9bd-2902-4c46-b592-bc87ac8c71f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = torch.stack(\n",
    "    [\n",
    "        intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.nears[..., None],\n",
    "        intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.fars[..., None],\n",
    "    ],\n",
    "    dim=1,\n",
    ").tolist()  # (num_rays, 2, 3)\n",
    "lines = torch.tensor(random.sample(lines, k=100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "991df6a2-efe3-4412-ba41-9a244e2e90cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "data += get_line_segments_from_lines(lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b2404d8-ca9e-4630-9f67-7c68d19458cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "layout = go.Layout(\n",
    "    autosize=False,\n",
    "    width=1000,\n",
    "    height=1000,\n",
    "    margin=go.layout.Margin(l=50, r=50, b=100, t=100, pad=4),\n",
    "    scene=go.layout.Scene(\n",
    "        aspectmode=\"data\",\n",
    "        camera=dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=0), eye=dict(x=1.25, y=1.25, z=1.25)),\n",
    "    ),\n",
    ")\n",
    "fig = go.Figure(data=data, layout=layout)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd04703a-aa50-45e4-95fd-4d881ff18d61",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "4a749acc7d255f078aee52e0584cc77b3cb5aaed1b3a7407ec4262c1bf6cb526"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}