notebooks/walkthrough_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03abff27",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random\n",
    "import urllib.request\n",
    "import imageio as imio\n",
    "from einops import rearrange\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output, HTML, display\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "\n",
    "from models.model_utils import TrajectorySampler\n",
    "from builders.builders import build_dataloader\n",
    "from utils.utils import instantiate_from_config\n",
    "from notebooks.walkthrough_utils import get_smooth_trajectory\n",
    "from utils.camera_trajectory import rotate_n, go_forward, go_backward"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bde26aaf",
   "metadata": {},
   "source": [
    "## Load Pre-trained Model and Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cabb8078",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download pretrained model (takes anywhere from a few seconds to a few minutes depending on internet speed)\n",
    "model_url = 'https://docs-assets.developer.apple.com/ml-research/models/gsn/replica_128x128.ckpt'\n",
    "checkpoint_filename = 'replica_128x128.ckpt'\n",
    "urllib.request.urlretrieve(model_url, checkpoint_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34ee4ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: adjust the data path here if data is in a different location than default\n",
    "data_path = '../data/replica_all'\n",
    "\n",
    "checkpoint = torch.load(checkpoint_filename)\n",
    "state_dict = checkpoint['state_dict']\n",
    "\n",
    "# get rid of all the inception params which are leftover from FID metric\n",
    "keys_for_deletion = []\n",
    "for key in state_dict.keys():\n",
    "    if 'fid' in key:\n",
    "        keys_for_deletion.append(key)\n",
    "for key in keys_for_deletion:\n",
    "    del state_dict[key]\n",
    "\n",
    "opt = checkpoint['opt']\n",
    "opt.data_config.data_dir = data_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57db11fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataloader\n",
    "data_module = build_dataloader(opt.data_config, verbose=False)\n",
    "\n",
    "# collect a set of real trajectories from the dataset\n",
    "real_Rts = data_module.train_loader.dataset.get_trajectory_Rt()\n",
    "trajectory_sampler = TrajectorySampler(real_Rts=real_Rts, mode=opt.model_config.params.trajectory_mode)\n",
    "\n",
    "# initialize model and load state\n",
    "gsn = instantiate_from_config(opt.model_config).cuda().eval()\n",
    "gsn.set_trajectory_sampler(trajectory_sampler=trajectory_sampler)\n",
    "gsn.load_state_dict(state_dict, strict=True)\n",
    "\n",
    "# increase nerf_out_res after loading model to generate higher res samples (i.e., implicit upsampling)\n",
    "gsn.generator_config.params.nerf_out_res *= 2\n",
    "# trade render speed for memory by rendering smaller patches at a time\n",
    "gsn.patch_size = 32\n",
    "\n",
    "# load a batch of data so that we can use the camera parameters\n",
    "real_data = next(iter(data_module.train_dataloader()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a22fd1fb",
   "metadata": {},
   "source": [
    "## Interactive Walkthrough\n",
    "Use wasd to explore the scene. Enter 'q' to stop exploring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f730fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'\n",
    "z_dim = opt.model_config.params.decoder_config.params.z_dim\n",
    "z = torch.randn(1, z_dim, device=device)\n",
    "\n",
    "# intrinsic camera matrix from real data\n",
    "K_current = real_data['K'].to(device)\n",
    "# initialize extrinsic camera matrix at the center of the scene\n",
    "Rt_current = torch.eye(4, device=device).view(1, 1, 4, 4)\n",
    "\n",
    "trajectory = {'rgb': [], 'depth': [], 'Rt': [], 'K': []}\n",
    "\n",
    "while True:\n",
    "    camera_params = {'K': K_current, 'Rt': Rt_current}\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        fake_rgb, fake_depth, Rt, K = gsn(z, camera_params=camera_params)\n",
    "    \n",
    "    trajectory['rgb'].append(fake_rgb)\n",
    "    trajectory['depth'].append(fake_depth)\n",
    "    trajectory['Rt'].append(Rt)\n",
    "    trajectory['K'].append(K)\n",
    "    \n",
    "    rgb_current = rearrange(fake_rgb, 'b t c h w -> (b t h) w c').cpu()\n",
    "    \n",
    "    clear_output()\n",
    "    fig = plt.figure(figsize = (8, 8)) \n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.imshow(rgb_current, interpolation='bilinear')\n",
    "    ax.set_title('Current observation');\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "    action = input(\"a/w/d/s\")\n",
    "    if action == 'a':\n",
    "        # Turn left\n",
    "        Rt = rotate_n(n=-30.0).to(device)\n",
    "        Rt_current = torch.bmm(Rt.unsqueeze(0), Rt_current[0]).unsqueeze(0)\n",
    "    if action == 'd':\n",
    "        # Turn right\n",
    "        Rt = rotate_n(n=30.0).to(device)\n",
    "        Rt_current = torch.bmm(Rt.unsqueeze(0), Rt_current[0]).unsqueeze(0)\n",
    "    if action == 'w':\n",
    "        # Go forward\n",
    "        Rt_current = go_forward(Rt_current, step=opt.model_config.params.voxel_size / 0.6)\n",
    "    if action == 's':\n",
    "        # Go backward\n",
    "        Rt_current = go_backward(Rt_current, step=opt.model_config.params.voxel_size / 0.6)\n",
    "    if action == 'q':\n",
    "        break\n",
    "        \n",
    "for key in trajectory.keys():\n",
    "    trajectory[key] = torch.cat(trajectory[key], dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3603cd5",
   "metadata": {},
   "source": [
    "## Generate Animation\n",
    "Fit a spline to the camera trajectory and uniformly render frames along it to produce a smooth animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a33126b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# jitter camera pose a tiny amount to make sure each pose is unique\n",
    "# (to avoid problems with trajectory smoothing)\n",
    "trajectory['Rt'] = trajectory['Rt'] + torch.rand_like(trajectory['Rt']) * 1e-5\n",
    "\n",
    "# fit a smooth spline to the trajectory keypoints\n",
    "n_keypoints = len(trajectory['Rt'][0])\n",
    "new_Rts = get_smooth_trajectory(Rt=trajectory['Rt'][0], n_frames=5 * n_keypoints, subsample=3)\n",
    "n_steps = len(new_Rts)\n",
    "\n",
    "# render frames along new smooth trajectory\n",
    "resampled_trajectory = {'rgb': [], 'depth': [], 'Rt': [], 'K': []}\n",
    "for i in range(n_steps):\n",
    "    if i % 10 == 0:\n",
    "        print('Rendering frame {}/{}'.format(i, n_steps))\n",
    "    \n",
    "    camera_params = {'K': K_current[:1, :1], 'Rt': new_Rts[i:i+1].to(device)}\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        fake_rgb, fake_depth, Rt, K = gsn(z, camera_params=camera_params)\n",
    "\n",
    "    resampled_trajectory['rgb'].append(fake_rgb)\n",
    "    resampled_trajectory['depth'].append(fake_depth)\n",
    "    resampled_trajectory['Rt'].append(Rt)\n",
    "    resampled_trajectory['K'].append(K)\n",
    "\n",
    "for key in resampled_trajectory.keys():\n",
    "    resampled_trajectory[key] = torch.cat(resampled_trajectory[key], dim=1)\n",
    "\n",
    "imgs = []\n",
    "for i in range(n_steps):\n",
    "    im = resampled_trajectory['rgb'][0, i].permute(1, 2, 0).detach().cpu()\n",
    "    imgs.append((im * 255).byte())\n",
    "\n",
    "# save gif with random unique name\n",
    "n = ''.join(random.choice('0123456789ABCDEF') for i in range(5))\n",
    "animation_filename = 'camera_trajectory_{}.gif'.format(n)\n",
    "print('Saving animation to {}\\n'.format(animation_filename))\n",
    "imio.mimsave(animation_filename, imgs, fps=30)\n",
    "\n",
    "display(HTML('<img src={}>'.format(animation_filename)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}