Upload HyperNeRF_Render_Video_clean.ipynb

HyperNeRF_Render_Video_clean.ipynb

@@ -0,0 +1,473 @@
+{
+ "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
+}