nerf/tiny_nerf_extended.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lLDTVWKq7-ei"
   },
   "source": [
    "# PyTorch Tiny NERF\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/dataflowr/website/master/modules/extras/nerf/pipeline.jpg'>\n",
    "\n",
    "> A neural radiance field is a simple fully connected network (weights are ~5MB) trained to reproduce input views of a single scene using a rendering loss. The network directly maps from spatial location and viewing direction (5D input) to color and opacity (4D output), acting as the \"volume\" so we can use volume rendering to differentiably render new views\n",
    "\n",
    "This is an extended version of [Tiny NeRF](https://github.com/bmild/nerf/blob/master/tiny_nerf.ipynb) from [NeRF: Neural Radiance Fields](https://github.com/bmild/nerf) implemented in PyTorch.\n",
    "\n",
    "Compared to the Tensorflow implementation of Tiny NeRF, this version includes\n",
    "*   5D input including view directions\n",
    "*   Hierarchical Sampling\n",
    "\n",
    "The overall architecture is adaptable to run on small GPUs (like colab's one).\n",
    "\n",
    "The original paper [**NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis**](https://arxiv.org/abs/2003.08934) by Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi and Ren Ng appeared at ECCV 2020. \n",
    "\n",
    "Full details can be found on the [Project Website](http://www.matthewtancik.com/nerf)\n",
    "\n",
    "There is also a PyTorch implementation [nerf-pytorch](https://github.com/yenchenlin/nerf-pytorch) by [Yen-Chen Lin](https://github.com/yenchenlin)\n",
    "\n",
    "This implementation of Tiny NERF is essentially due to [Virgile Foussereau](https://github.com/Virgile-Foussereau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bZNXlxmEj0FC",
    "outputId": "ec44de8c-c879-46cd-9f56-36bf6d065652"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "except:\n",
    "    IN_COLAB = False\n",
    "\n",
    "import os, sys\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "if IN_COLAB:\n",
    "  !pip install imageio-ffmpeg\n",
    "import imageio\n",
    "from IPython.display import HTML\n",
    "from base64 import b64encode\n",
    "from tqdm.notebook import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5mTxAwgrj4yn"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('tiny_nerf_data.npz'):\n",
    "    !wget http://cseweb.ucsd.edu/~viscomp/projects/LF/papers/ECCV20/nerf/tiny_nerf_data.npz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t2dgdCDi-m3T"
   },
   "source": [
    "# Load Input Images and Poses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 285
    },
    "id": "jj1lof2ej0FI",
    "outputId": "1d3ef586-7c39-45a3-c435-652af5c155f9"
   },
   "outputs": [],
   "source": [
    "data = np.load('tiny_nerf_data.npz')\n",
    "images = data['images']\n",
    "poses = data['poses']\n",
    "focal = data['focal']\n",
    "H, W = images.shape[1:3]\n",
    "print(images.shape, poses.shape, focal)\n",
    "\n",
    "testimg, testpose = images[101], poses[101]\n",
    "images = images[:100,...,:3]\n",
    "poses = poses[:100]\n",
    "\n",
    "plt.imshow(testimg)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jxDt192E-v6i"
   },
   "source": [
    "# Optimize NeRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R1avtwVoAQTu"
   },
   "outputs": [],
   "source": [
    "def embed_fn(x, L_embed=6):\n",
    "    rets = [x]\n",
    "    for i in range(L_embed):\n",
    "        for fn in [torch.sin, torch.cos]:\n",
    "            rets.append(fn(2.**i * x))\n",
    "    return torch.cat(rets, -1).to(device)\n",
    "\n",
    "  \n",
    "class NeRF(nn.Module):\n",
    "    def __init__(self, useViewDirs=False, D=8, W=256, skip=[4], pos_embed=10, view_embed=4):\n",
    "        super(NeRF, self).__init__()\n",
    "        self.useViewDirs = useViewDirs\n",
    "        inputSize = 3 + 3*2*pos_embed\n",
    "        d_viewdirs = 3 + 3*2*view_embed\n",
    "        self.inputLayer = nn.Linear(inputSize, W)\n",
    "        self.hiddenLayers = nn.ModuleList() \n",
    "        for i in range(D-1):\n",
    "            if i in skip:\n",
    "                self.hiddenLayers.append(nn.Linear(W+inputSize, W))\n",
    "            else:\n",
    "                self.hiddenLayers.append(nn.Linear(W, W))\n",
    "        \n",
    "        if useViewDirs:\n",
    "            self.alpha_out = nn.Linear(W, 1)\n",
    "            self.rgb_filters = nn.Linear(W, W)\n",
    "            self.branch = nn.Linear(W + d_viewdirs, W // 2)\n",
    "            self.output = nn.Linear(W // 2, 3)\n",
    "        else:\n",
    "            self.outputLayer = nn.Linear(W, 4)\n",
    "        self.skip = skip\n",
    "  \n",
    "    def forward(self, x, viewdirs=None):\n",
    "        x_initial = x\n",
    "        x = nn.functional.relu(self.inputLayer(x))\n",
    "        for i, layer in enumerate(self.hiddenLayers):\n",
    "            x = nn.functional.relu(layer(x))\n",
    "            if i+1 in self.skip:\n",
    "                x = torch.cat([x, x_initial], dim=-1)\n",
    "      \n",
    "        if self.useViewDirs:\n",
    "            alpha = self.alpha_out(x)\n",
    "\n",
    "            x = self.rgb_filters(x)\n",
    "            x = torch.concat([x, viewdirs], dim=-1)\n",
    "            x = nn.functional.relu(self.branch(x))\n",
    "            x = self.output(x)\n",
    "\n",
    "            x = torch.concat([x, alpha], dim=-1)\n",
    "        else:\n",
    "            x=self.outputLayer(x)\n",
    "        return x\n",
    "\n",
    "def sample_pdf(bins, weights, N_samples, det=False):\n",
    "\n",
    "    # Get pdf\n",
    "    weights += 1e-5  # prevent nans\n",
    "    pdf = weights / torch.sum(weights, -1, keepdims=True)\n",
    "    cdf = torch.cumsum(pdf, dim=-1)\n",
    "    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)\n",
    "\n",
    "    # Take uniform samples if det=False\n",
    "    if det:\n",
    "        u = torch.linspace(0., 1., N_samples, device=cdf.device)\n",
    "        u = torch.broadcast_to(u, list(cdf.shape[:-1]) + [N_samples]) \n",
    "    else:\n",
    "        u = torch.rand(list(cdf.shape[:-1]) + [N_samples], device=cdf.device) \n",
    "\n",
    "    # Invert CDF\n",
    "    u = u.contiguous() #otherwise torch.searchsorted is not happy\n",
    "    inds = torch.searchsorted(cdf, u, right=True) \n",
    "    below = torch.clamp(inds - 1, min=0)\n",
    "    above = torch.clamp(inds, max=cdf.shape[-1]-1)\n",
    "    inds_g = torch.stack([below, above], dim=-1) \n",
    "    shape = list(inds_g.shape[:-1]) + [cdf.shape[-1]]\n",
    "    cdf_g = torch.gather(cdf.unsqueeze(-2).expand(shape), dim=-1, index=inds_g)\n",
    "    bins_g = torch.gather(bins.unsqueeze(-2).expand(shape), dim=-1, index=inds_g)\n",
    "\n",
    "    denom = (cdf_g[..., 1] - cdf_g[..., 0])\n",
    "    denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)\n",
    "    t = (u - cdf_g[..., 0]) / denom\n",
    "    samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])\n",
    "\n",
    "    return samples "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R1avtwVoAQTu"
   },
   "outputs": [],
   "source": [
    "class full_NeRF(nn.Module):\n",
    "    def __init__(self, near: float, far: float, useViewDirs=True, useHierarchicalSampling=True, pos_embed=10, view_embed=4, chunksize=2**15, rand_stratified_sampling=True, D_coarse=8, D_fine=8, skip_coarse=[4], skip_fine=[4], W_coarse=256, W_fine=258):\n",
    "        super().__init__()\n",
    "        self.near = near\n",
    "        self.far = far\n",
    "        self.useViewDirs = useViewDirs\n",
    "        self.useHierarchicalSampling = useHierarchicalSampling\n",
    "        self.rand_stratified_sampling = rand_stratified_sampling\n",
    "        self.pos_embed = pos_embed\n",
    "        self.view_embed = view_embed if useViewDirs else None;\n",
    "        self.chunksize = chunksize\n",
    "        self.coarse_model = NeRF(useViewDirs=useViewDirs, D=D_coarse, W=W_coarse, skip=skip_coarse, pos_embed=pos_embed, view_embed=view_embed)\n",
    "        self.fine_model = NeRF(useViewDirs=useViewDirs, D=D_fine, W=W_fine, skip=skip_fine, pos_embed=pos_embed, view_embed=view_embed) if useHierarchicalSampling else None;\n",
    "    \n",
    "    def make_chunks_pos(self, points):\n",
    "        points = points.reshape((-1, 3))\n",
    "        points = embed_fn(points, L_embed=self.pos_embed)\n",
    "        return [points[i:i + self.chunksize] for i in range(0, points.shape[0], self.chunksize)]\n",
    "\n",
    "    def make_chunks_view(self, points, rays_d):\n",
    "        viewdirs = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)\n",
    "        viewdirs = viewdirs[:, None, ...].expand(points.shape).reshape((-1, 3))\n",
    "        viewdirs = embed_fn(viewdirs, L_embed=self.view_embed)\n",
    "        return [viewdirs[i:i + self.chunksize] for i in range(0, viewdirs.shape[0], self.chunksize)]\n",
    "\n",
    "    def render_rays(self, raw, z_vals, rays_d):\n",
    "        # Compute opacities and colors\n",
    "        sigma_a = nn.functional.relu(raw[...,3])\n",
    "        rgb = torch.sigmoid(raw[...,:3]) \n",
    "      \n",
    "        # Do volume rendering\n",
    "        dists = z_vals[..., 1:] - z_vals[..., :-1]\n",
    "        dists = torch.cat([dists, torch.full(dists[..., :1].shape, 1e10, device=device)], dim=-1)\n",
    "        dists = dists * torch.norm(rays_d[..., None, :], dim=-1) \n",
    "        alpha = 1.-torch.exp(-sigma_a * dists)  \n",
    "        weights = torch.cumprod(1.-alpha + 1e-10, -1)\n",
    "        weights = torch.roll(weights, 1, -1)\n",
    "        weights[..., 0] = 1.\n",
    "        weights =  alpha * weights\n",
    "\n",
    "        rgb_map = torch.sum(weights[...,None] * rgb, -2) \n",
    "        return rgb_map, weights\n",
    "    \n",
    "    def forward(self, rays_o: torch.Tensor, rays_d: torch.Tensor):\n",
    "        #Stratified sampling\n",
    "        z_vals = torch.linspace(self.near, self.far, N_samples, device=device)\n",
    "        z_vals = torch.broadcast_to(z_vals,list(rays_o.shape[:-1]) + [N_samples]).clone()\n",
    "        if self.rand_stratified_sampling:\n",
    "            z_vals += torch.rand(list(rays_o.shape[:-1]) + [N_samples], device=device) * (self.far-self.near)/N_samples\n",
    "        pts = rays_o[...,None,:] + rays_d[...,None,:] * z_vals[...,:,None]\n",
    "\n",
    "        chunks_pos = self.make_chunks_pos(pts)\n",
    "        if self.useViewDirs:\n",
    "            chunks_view = self.make_chunks_view(pts, rays_d)\n",
    "\n",
    "        #coarse model pass\n",
    "        predictions = []\n",
    "\n",
    "        if self.useViewDirs:\n",
    "            for chunk_pos, chunk_view in zip(chunks_pos, chunks_view):\n",
    "                predictions.append(self.coarse_model(chunk_pos, chunk_view))\n",
    "        \n",
    "        else:\n",
    "            for chunk_pos in chunks_pos:\n",
    "                  predictions.append(self.coarse_model(chunk_pos))\n",
    "      \n",
    "        raw = torch.cat(predictions, dim=0)\n",
    "        raw = raw.reshape(list(pts.shape[:2]) + [raw.shape[-1]])\n",
    "\n",
    "        rgb_map, weights = self.render_rays(raw, z_vals, rays_d)\n",
    "\n",
    "        if not self.useHierarchicalSampling:\n",
    "            del weights\n",
    "        return rgb_map\n",
    "\n",
    "        #Hierarchical sampling\n",
    "\n",
    "        # Obtain additional integration times to evaluate based on the weights\n",
    "        # assigned to colors in the coarse model.\n",
    "        z_vals_mid = .5 * (z_vals[..., 1:] + z_vals[..., :-1])\n",
    "        z_samples = sample_pdf(z_vals_mid, weights[..., 1:-1], N_samples_hierarchical, det=True)\n",
    "        z_samples = z_samples.detach() #equivalent to tf.stop_gradient(z_samples)\n",
    "\n",
    "        # Obtain all points to evaluate color, density at.\n",
    "        z_vals_combined, _ = torch.sort(torch.cat([z_vals, z_samples], dim=-1), dim=-1)\n",
    "        pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals_combined[..., :, None]  # [N_rays, N_samples + N_samples_hierarchical, 3]\n",
    "\n",
    "        del rgb_map, weights \n",
    "\n",
    "        chunks_pos = self.make_chunks_pos(pts)\n",
    "        if self.useViewDirs:\n",
    "            chunks_view = self.make_chunks_view(pts, rays_d)\n",
    "\n",
    "        #fine model pass\n",
    "        predictions = []\n",
    "        if self.useViewDirs:\n",
    "            for chunk_pos, chunk_view in zip(chunks_pos, chunks_view):\n",
    "                predictions.append(self.fine_model(chunk_pos, chunk_view))\n",
    "        else:\n",
    "            for chunk_pos in chunks_pos:\n",
    "                predictions.append(self.fine_model(chunk_pos))\n",
    "\n",
    "        raw = torch.cat(predictions, dim=0)\n",
    "        raw = raw.reshape(list(pts.shape[:2]) + [raw.shape[-1]]) \n",
    "\n",
    "        rgb_map, _ = self.render_rays(raw, z_vals_combined, rays_d)\n",
    "\n",
    "        return rgb_map\n",
    "\n",
    "\n",
    "\n",
    "def get_rays(H, W, focal, c2w):\n",
    "    c2w = torch.from_numpy(c2w).to(device)\n",
    "    focal = torch.from_numpy(focal).to(device)\n",
    "    i, j = torch.meshgrid(torch.arange(W, dtype=torch.float32, device=device), torch.arange(H, dtype=torch.float32, device=device), indexing=\"xy\")\n",
    "    dirs = torch.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -torch.ones_like(i)], -1)\n",
    "    rays_d = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)\n",
    "    rays_o = torch.broadcast_to(c2w[:3,-1], rays_d.size())\n",
    "    return rays_o, rays_d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3TSAyVcKAiyI"
   },
   "source": [
    "Here we optimize the model. We plot a rendered holdout view and its PSNR every 50 iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 988
    },
    "id": "6XurcHoCj0FQ",
    "outputId": "c014fc12-8f00-4127-d12b-dbc31248d250"
   },
   "outputs": [],
   "source": [
    "N_iters = 1000\n",
    "psnrs = []\n",
    "iternums = []\n",
    "i_plot = 25\n",
    "loss_f = torch.nn.MSELoss()\n",
    "load_model = False\n",
    "true_NeRF = False\n",
    "chunksize = 2**15\n",
    "\n",
    "\n",
    "if not true_NeRF:\n",
    "    ### Light model if on colab or on cpu (cpu is not recommended unless you have a lot of time)\n",
    "    N_samples = 32\n",
    "    N_samples_hierarchical = 32\n",
    "    model = full_NeRF(2.0, 6.0, pos_embed=6, useViewDirs=False, useHierarchicalSampling=True, chunksize=chunksize, D_coarse=2, skip_coarse=[], D_fine=6, skip_fine=[3], W_coarse=128, W_fine=128, rand_stratified_sampling=False).to(device)\n",
    "else:\n",
    "    if device == 'cpu' or IN_COLAB:\n",
    "        print(\"Warning: using full NeRF architecture on cpu or colab is not recommended. Set true_NeRF to False.\")\n",
    "    ### True NeRF model for strong GPUs\n",
    "    N_samples = 64\n",
    "    N_samples_hierarchical = 128\n",
    "    model = full_NeRF(2.0, 6.0).to(device)\n",
    "\n",
    "lr = 5e-4\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "\n",
    "if load_model:\n",
    "    model.load_state_dict(torch.load('tiny_nerf_extended_trained.pt', map_location=device))\n",
    "    model.eval()\n",
    "    print('Model loaded')\n",
    "else:\n",
    "    print('Training model')\n",
    "    import time\n",
    "    t = time.time()\n",
    "    for i in range(N_iters+1):\n",
    "        model.train()\n",
    "        img_i = np.random.randint(images.shape[0])\n",
    "        target = images[img_i]\n",
    "        pose = poses[img_i]\n",
    "        height, width = target.shape[:2]\n",
    "        target = torch.from_numpy(target).to(device)    \n",
    "        rays_o, rays_d = get_rays(H, W, focal, pose)\n",
    "        rays_o = rays_o.reshape(-1,3)\n",
    "        rays_d = rays_d.reshape(-1,3)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        rgb = model(rays_o, rays_d)\n",
    "        rgb = rgb.reshape([height, width, 3])\n",
    "        loss = loss_f(rgb, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        del loss, rgb, target, rays_o, rays_d\n",
    "        \n",
    "        if i%i_plot==0:\n",
    "            model.eval()\n",
    "            print(i, (time.time() - t) / i_plot, 'secs per iter')\n",
    "            t = time.time()\n",
    "            \n",
    "            # Render the holdout view for logging\n",
    "            rays_o, rays_d = get_rays(H, W, focal, testpose)\n",
    "            rays_o = rays_o.reshape(-1,3)\n",
    "            rays_d = rays_d.reshape(-1,3)\n",
    "            rgb = model(rays_o, rays_d)\n",
    "            rgb = rgb.reshape([height, width, 3])\n",
    "            loss = loss_f(rgb, torch.from_numpy(testimg).to(device))\n",
    "            psnr = -10. * np.log10(loss.item())\n",
    "\n",
    "            psnrs.append(psnr)\n",
    "            iternums.append(i)\n",
    "            \n",
    "            plt.figure(figsize=(10,4))\n",
    "            plt.subplot(121)\n",
    "            plt.imshow(rgb.cpu().detach().numpy())\n",
    "            plt.title(f'Iteration: {i}')\n",
    "            plt.subplot(122)\n",
    "            plt.plot(iternums, psnrs)\n",
    "            plt.title('PSNR')\n",
    "            plt.show()\n",
    "            del rgb, loss, psnr, rays_o, rays_d\n",
    "\n",
    "    print('Done')\n",
    "    #save model\n",
    "    torch.save(model.state_dict(), 'tiny_nerf_extended_trained.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bZLEFNox_UVK"
   },
   "source": [
    "# Interactive Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L92jHDI7j0FT"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from ipywidgets import interactive, widgets\n",
    "\n",
    "trans_t = lambda t : np.array([\n",
    "    [1,0,0,0],\n",
    "    [0,1,0,0],\n",
    "    [0,0,1,t],\n",
    "    [0,0,0,1],\n",
    "], dtype=np.float32)\n",
    "\n",
    "rot_phi = lambda phi : np.array([\n",
    "    [1,0,0,0],\n",
    "    [0,np.cos(phi),-np.sin(phi),0],\n",
    "    [0,np.sin(phi), np.cos(phi),0],\n",
    "    [0,0,0,1],\n",
    "], dtype=np.float32)\n",
    "\n",
    "rot_theta = lambda th : np.array([\n",
    "    [np.cos(th),0,-np.sin(th),0],\n",
    "    [0,1,0,0],\n",
    "    [np.sin(th),0, np.cos(th),0],\n",
    "    [0,0,0,1],\n",
    "], dtype=np.float32)\n",
    "\n",
    "def pose_spherical(theta, phi, radius):\n",
    "    c2w = trans_t(radius)\n",
    "    c2w = rot_phi(phi/180.*np.pi) @ c2w\n",
    "    c2w = rot_theta(theta/180.*np.pi) @ c2w\n",
    "    c2w = np.array([[-1,0,0,0],[0,0,1,0],[0,1,0,0],[0,0,0,1]], dtype=np.float32) @ c2w\n",
    "    return c2w\n",
    "\n",
    "\n",
    "def f(**kwargs):\n",
    "    c2w = pose_spherical(**kwargs)\n",
    "    rays_o, rays_d = get_rays(H, W, focal, c2w[:3,:4])\n",
    "    rays_o = rays_o.reshape(-1,3)\n",
    "    rays_d = rays_d.reshape(-1,3)\n",
    "    rgb = model(rays_o, rays_d)\n",
    "    rgb = rgb.reshape([height, width, 3])\n",
    "    rgb = rgb.cpu().detach().numpy()\n",
    "    img = np.clip(rgb,0,1)\n",
    "    \n",
    "    plt.figure(2, figsize=(20,6))\n",
    "    plt.imshow(img)\n",
    "    plt.show()\n",
    "    \n",
    "\n",
    "sldr = lambda v, mi, ma: widgets.FloatSlider(\n",
    "    value=v,\n",
    "    min=mi,\n",
    "    max=ma,\n",
    "    step=.01,\n",
    ")\n",
    "\n",
    "names = [\n",
    "    ['theta', [100., 0., 360]],\n",
    "    ['phi', [-30., -90, 0]],\n",
    "    ['radius', [4., 3., 5.]],\n",
    "]\n",
    "\n",
    "interactive_plot = interactive(f, **{s[0] : sldr(*s[1]) for s in names})\n",
    "output = interactive_plot.children[-1]\n",
    "output.layout.height = '350px'\n",
    "interactive_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PpKhAn2a__Iu"
   },
   "source": [
    "# Render 360 Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8Sg4aV0cmVPs"
   },
   "outputs": [],
   "source": [
    "frames = []\n",
    "for th in tqdm(np.linspace(0., 360., 120, endpoint=False)):\n",
    "    c2w = pose_spherical(th, -30., 4.)\n",
    "    rays_o, rays_d = get_rays(H, W, focal, c2w[:3,:4])\n",
    "    rays_o = rays_o.reshape(-1,3)\n",
    "    rays_d = rays_d.reshape(-1,3)\n",
    "    rgb = model(rays_o, rays_d)\n",
    "    rgb = rgb.reshape([height, width, 3])\n",
    "    rgb = rgb.cpu().detach().numpy()\n",
    "    frames.append((255*np.clip(rgb,0,1)).astype(np.uint8))\n",
    "\n",
    "    del rgb, rays_d, rays_o\n",
    "\n",
    "f = 'video.mp4'\n",
    "imageio.mimwrite(f, frames, fps=30, quality=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OQ_ms-YMyFly"
   },
   "outputs": [],
   "source": [
    "mp4 = open('video.mp4','rb').read()\n",
    "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "HTML(\"\"\"\n",
    "<video width=400 controls autoplay loop>\n",
    "      <source src=\"%s\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\" % data_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GvR-v3uzCFYQ"
   },
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CT38DMfyCICr"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "dldiy",
   "language": "python",
   "name": "dldiy"
  },
  "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.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "693d9fa031ebe6d3f2c6244d39ef7f16e64fce60ab1fcc0674732e684efeda0d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}