{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QMMWf9AQcdlp"
   },
   "source": [
    "# Render a HyperNeRF video!\n",
    "\n",
    "**Author**: [Keunhong Park](https://keunhong.com)\n",
    "\n",
    "[[Project Page](https://hypernerf.github.io)]\n",
    "[[Paper](https://arxiv.org/abs/2106.13228)]\n",
    "[[GitHub](https://github.com/google/hypernerf)]\n",
    "\n",
    "This notebook renders a video using the test cameras generated in the capture processing notebook.\n",
    "\n",
    "You can also load your own custom cameras by modifying the code slightly.\n",
    "\n",
    "### Instructions\n",
    "\n",
    "1. Convert a video into our dataset format using the [capture processing notebook](https://colab.sandbox.google.com/github/google/nerfies/blob/main/notebooks/Nerfies_Capture_Processing.ipynb).\n",
    "2. Train a HyperNeRF model using the [training notebook](https://colab.sandbox.google.com/github/google/hypernerf/blob/main/notebooks/HyperNeRF_Training.ipynb)\n",
    "3. Run this notebook!\n",
    "\n",
    "\n",
    "### Notes\n",
    " * Please report issues on the [GitHub issue tracker](https://github.com/google/hypernerf/issues)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gHqkIo4hcGou"
   },
   "source": [
    "## Environment Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ws81Eje47SuV",
    "outputId": "2fa89ef5-4030-46d4-e2d9-2eebffd1b0f9"
   },
   "outputs": [],
   "source": [
    "#!wget https://raw.githubusercontent.com/google/hypernerf/main/requirements.txt\n",
    "!wget https://raw.githubusercontent.com/xieyizheng/hypernerf/main/requirements.txt\n",
    "!python --version\n",
    "!pip install -r requirements.txt\n",
    "\n",
    "#if freshly installed, recommend to restart the runtime!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-3T2lBKBcIGP",
    "outputId": "6bcc5d9c-108a-4c2b-bef5-fe140c87b3fb"
   },
   "outputs": [],
   "source": [
    "# @title Configure notebook runtime\n",
    "# @markdown If you would like to use a GPU runtime instead, change the runtime type by going to `Runtime > Change runtime type`. \n",
    "# @markdown You will have to use a smaller batch size on GPU.\n",
    "import jax\n",
    "runtime_type = 'gpu'  # @param ['gpu', 'tpu']\n",
    "if runtime_type == 'tpu':\n",
    "  import jax.tools.colab_tpu\n",
    "  jax.tools.colab_tpu.setup_tpu()\n",
    "\n",
    "print('Detected Devices:', jax.devices())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "82kU-W1NcNTW",
    "outputId": "08a21bab-c3cc-43a0-9f1f-fb7e843a8aaa"
   },
   "outputs": [],
   "source": [
    "# @title Mount Google Drive\n",
    "# @markdown Mount Google Drive onto `/content/gdrive`. You can skip this if running locally.\n",
    "\n",
    "from google.colab import drive\n",
    "drive.mount('/content/gdrive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YIDbV769cPn1"
   },
   "outputs": [],
   "source": [
    "# @title Define imports and utility functions.\n",
    "\n",
    "import jax\n",
    "from jax.config import config as jax_config\n",
    "import jax.numpy as jnp\n",
    "from jax import grad, jit, vmap\n",
    "from jax import random\n",
    "\n",
    "import flax\n",
    "import flax.linen as nn\n",
    "from flax import jax_utils\n",
    "from flax import optim\n",
    "from flax.metrics import tensorboard\n",
    "from flax.training import checkpoints\n",
    "\n",
    "from absl import logging\n",
    "from io import BytesIO\n",
    "import random as pyrandom\n",
    "import numpy as np\n",
    "import PIL\n",
    "import IPython\n",
    "import tempfile\n",
    "import imageio\n",
    "import mediapy\n",
    "from IPython.display import display, HTML\n",
    "from base64 import b64encode\n",
    "\n",
    "\n",
    "# Monkey patch logging.\n",
    "def myprint(msg, *args, **kwargs):\n",
    " print(msg % args)\n",
    "\n",
    "logging.info = myprint \n",
    "logging.warn = myprint\n",
    "logging.error = myprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "2QYJ7dyMcw2f",
    "outputId": "73b49855-05f6-4377-a57a-3a5a061c980a"
   },
   "outputs": [],
   "source": [
    "# @title Model and dataset configuration\n",
    "# @markdown Change the directories to where you saved your capture and experiment.\n",
    "\n",
    "\n",
    "from pathlib import Path\n",
    "from pprint import pprint\n",
    "import gin\n",
    "from IPython.display import display, Markdown\n",
    "\n",
    "from hypernerf import models\n",
    "from hypernerf import modules\n",
    "from hypernerf import warping\n",
    "from hypernerf import datasets\n",
    "from hypernerf import configs\n",
    "\n",
    "\n",
    "# @markdown The working directory where the trained model is.\n",
    "train_dir = '/content/gdrive/My Drive/nerfies/hypernerf_experiments/dvd/exp2'  # @param {type: \"string\"}\n",
    "# @markdown The directory to the dataset capture.\n",
    "data_dir = '/content/gdrive/My Drive/nerfies/captures/dvd'  # @param {type: \"string\"}\n",
    "\n",
    "checkpoint_dir = Path(train_dir, 'checkpoints')\n",
    "checkpoint_dir.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "config_path = Path(train_dir, 'config.gin')\n",
    "with open(config_path, 'r') as f:\n",
    "  logging.info('Loading config from %s', config_path)\n",
    "  config_str = f.read()\n",
    "gin.parse_config(config_str)\n",
    "\n",
    "config_path = Path(train_dir, 'config.gin')\n",
    "with open(config_path, 'w') as f:\n",
    "  logging.info('Saving config to %s', config_path)\n",
    "  f.write(config_str)\n",
    "\n",
    "exp_config = configs.ExperimentConfig()\n",
    "train_config = configs.TrainConfig()\n",
    "eval_config = configs.EvalConfig()\n",
    "\n",
    "display(Markdown(\n",
    "    gin.config.markdown(gin.config_str())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 439
    },
    "id": "6T7LQ5QSmu4o",
    "outputId": "399c441e-b125-4a99-b36e-7b58e0256858"
   },
   "outputs": [],
   "source": [
    "# @title Create datasource and show an example.\n",
    "\n",
    "from hypernerf import datasets\n",
    "from hypernerf import image_utils\n",
    "\n",
    "dummy_model = models.NerfModel({}, 0, 0)\n",
    "datasource = exp_config.datasource_cls(\n",
    "    image_scale=exp_config.image_scale,\n",
    "    random_seed=exp_config.random_seed,\n",
    "    # Enable metadata based on model needs.\n",
    "    use_warp_id=dummy_model.use_warp,\n",
    "    use_appearance_id=(\n",
    "        dummy_model.nerf_embed_key == 'appearance'\n",
    "        or dummy_model.hyper_embed_key == 'appearance'),\n",
    "    use_camera_id=dummy_model.nerf_embed_key == 'camera',\n",
    "    use_time=dummy_model.warp_embed_key == 'time')\n",
    "\n",
    "mediapy.show_image(datasource.load_rgb(datasource.train_ids[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jEO3xcxpnCqx",
    "outputId": "15e2646e-cf00-4c86-f110-86e21b686813"
   },
   "outputs": [],
   "source": [
    "# @title Load model\n",
    "# @markdown Defines the model and initializes its parameters.\n",
    "\n",
    "from flax.training import checkpoints\n",
    "from hypernerf import models\n",
    "from hypernerf import model_utils\n",
    "from hypernerf import schedules\n",
    "from hypernerf import training\n",
    "\n",
    "rng = random.PRNGKey(exp_config.random_seed)\n",
    "np.random.seed(exp_config.random_seed + jax.process_index())\n",
    "devices_to_use = jax.devices()\n",
    "\n",
    "learning_rate_sched = schedules.from_config(train_config.lr_schedule)\n",
    "nerf_alpha_sched = schedules.from_config(train_config.nerf_alpha_schedule)\n",
    "warp_alpha_sched = schedules.from_config(train_config.warp_alpha_schedule)\n",
    "elastic_loss_weight_sched = schedules.from_config(\n",
    "train_config.elastic_loss_weight_schedule)\n",
    "hyper_alpha_sched = schedules.from_config(train_config.hyper_alpha_schedule)\n",
    "hyper_sheet_alpha_sched = schedules.from_config(\n",
    "    train_config.hyper_sheet_alpha_schedule)\n",
    "\n",
    "rng, key = random.split(rng)\n",
    "params = {}\n",
    "model, params['model'] = models.construct_nerf(\n",
    "      key,\n",
    "      batch_size=train_config.batch_size,\n",
    "      embeddings_dict=datasource.embeddings_dict,\n",
    "      near=datasource.near,\n",
    "      far=datasource.far)\n",
    "\n",
    "optimizer_def = optim.Adam(learning_rate_sched(0))\n",
    "optimizer = optimizer_def.create(params)\n",
    "\n",
    "state = model_utils.TrainState(\n",
    "    optimizer=optimizer,\n",
    "    nerf_alpha=nerf_alpha_sched(0),\n",
    "    warp_alpha=warp_alpha_sched(0),\n",
    "    hyper_alpha=hyper_alpha_sched(0),\n",
    "    hyper_sheet_alpha=hyper_sheet_alpha_sched(0))\n",
    "scalar_params = training.ScalarParams(\n",
    "    learning_rate=learning_rate_sched(0),\n",
    "    elastic_loss_weight=elastic_loss_weight_sched(0),\n",
    "    warp_reg_loss_weight=train_config.warp_reg_loss_weight,\n",
    "    warp_reg_loss_alpha=train_config.warp_reg_loss_alpha,\n",
    "    warp_reg_loss_scale=train_config.warp_reg_loss_scale,\n",
    "    background_loss_weight=train_config.background_loss_weight,\n",
    "    hyper_reg_loss_weight=train_config.hyper_reg_loss_weight)\n",
    "\n",
    "logging.info('Restoring checkpoint from %s', checkpoint_dir)\n",
    "state = checkpoints.restore_checkpoint(checkpoint_dir, state)\n",
    "step = state.optimizer.state.step + 1\n",
    "state = jax_utils.replicate(state, devices=devices_to_use)\n",
    "del params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "2KYhbpsklwAy"
   },
   "outputs": [],
   "source": [
    "# @title Define pmapped render function.\n",
    "\n",
    "import functools\n",
    "from hypernerf import evaluation\n",
    "\n",
    "devices = jax.devices()\n",
    "\n",
    "\n",
    "def _model_fn(key_0, key_1, params, rays_dict, extra_params):\n",
    "  out = model.apply({'params': params},\n",
    "                    rays_dict,\n",
    "                    extra_params=extra_params,\n",
    "                    rngs={\n",
    "                        'coarse': key_0,\n",
    "                        'fine': key_1\n",
    "                    },\n",
    "                    mutable=False)\n",
    "  return jax.lax.all_gather(out, axis_name='batch')\n",
    "\n",
    "pmodel_fn = jax.pmap(\n",
    "    # Note rng_keys are useless in eval mode since there's no randomness.\n",
    "    _model_fn,\n",
    "    in_axes=(0, 0, 0, 0, 0),  # Only distribute the data input.\n",
    "    devices=devices_to_use,\n",
    "    axis_name='batch',\n",
    ")\n",
    "\n",
    "render_fn = functools.partial(evaluation.render_image,\n",
    "                              model_fn=pmodel_fn,\n",
    "                              device_count=len(devices),\n",
    "                              chunk=eval_config.chunk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "73Fq0kNcmAra",
    "outputId": "01f7bcee-833f-47fb-d2ab-0a9a2c15837f"
   },
   "outputs": [],
   "source": [
    "# @title Load cameras.\n",
    "\n",
    "from hypernerf import utils\n",
    "\n",
    "camera_path = 'camera-paths/orbit-mild'  # @param {type: 'string'}\n",
    "\n",
    "camera_dir = Path(data_dir, camera_path)\n",
    "print(f'Loading cameras from {camera_dir}')\n",
    "test_camera_paths = datasource.glob_cameras(camera_dir)\n",
    "test_cameras = utils.parallel_map(datasource.load_camera, test_camera_paths, show_pbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "aP9LjiAZmoRc",
    "outputId": "811dfbc3-ccbc-4748-dee8-92281ea01b2c"
   },
   "outputs": [],
   "source": [
    "# @title Render video frames.\n",
    "from hypernerf import visualization as viz\n",
    "\n",
    "\n",
    "rng = rng + jax.process_index()  # Make random seed separate across hosts.\n",
    "keys = random.split(rng, len(devices))\n",
    "\n",
    "results = []\n",
    "for i in range(len(test_cameras)):\n",
    "  print(f'Rendering frame {i+1}/{len(test_cameras)}')\n",
    "  camera = test_cameras[i]\n",
    "  batch = datasets.camera_to_rays(camera)\n",
    "  batch['metadata'] = {\n",
    "      'appearance': jnp.zeros_like(batch['origins'][..., 0, jnp.newaxis], jnp.uint32),\n",
    "      'warp': jnp.zeros_like(batch['origins'][..., 0, jnp.newaxis], jnp.uint32),\n",
    "  }\n",
    "  #these two are the \"ambient dimensions\" or \"time axis\" for rendering\n",
    "  batch['metadata']['appearance'] += i\n",
    "  batch['metadata']['warp'] += i\n",
    "\n",
    "  render = render_fn(state, batch, rng=rng)\n",
    "  rgb = np.array(render['rgb'])\n",
    "  depth_med = np.array(render['med_depth'])\n",
    "  results.append((rgb, depth_med))\n",
    "  depth_viz = viz.colorize(depth_med.squeeze(), cmin=datasource.near, cmax=datasource.far, invert=True)\n",
    "  mediapy.show_images([rgb, depth_viz])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_5hHR9XVm8Ix"
   },
   "outputs": [],
   "source": [
    "# @title Show rendered video.\n",
    "\n",
    "fps = 30  # @param {type:'number'}\n",
    "\n",
    "frames = []\n",
    "for rgb, depth in results:\n",
    "  depth_viz = viz.colorize(depth.squeeze(), cmin=datasource.near, cmax=datasource.far, invert=True)\n",
    "  frame = np.concatenate([rgb, depth_viz], axis=1)\n",
    "  frames.append(image_utils.image_to_uint8(frame))\n",
    "\n",
    "mediapy.show_video(frames, fps=fps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WW32AVGR0Vwh"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "machine_shape": "hm",
   "provenance": []
  },
  "gpuClass": "standard",
  "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}