{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000001",
   "metadata": {},
   "source": [
    "# Working with 3D Gaussian Splat Models\n",
    "\n",
    "This notebook demonstrates key features of Kaolin's `GaussianSplatModel` — a data container for **3D Gaussian Splatting** representations. We'll walk through reading and writing models from PLY and USD formats, inspecting their contents, moving tensors between devices, applying spatial transforms, concatenating scenes, and filtering by opacity.\n",
    "\n",
    "**Sections:**\n",
    "- [Reading from/to USD and PLY](#Reading-from/to-USD-and-PLY)\n",
    "- [Inspecting a GaussianSplatModel](#Inspecting-a-GaussianSplatModel)\n",
    "- [Tensor Management](#Tensor-Management)\n",
    "- [Applying Transforms](#Applying-Transforms)\n",
    "- [Concatenating Gaussian Models](#Concatenating-Gaussian-Models)\n",
    "- [Masking and Filtering](#Masking-and-Filtering)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000002",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q matplotlib\n",
    "!pip install gsplat\n",
    "\n",
    "import copy\n",
    "import gsplat.rendering\n",
    "import math\n",
    "import os\n",
    "import torch\n",
    "\n",
    "import kaolin\n",
    "from kaolin.rep import GaussianSplatModel\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",
    "%matplotlib inline\n",
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000003",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_scanned_toys_dataset()  # only re-downloads if needed\n",
    "\n",
    "TOY_A = SCANNED_TOYS_NAMES[0]  # 'BluehairRagdoll'\n",
    "TOY_B = SCANNED_TOYS_NAMES[1]  # 'bublik_octopus'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000004",
   "metadata": {},
   "source": [
    "### Reading from/to USD and PLY\n",
    "\n",
    "`kaolin.io.import_gaussiancloud` auto-detects the file format (PLY or USD/USDC). Both formats store the same Gaussian attributes; USD additionally supports scene hierarchy and transforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000005",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load from PLY\n",
    "ply_path = os.path.join(SCANNED_TOYS_PATH, f'{TOY_A}.ply')\n",
    "gsmodel_a = kaolin.io.import_gaussiancloud(ply_path).to(device)\n",
    "\n",
    "# Load from USD\n",
    "usd_path = os.path.join(SCANNED_TOYS_PATH, f'{TOY_A}.usdc')\n",
    "gsmodel_a_usd = kaolin.io.import_gaussiancloud(usd_path).to(device)\n",
    "\n",
    "print(f'Loaded from PLY: {len(gsmodel_a)} Gaussians')\n",
    "print(f'Loaded from USD: {len(gsmodel_a_usd)} Gaussians')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000006",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export to PLY (works seamlessly with individual tensor attributes)\n",
    "out_ply = os.path.join('/tmp', f'{TOY_A}_roundtrip.ply')\n",
    "kaolin.io.ply.export_gaussiancloud(\n",
    "    out_ply,\n",
    "    positions=gsmodel_a.positions,\n",
    "    orientations=gsmodel_a.orientations,\n",
    "    scales=gsmodel_a.scales,\n",
    "    opacities=gsmodel_a.opacities,\n",
    "    sh_coeff=gsmodel_a.sh_coeff,\n",
    "    overwrite=True\n",
    ")\n",
    "print(f'Saved to {out_ply}')\n",
    "\n",
    "# Verify roundtrip\n",
    "gsmodel_roundtrip = kaolin.io.import_gaussiancloud(out_ply)\n",
    "print(f'Roundtrip: {len(gsmodel_roundtrip)} Gaussians (original: {len(gsmodel_a)})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000007",
   "metadata": {},
   "source": [
    "### Inspecting a GaussianSplatModel\n",
    "\n",
    "`print()` gives a compact summary. `to_string(print_stats=True)` adds per-tensor statistics (min, max, mean) — useful for spotting scaling issues or NaNs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000008",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gsmodel_a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000009",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gsmodel_a.to_string(print_stats=True))\n",
    "print(f'Number of Gaussians: {len(gsmodel_a)}')\n",
    "print(f'SH degree: {gsmodel_a.sh_degree}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000010",
   "metadata": {},
   "source": [
    "### Tensor Management\n",
    "\n",
    "`GaussianSplatModel` supports the same `.to()`, `.cuda()`, `.cpu()` API as `torch.nn.Module`, applied to all tensor attributes at once. `check_sanity()` validates internal consistency of shapes and dtypes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000011",
   "metadata": {},
   "outputs": [],
   "source": [
    "gsmodel_cpu = gsmodel_a.cpu()\n",
    "print('After .cpu():')\n",
    "print(gsmodel_cpu.to_string(detailed=True))\n",
    "\n",
    "gsmodel_back = gsmodel_cpu.cuda()\n",
    "print('After .cuda():')\n",
    "print(gsmodel_back.to_string(detailed=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000012",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cast to float16 for memory efficiency\n",
    "gsmodel_fp16 = gsmodel_a.to(dtype=torch.float16)\n",
    "log_tensor(gsmodel_fp16.positions, 'positions (fp16)', print_stats=True)\n",
    "log_tensor(gsmodel_fp16.sh_coeff, 'sh_coeff (fp16)', print_stats=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000013",
   "metadata": {},
   "outputs": [],
   "source": [
    "gsmodel_a.check_sanity()\n",
    "print('Sanity check passed')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000014",
   "metadata": {},
   "source": [
    "### Applying Transforms\n",
    "\n",
    "`as_transformed(transform)` accepts an affine `(4, 4)` matrix and returns a **new** model with positions, orientations, scales, and SH coefficients all consistently updated. The original model is unchanged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000015",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a translation: shift model along X axis by 1.5 units\n",
    "translation = torch.eye(4, device=device)\n",
    "translation[0, 3] = 1.5\n",
    "\n",
    "gsmodel_a_shifted = gsmodel_a.as_transformed(translation)\n",
    "\n",
    "print('Original centroid:', [f'{v:.4f}' for v in gsmodel_a.positions.mean(0).cpu().tolist()])\n",
    "print('Shifted centroid: ', [f'{v:.4f}' for v in gsmodel_a_shifted.positions.mean(0).cpu().tolist()])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000016",
   "metadata": {},
   "source": [
    "### Concatenating Gaussian Models\n",
    "\n",
    "`GaussianSplatModel.cat()` merges a list of models along the point dimension. Any stored transform on each model is applied first, so the result always has `transform=None`. If models have different SH degrees, the result is capped to the minimum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000017",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load second toy and shift it to sit beside TOY_A\n",
    "ply_path_b = os.path.join(SCANNED_TOYS_PATH, f'{TOY_B}.ply')\n",
    "gsmodel_b = kaolin.io.import_gaussiancloud(ply_path_b).to(device)\n",
    "\n",
    "translation_b = torch.eye(4, device=device)\n",
    "translation_b[0, 3] = -1.5\n",
    "gsmodel_b_shifted = gsmodel_b.as_transformed(translation_b)\n",
    "\n",
    "# Concatenate both models into a single scene\n",
    "combined = GaussianSplatModel.cat([gsmodel_a, gsmodel_b_shifted])\n",
    "print(f'Combined: {len(combined)} Gaussians (A={len(gsmodel_a)}, B={len(gsmodel_b_shifted)})')\n",
    "print(combined.to_string(print_stats=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000018",
   "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",
    "center = combined.positions.mean(0)\n",
    "kaolin_cam = kaolin.render.camera.Camera.from_args(\n",
    "    eye=center + torch.tensor([0., -2.0, 0.5], device=device),\n",
    "    at=center,\n",
    "    up=torch.tensor([0., 0., 1.], device=device),\n",
    "    fov=math.pi * 60 / 180,\n",
    "    height=512, width=512).to(device)\n",
    "\n",
    "colors, alphas, _ = render_with_gsplat(kaolin_cam, combined)\n",
    "quick_viz(colors.permute(0, 3, 1, 2), inches=10).set_title('Combined Scene')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000019",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tst_render(in_cam):\n",
    "    colors, _, _ = render_with_gsplat(in_cam, combined)\n",
    "    return (colors[0] * 255).clip(0, 255).to(torch.uint8)\n",
    "\n",
    "def tst_render_lowres(in_cam):\n",
    "    cam = copy.deepcopy(in_cam)\n",
    "    cam.width = in_cam.width // 4\n",
    "    cam.height = in_cam.height // 4\n",
    "    return tst_render(cam)\n",
    "\n",
    "visualizer = kaolin.visualize.IpyTurntableVisualizer(\n",
    "    kaolin_cam.height, kaolin_cam.width, copy.deepcopy(kaolin_cam),\n",
    "    tst_render, fast_render=tst_render_lowres,\n",
    "    max_fps=24, world_up_axis=2, focus_at=center)\n",
    "visualizer.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0001-0001-0001-000000000020",
   "metadata": {},
   "source": [
    "### Masking and Filtering\n",
    "\n",
    "Boolean indexing (`model[mask]`) returns a new `GaussianSplatModel` containing only the selected Gaussians. Here we remove semi-transparent Gaussians (low opacity) to get a cleaner, more compact model, then save to disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000021",
   "metadata": {},
   "outputs": [],
   "source": [
    "opacity_threshold = 0.5\n",
    "mask = combined.opacities > opacity_threshold\n",
    "print(f'Before masking: {len(combined)} Gaussians')\n",
    "filtered = combined[mask]\n",
    "print(f'After masking:  {len(filtered)} Gaussians  (kept {100 * mask.float().mean():.1f}%)')\n",
    "print(filtered.to_string(print_stats=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000022",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare before and after masking side by side\n",
    "colors_filtered, _, _ = render_with_gsplat(kaolin_cam, filtered)\n",
    "quick_viz(\n",
    "    torch.cat([colors, colors_filtered], dim=0).permute(0, 3, 1, 2),\n",
    "    inches=20).set_title('Before vs. After Opacity Masking')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b2c3d4-0001-0001-0001-000000000023",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the filtered model to disk\n",
    "out_masked_ply = os.path.join('/tmp', 'combined_filtered.ply')\n",
    "kaolin.io.ply.export_gaussiancloud(\n",
    "    out_masked_ply,\n",
    "    positions=filtered.positions,\n",
    "    orientations=filtered.orientations,\n",
    "    scales=filtered.scales,\n",
    "    opacities=filtered.opacities,\n",
    "    sh_coeff=filtered.sh_coeff,\n",
    "    overwrite=True\n",
    ")\n",
    "print(f'Filtered model saved to {out_masked_ply}')"
   ]
  }
 ],
 "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
}