{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2dc93f2-5194-466f-8400-bfa78758c570",
   "metadata": {},
   "source": [
    "# Interactive Visualization of 3D Gaussian Splats with nerfstudio's gsplat renderer\n",
    "\n",
    "Here we show how to use Kaolin's **camera conversion**, **ply/USD readers** and **interactive ipython visualizer** to view a 3D Guassian Splat model rendered with [nerfstudio's gsplat library](https://github.com/nerfstudio-project/gsplat). \n",
    "\n",
    "Specifically, we'll show how to interactively control aligned mesh and splat renderings using same camera conventions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27a39f75-4e1c-41e6-af0f-4b80440e343b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q matplotlib\n",
    "!pip install gsplat\n",
    "\n",
    "import copy\n",
    "import gsplat.rendering\n",
    "import ipywidgets\n",
    "import ipyevents\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import torch\n",
    "\n",
    "import kaolin\n",
    "from kaolin.render.camera import kaolin_camera_to_gsplat_nerfstudio\n",
    "from kaolin.utils.bundled_data import SCANNED_TOYS_PATH, SCANNED_TOYS_NAMES, download_scanned_toys_dataset\n",
    "from kaolin.utils.log import log_tensor\n",
    "from kaolin.visualize.ipython import quick_viz\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "kaolin.utils.log.default_log_setup(logging.DEBUG)\n",
    "    \n",
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6479dbe6-d875-481c-89bd-16d9c77b71ff",
   "metadata": {},
   "source": [
    "### Preliminaries\n",
    "\n",
    "Let's download sample data and install gsplat. Note that our data contains an aligned **3D Gaussian Splat** model and a **low-poly mesh** approximation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "020694d6-48b1-4825-b916-6dc7625f3914",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_scanned_toys_dataset()  # only re-downloads if needed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22bfd3ea-945d-4ffe-ba26-fc905f8f8e9c",
   "metadata": {},
   "source": [
    "### Loading Data\n",
    "\n",
    "Now let's load a gaussian splat model and mesh into torch tensors. \n",
    "\n",
    "We'll examine tensors and their stats: note that gaussian `positions` and mesh `vertices` have similar stats, - as these two are already aligned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c6eecca-12de-4fce-a657-24087ba19350",
   "metadata": {},
   "outputs": [],
   "source": [
    "TOY_NAME = SCANNED_TOYS_NAMES[0]\n",
    "GSMODEL = os.path.join(SCANNED_TOYS_PATH, f'{TOY_NAME}.ply')  # USD also downloaded and supported for splats\n",
    "MESH = os.path.join(SCANNED_TOYS_PATH, f'mesh.{TOY_NAME}.usd')\n",
    "\n",
    "gsmodel = kaolin.io.import_gaussiancloud(GSMODEL).to(device)\n",
    "print(gsmodel.to_string(print_stats=True))\n",
    "\n",
    "mesh = kaolin.io.import_mesh(MESH).to(device)\n",
    "print(mesh.to_string(print_stats=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af749f74-c84c-43ba-ad52-e795d4084a9b",
   "metadata": {},
   "source": [
    "### Filter by Opacity\n",
    "\n",
    "Our `GaussianSplatModel` class makes some other operations easy, for example filtering by a mask (let's filter by opacity)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b8b11dc-9121-45d7-a10a-d249bec24f8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = gsmodel.opacities > 0.1\n",
    "print(f'Keeping {mask.sum()} gaussians out of {mask.shape[0]}')\n",
    "\n",
    "gsmodel = gsmodel[mask]\n",
    "print(gsmodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c790369f-021f-48ac-aa21-060a6bc31a96",
   "metadata": {},
   "source": [
    "### Basic Rendering (Aligned Mesh and Splats)\n",
    "\n",
    "Now we'll use Kaolin's `easy_render` module to render mesh using Kaolin camera, and we will use kaolin's converter to render using nerfstudio gsplat library in the same coordinate frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f644ede-84b0-441b-9550-38c72e4e34e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def render_with_gsplat(kal_cam, gsmodel):\n",
    "    gsplat_cam_params = kaolin_camera_to_gsplat_nerfstudio(kal_cam)\n",
    "    render_colors, render_alphas, info = gsplat.rendering.rasterization(\n",
    "            gsmodel.positions,  # [N, 3]\n",
    "            gsmodel.orientations,  # [N, 4]\n",
    "            gsmodel.scales,  # [N, 3]\n",
    "            gsmodel.opacities,  # [N]\n",
    "            gsmodel.sh_coeff,  # [N, S, 3]\n",
    "            sh_degree=gsmodel.sh_degree,  \n",
    "            **gsplat_cam_params)\n",
    "    return render_colors, render_alphas, info\n",
    "\n",
    "kaolin_cam = kaolin.render.camera.Camera.from_args(\n",
    "     eye=mesh.vertices.mean(dim=0) + torch.tensor([0, -1.3, 0], device=device), at=mesh.vertices.mean(dim=0), up=torch.tensor([0., 0.0, 1.0], device=device),\n",
    "     fov=math.pi * 50 / 180, height=512, width=512).to(device)\n",
    "\n",
    "colors, alphas, info = render_with_gsplat(kaolin_cam, gsmodel)\n",
    "log_tensor(colors, 'colors', print_stats=True)\n",
    "log_tensor(alphas, 'alphas', print_stats=True)\n",
    "\n",
    "mesh_render = kaolin.render.easy_render.render_mesh(kaolin_cam, mesh)\n",
    "log_tensor(mesh_render, 'mesh_render', print_stats=True)\n",
    "\n",
    "kaolin.visualize.ipython.quick_viz(\n",
    "    torch.cat([colors, mesh_render[kaolin.render.easy_render.RenderPass.albedo]], dim=0).permute(0, 3, 1, 2), inches=20).set_title('3D Gaussian Splat Rendering vs. Aligned Low-Poly Mesh')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec0ce408-7d68-4b3f-bc31-79846e06f2e7",
   "metadata": {},
   "source": [
    "### Interactive Visualization\n",
    "\n",
    "Now we'll use Kaolin's ipython utility `IpyTurntableVisualizer` to allow interactive control of camera."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8879ba28-91a6-4970-9d12-4b1046c1a47d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tst_render(in_cam):\n",
    "    colors, _, _ = render_with_gsplat(in_cam, gsmodel)\n",
    "    colors = (colors[0, ...] * 255).clip(0, 255).to(torch.uint8)\n",
    "    return colors\n",
    "\n",
    "def tst_render_lowres(in_cam):\n",
    "    new_cam = copy.deepcopy(in_cam)\n",
    "    new_cam.width = in_cam.width // 4\n",
    "    new_cam.height = in_cam.height // 4\n",
    "    return tst_render(new_cam)\n",
    "\n",
    "visualizer = kaolin.visualize.IpyTurntableVisualizer(\n",
    "    kaolin_cam.height, kaolin_cam.width, copy.deepcopy(kaolin_cam), tst_render, fast_render=tst_render_lowres,\n",
    "    max_fps=24, world_up_axis=2, focus_at=torch.tensor([0, 0, 0.5], device=device))\n",
    "visualizer.show()\n"
   ]
  }
 ],
 "metadata": {
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}