notebooks/Capture_Processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "saNBv0dY-Eef"
   },
   "source": [
    "# Nerfies Dataset Processing.\n",
    "\n",
    "**Author**: [Keunhong Park](https://keunhong.com)\n",
    "\n",
    "[[Project Page](https://nerfies.github.io)]\n",
    "[[Paper](https://storage.googleapis.com/nerfies-public/videos/nerfies_paper.pdf)]\n",
    "[[Video](https://www.youtube.com/watch?v=MrKrnHhk8IA)]\n",
    "[[GitHub](https://github.com/google/nerfies)]\n",
    "\n",
    "This notebook contains an example workflow for converting a video file to a Nerfies dataset.\n",
    "\n",
    "### Instructions\n",
    "\n",
    "1. Convert a video into our dataset format using this notebook.\n",
    "2. Train a Nerfie using the [training notebook](https://colab.sandbox.google.com/github/google/nerfies/blob/main/notebooks/Nerfies_Training.ipynb).\n",
    "\n",
    "\n",
    "### Notes\n",
    "* While this will work for small datasets in a Colab runtime, larger datasets will require more compute power.\n",
    "* If you would like to train a model on a serious dataset, you should consider copying this to your own workstation and running it there. Some minor modifications will be required, and you will have to install the dependencies separately.\n",
    "* Please report issues on the [GitHub issue tracker](https://github.com/google/nerfies/issues).\n",
    "\n",
    "If you find this work useful, please consider citing:\n",
    "```bibtex\n",
    "@article{park2021nerfies\n",
    "  author    = {Park, Keunhong \n",
    "               and Sinha, Utkarsh \n",
    "               and Barron, Jonathan T. \n",
    "               and Bouaziz, Sofien \n",
    "               and Goldman, Dan B \n",
    "               and Seitz, Steven M. \n",
    "               and Martin-Brualla, Ricardo},\n",
    "  title     = {Nerfies: Deformable Neural Radiance Fields},\n",
    "  journal   = {ICCV},\n",
    "  year      = {2021},\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cbXoNhFF-D8Q"
   },
   "source": [
    "## Install dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8QlvguTr92ko"
   },
   "outputs": [],
   "source": [
    "!pip install mediapipe\n",
    "!pip install tensorflow_graphics\n",
    "!pip install git+https://github.com/google/nerfies.git@v2\n",
    "!pip install \"git+https://github.com/google/nerfies.git#egg=pycolmap&subdirectory=third_party/pycolmap\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!which jupyter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7Z-ASlgBUPXJ"
   },
   "source": [
    "## Configuration.\n",
    "\n",
    "Mount Google Drive onto `/content/gdrive`. You can skip this if you want to run this locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "5NR5OGyeUOKU"
   },
   "outputs": [],
   "source": [
    "# @title Configure dataset directories\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "import nerfies\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "# @markdown The base directory for all captures. This can be anything if you're running this notebook on your own Jupyter runtime.\n",
    "save_dir = '../captures'  # @param {type: 'string'}\n",
    "# @markdown The name of this capture. The working directory will be `$save_dir/$capture_name`. **Make sure you change this** when processing a new video.\n",
    "capture_name = '2-metronomes-iphone'  # @param {type: 'string'}\n",
    "# The root directory for this capture.\n",
    "root_dir = Path(save_dir, capture_name)\n",
    "# Where to save RGB images.\n",
    "rgb_dir = root_dir / 'rgb'\n",
    "rgb_raw_dir = root_dir / 'rgb-raw'\n",
    "# Where to store annotations\n",
    "annotations_dir = root_dir / 'annotations'\n",
    "annotations_file = root_dir / 'annotations.yml'\n",
    "mapping_file = root_dir / \"mapping.yml\"\n",
    "# Where to save the COLMAP outputs.\n",
    "colmap_dir = root_dir / 'colmap'\n",
    "colmap_db_path = colmap_dir / 'database.db'\n",
    "colmap_out_path = colmap_dir / 'sparse'\n",
    "\n",
    "colmap_out_path.mkdir(exist_ok=True, parents=True)\n",
    "rgb_raw_dir.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "print(f\"\"\"Directories configured:\n",
    "  root_dir = {root_dir}\n",
    "  rgb_raw_dir = {rgb_raw_dir}\n",
    "  rgb_dir = {rgb_dir}\n",
    "  colmap_dir = {colmap_dir}\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "to4QpKLFHf2s"
   },
   "source": [
    "## Dataset Processing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nscgY8DW-DHk"
   },
   "source": [
    "### Load Video.\n",
    "\n",
    "In this step we upload a video file and flatten it into PNG files using ffmpeg."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "SFzPpUoM99nd"
   },
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "file_paths = list(Path(\"../captures/\").glob(\"*.mp4\"))\n",
    "\n",
    "dropdown = widgets.Dropdown(\n",
    "  options=file_paths,\n",
    "  index=0,\n",
    ")\n",
    "\n",
    "video_path = None\n",
    "\n",
    "@widgets.interact(file_path=dropdown)\n",
    "def load(file_path):\n",
    "  global video_path\n",
    "  video_path = file_path.as_posix()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "rjnL6FdlCGhE"
   },
   "outputs": [],
   "source": [
    "# @title Flatten into images.\n",
    "import cv2\n",
    "\n",
    "\n",
    "# @markdown Flattens the video into images. The results will be saved to `rgb_raw_dir`.\n",
    "print(video_path)\n",
    "\n",
    "# @markdown Adjust `max_scale` to something smaller for faster processing.\n",
    "max_scale = 1.0  # @param {type:'number'}\n",
    "# @markdown A smaller FPS will be much faster for bundle adjustment, but at the expensive of a lower sampling density for training. For the paper we used ~15 fps but we default to something lower here to get you started faster.\n",
    "# @markdown If given an fps of -1 we will try to auto-compute it.\n",
    "fps = 15  # @param {type:'number'}\n",
    "target_num_frames = 900 # @param {type: 'number'}\n",
    "\n",
    "cap = cv2.VideoCapture(video_path)\n",
    "input_fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "\n",
    "target_num_frames = num_frames\n",
    "if num_frames < target_num_frames:\n",
    "  raise RuntimeError(\n",
    "      'The video is too short and has fewer frames than the target.')\n",
    "if fps == -1:\n",
    "  fps = int(target_num_frames / num_frames * input_fps)\n",
    "  print(f\"Auto-computed FPS = {fps}\")\n",
    "\n",
    "# @markdown Check this if you want to reprocess the frames.\n",
    "overwrite = True  # @param {type:'boolean'}\n",
    "\n",
    "tmp_rgb_raw_dir = 'rgb-raw'\n",
    "out_pattern = str('rgb-raw/%06d.png')\n",
    "if (rgb_dir / '1x').exists() and not overwrite:\n",
    "  raise RuntimeError(\n",
    "      f'The RGB frames have already been processed. Check `overwrite` and run again if you really meant to do this.')\n",
    "else:\n",
    "  filters = f\"'mpdecimate,setpts=N/FRAME_RATE/TB,scale=iw*{max_scale}:ih*{max_scale}'\"\n",
    "  !mkdir -p \"$tmp_rgb_raw_dir\"\n",
    "  !ffmpeg -i \"$video_path\" -r $fps -vf $filters \"$out_pattern\"\n",
    "  !mkdir -p \"$rgb_raw_dir\"\n",
    "  !rsync -av \"$tmp_rgb_raw_dir/\" \"$rgb_raw_dir/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continue processing the video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "5YsXeX4ckaKJ"
   },
   "outputs": [],
   "source": [
    "# @title Resize images into different scales.\n",
    "# @markdown Here we save the input images at various resolutions (downsample by a factor of 1, 2, 4, 8). We use area relation interpolation to prevent moire artifacts.\n",
    "import concurrent.futures\n",
    "import numpy as np\n",
    "import cv2\n",
    "import imageio\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def save_image(path, image: np.ndarray) -> None:\n",
    "  print(f'Saving {path}')\n",
    "  if not path.parent.exists():\n",
    "    path.parent.mkdir(exist_ok=True, parents=True)\n",
    "  with path.open('wb') as f:\n",
    "    image = Image.fromarray(np.asarray(image))\n",
    "    image.save(f, format=path.suffix.lstrip('.'))\n",
    "\n",
    "\n",
    "def image_to_uint8(image: np.ndarray) -> np.ndarray:\n",
    "  \"\"\"Convert the image to a uint8 array.\"\"\"\n",
    "  if image.dtype == np.uint8:\n",
    "    return image\n",
    "  if not issubclass(image.dtype.type, np.floating):\n",
    "    raise ValueError(\n",
    "        f'Input image should be a floating type but is of type {image.dtype!r}')\n",
    "  return (image * 255).clip(0.0, 255).astype(np.uint8)\n",
    "\n",
    "\n",
    "def make_divisible(image: np.ndarray, divisor: int) -> np.ndarray:\n",
    "  \"\"\"Trim the image if not divisible by the divisor.\"\"\"\n",
    "  height, width = image.shape[:2]\n",
    "  if height % divisor == 0 and width % divisor == 0:\n",
    "    return image\n",
    "\n",
    "  new_height = height - height % divisor\n",
    "  new_width = width - width % divisor\n",
    "\n",
    "  return image[:new_height, :new_width]\n",
    "\n",
    "\n",
    "def downsample_image(image: np.ndarray, scale: int) -> np.ndarray:\n",
    "  \"\"\"Downsamples the image by an integer factor to prevent artifacts.\"\"\"\n",
    "  if scale == 1:\n",
    "    return image\n",
    "\n",
    "  height, width = image.shape[:2]\n",
    "  if height % scale > 0 or width % scale > 0:\n",
    "    raise ValueError(f'Image shape ({height},{width}) must be divisible by the'\n",
    "                     f' scale ({scale}).')\n",
    "  out_height, out_width = height // scale, width // scale\n",
    "  resized = cv2.resize(image, (out_width, out_height), cv2.INTER_AREA)\n",
    "  return resized\n",
    "\n",
    "\n",
    "\n",
    "image_scales = \"1,2,4,8\"  # @param {type: \"string\"}\n",
    "image_scales = [int(x) for x in image_scales.split(',')]\n",
    "\n",
    "tmp_rgb_dir = Path('rgb')\n",
    "\n",
    "for image_path in Path(tmp_rgb_raw_dir).glob('*.png'):\n",
    "  image = make_divisible(imageio.imread(image_path), max(image_scales))\n",
    "  for scale in image_scales:\n",
    "    save_image(\n",
    "        tmp_rgb_dir / f'{scale}x/{image_path.stem}.png',\n",
    "        image_to_uint8(downsample_image(image, scale)))\n",
    "\n",
    "!rsync -av \"$tmp_rgb_dir/\" \"$rgb_dir/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "ql9r4rufLQue"
   },
   "outputs": [],
   "source": [
    "# @title Example frame.\n",
    "# @markdown Make sure that the video was processed correctly.\n",
    "# @markdown If this gives an exception, try running the preceding cell one more time--sometimes uploading to Google Drive can fail.\n",
    "\n",
    "from pathlib import Path\n",
    "import imageio\n",
    "from PIL import Image\n",
    "\n",
    "image_paths = list((rgb_dir / '1x').iterdir())\n",
    "Image.open(image_paths[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0YnhY66zOShI"
   },
   "source": [
    "### Camera registration with COLMAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "T2xqbzxILqZO"
   },
   "outputs": [],
   "source": [
    "# @title Extract features.\n",
    "# @markdown Computes SIFT features and saves them to the COLMAP DB.\n",
    "share_intrinsics = True  # @param {type: 'boolean'}\n",
    "assume_upright_cameras = True  # @param {type: 'boolean'}\n",
    "\n",
    "# @markdown This sets the scale at which we will run COLMAP. A scale of 1 will be more accurate but will be slow.\n",
    "colmap_image_scale = 4  # @param {type: 'number'}\n",
    "colmap_rgb_dir = rgb_dir / f'{colmap_image_scale}x'\n",
    "\n",
    "# @markdown Check this if you want to re-process SfM.\n",
    "overwrite = True  # @param {type: 'boolean'}\n",
    "\n",
    "if overwrite and colmap_db_path.exists():\n",
    "  colmap_db_path.unlink()\n",
    "\n",
    "!colmap feature_extractor \\\n",
    "--SiftExtraction.use_gpu 0 \\\n",
    "--SiftExtraction.upright {int(assume_upright_cameras)} \\\n",
    "--ImageReader.camera_model OPENCV \\\n",
    "--ImageReader.single_camera {int(share_intrinsics)} \\\n",
    "--database_path \"{str(colmap_db_path)}\" \\\n",
    "--image_path \"{str(colmap_rgb_dir)}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "7f_n95abLqw6"
   },
   "outputs": [],
   "source": [
    "# @title Match features.\n",
    "# @markdown Match the SIFT features between images. Use `exhaustive` if you only have a few images and use `vocab_tree` if you have a lot of images.\n",
    "\n",
    "match_method = 'vocab_tree'  # @param [\"exhaustive\", \"vocab_tree\"]\n",
    "\n",
    "if match_method == 'exhaustive':\n",
    "  !colmap exhaustive_matcher \\\n",
    "    --SiftMatching.use_gpu 0 \\\n",
    "    --database_path \"{str(colmap_db_path)}\"\n",
    "else:\n",
    "  # Use this if you have lots of frames.\n",
    "  !wget https://demuc.de/colmap/vocab_tree_flickr100K_words32K.bin\n",
    "  !colmap vocab_tree_matcher \\\n",
    "    --VocabTreeMatching.vocab_tree_path vocab_tree_flickr100K_words32K.bin \\\n",
    "    --SiftMatching.use_gpu 0 \\\n",
    "    --database_path \"{str(colmap_db_path)}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "aR52ZlXJOAn3"
   },
   "outputs": [],
   "source": [
    "# @title Reconstruction.\n",
    "# @markdown Run structure-from-motion to compute camera parameters.\n",
    "\n",
    "refine_principal_point = True  #@param {type:\"boolean\"}\n",
    "min_num_matches =   32 # @param {type: 'number'}\n",
    "filter_max_reproj_error = 2  # @param {type: 'number'}\n",
    "tri_complete_max_reproj_error = 2  # @param {type: 'number'}\n",
    "\n",
    "!colmap mapper \\\n",
    "  --Mapper.ba_refine_principal_point {int(refine_principal_point)} \\\n",
    "  --Mapper.filter_max_reproj_error $filter_max_reproj_error \\\n",
    "  --Mapper.tri_complete_max_reproj_error $tri_complete_max_reproj_error \\\n",
    "  --Mapper.min_num_matches $min_num_matches \\\n",
    "  --database_path \"{str(colmap_db_path)}\" \\\n",
    "  --image_path \"{str(colmap_rgb_dir)}\" \\\n",
    "  --output_path \"{str(colmap_out_path)}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "1ckBrtc9O4s4"
   },
   "outputs": [],
   "source": [
    "# @title Verify that SfM worked.\n",
    "\n",
    "if not colmap_db_path.exists():\n",
    "  raise RuntimeError(f'The COLMAP DB does not exist, did you run the reconstruction?')\n",
    "elif not (colmap_dir / 'sparse/0/cameras.bin').exists():\n",
    "  raise RuntimeError(\"\"\"\n",
    "SfM seems to have failed. Try some of the following options:\n",
    " - Increase the FPS when flattenting to images. There should be at least 50-ish images.\n",
    " - Decrease `min_num_matches`.\n",
    " - If you images aren't upright, uncheck `assume_upright_cameras`.\n",
    "\"\"\")\n",
    "else:\n",
    "  print(\"Everything looks good!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DqpRdhDBdRjT"
   },
   "source": [
    "## Parse Data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "5LuJwJawdXKw"
   },
   "outputs": [],
   "source": [
    "# @title Define Scene Manager.\n",
    "from absl import logging\n",
    "from typing import Dict\n",
    "import numpy as np\n",
    "from nerfies.camera import Camera\n",
    "import pycolmap\n",
    "from pycolmap import Quaternion\n",
    "\n",
    "\n",
    "def convert_colmap_camera(colmap_camera, colmap_image):\n",
    "  \"\"\"Converts a pycolmap `image` to an SFM camera.\"\"\"\n",
    "  camera_rotation = colmap_image.R()\n",
    "  camera_position = -(colmap_image.t @ camera_rotation)\n",
    "  new_camera = Camera(\n",
    "      orientation=camera_rotation,\n",
    "      position=camera_position,\n",
    "      focal_length=colmap_camera.fx,\n",
    "      pixel_aspect_ratio=colmap_camera.fx / colmap_camera.fx,\n",
    "      principal_point=np.array([colmap_camera.cx, colmap_camera.cy]),\n",
    "      radial_distortion=np.array([colmap_camera.k1, colmap_camera.k2, 0.0]),\n",
    "      tangential_distortion=np.array([colmap_camera.p1, colmap_camera.p2]),\n",
    "      skew=0.0,\n",
    "      image_size=np.array([colmap_camera.width, colmap_camera.height])\n",
    "  )\n",
    "  return new_camera\n",
    "\n",
    "\n",
    "def filter_outlier_points(points, inner_percentile):\n",
    "  \"\"\"Filters outlier points.\"\"\"\n",
    "  outer = 1.0 - inner_percentile\n",
    "  lower = outer / 2.0\n",
    "  upper = 1.0 - lower\n",
    "  centers_min = np.quantile(points, lower, axis=0)\n",
    "  centers_max = np.quantile(points, upper, axis=0)\n",
    "  result = points.copy()\n",
    "\n",
    "  too_near = np.any(result < centers_min[None, :], axis=1)\n",
    "  too_far = np.any(result > centers_max[None, :], axis=1)\n",
    "\n",
    "  return result[~(too_near | too_far)]\n",
    "\n",
    "\n",
    "def average_reprojection_errors(points, pixels, cameras):\n",
    "  \"\"\"Computes the average reprojection errors of the points.\"\"\"\n",
    "  cam_errors = []\n",
    "  for i, camera in enumerate(cameras):\n",
    "    cam_error = reprojection_error(points, pixels[:, i], camera)\n",
    "    cam_errors.append(cam_error)\n",
    "  cam_error = np.stack(cam_errors)\n",
    "\n",
    "  return cam_error.mean(axis=1)\n",
    "\n",
    "\n",
    "def _get_camera_translation(camera):\n",
    "  \"\"\"Computes the extrinsic translation of the camera.\"\"\"\n",
    "  rot_mat = camera.orientation\n",
    "  return -camera.position.dot(rot_mat.T)\n",
    "\n",
    "\n",
    "def _transform_camera(camera, transform_mat):\n",
    "  \"\"\"Transforms the camera using the given transformation matrix.\"\"\"\n",
    "  # The determinant gives us volumetric scaling factor.\n",
    "  # Take the cube root to get the linear scaling factor.\n",
    "  scale = np.cbrt(linalg.det(transform_mat[:, :3]))\n",
    "  quat_transform = ~Quaternion.FromR(transform_mat[:, :3] / scale)\n",
    "\n",
    "  translation = _get_camera_translation(camera)\n",
    "  rot_quat = Quaternion.FromR(camera.orientation)\n",
    "  rot_quat *= quat_transform\n",
    "  translation = scale * translation - rot_quat.ToR().dot(transform_mat[:, 3])\n",
    "  new_transform = np.eye(4)\n",
    "  new_transform[:3, :3] = rot_quat.ToR()\n",
    "  new_transform[:3, 3] = translation\n",
    "\n",
    "  rotation = rot_quat.ToR()\n",
    "  new_camera = camera.copy()\n",
    "  new_camera.orientation = rotation\n",
    "  new_camera.position = -(translation @ rotation)\n",
    "  return new_camera\n",
    "\n",
    "\n",
    "def _pycolmap_to_sfm_cameras(manager: pycolmap.SceneManager) -> Dict[int, Camera]:\n",
    "  \"\"\"Creates SFM cameras.\"\"\"\n",
    "  # Use the original filenames as indices.\n",
    "  # This mapping necessary since COLMAP uses arbitrary numbers for the\n",
    "  # image_id.\n",
    "  image_id_to_colmap_id = {\n",
    "      image.name.split('.')[0]: image_id\n",
    "      for image_id, image in manager.images.items()\n",
    "  }\n",
    "\n",
    "  sfm_cameras = {}\n",
    "  for image_id in image_id_to_colmap_id:\n",
    "    colmap_id = image_id_to_colmap_id[image_id]\n",
    "    image = manager.images[colmap_id]\n",
    "    camera = manager.cameras[image.camera_id]\n",
    "    sfm_cameras[image_id] = convert_colmap_camera(camera, image)\n",
    "\n",
    "  return sfm_cameras\n",
    "\n",
    "\n",
    "class SceneManager:\n",
    "  \"\"\"A thin wrapper around pycolmap.\"\"\"\n",
    "\n",
    "  @classmethod\n",
    "  def from_pycolmap(cls, colmap_path, image_path, min_track_length=10):\n",
    "    \"\"\"Create a scene manager using pycolmap.\"\"\"\n",
    "    manager = pycolmap.SceneManager(str(colmap_path))\n",
    "    manager.load_cameras()\n",
    "    manager.load_images()\n",
    "    manager.load_points3D()\n",
    "    manager.filter_points3D(min_track_len=min_track_length)\n",
    "    sfm_cameras = _pycolmap_to_sfm_cameras(manager)\n",
    "    return cls(sfm_cameras, manager.get_filtered_points3D(), image_path)\n",
    "\n",
    "  def __init__(self, cameras, points, image_path):\n",
    "    self.image_path = Path(image_path)\n",
    "    self.camera_dict = cameras\n",
    "    self.points = points\n",
    "\n",
    "    logging.info('Created scene manager with %d cameras', len(self.camera_dict))\n",
    "\n",
    "  def __len__(self):\n",
    "    return len(self.camera_dict)\n",
    "\n",
    "  @property\n",
    "  def image_ids(self):\n",
    "    return sorted(self.camera_dict.keys())\n",
    "\n",
    "  @property\n",
    "  def camera_list(self):\n",
    "    return [self.camera_dict[i] for i in self.image_ids]\n",
    "\n",
    "  @property\n",
    "  def camera_positions(self):\n",
    "    \"\"\"Returns an array of camera positions.\"\"\"\n",
    "    return np.stack([camera.position for camera in self.camera_list])\n",
    "\n",
    "  def load_image(self, image_id):\n",
    "    \"\"\"Loads the image with the specified image_id.\"\"\"\n",
    "    path = self.image_path / f'{image_id}.png'\n",
    "    with path.open('rb') as f:\n",
    "      return imageio.imread(f)\n",
    "\n",
    "  def triangulate_pixels(self, pixels):\n",
    "    \"\"\"Triangulates the pixels across all cameras in the scene.\n",
    "\n",
    "    Args:\n",
    "      pixels: the pixels to triangulate. There must be the same number of pixels\n",
    "        as cameras in the scene.\n",
    "\n",
    "    Returns:\n",
    "      The 3D points triangulated from the pixels.\n",
    "    \"\"\"\n",
    "    if pixels.shape != (len(self), 2):\n",
    "      raise ValueError(\n",
    "          f'The number of pixels ({len(pixels)}) must be equal to the number '\n",
    "          f'of cameras ({len(self)}).')\n",
    "\n",
    "    return triangulate_pixels(pixels, self.camera_list)\n",
    "\n",
    "  def change_basis(self, axes, center):\n",
    "    \"\"\"Change the basis of the scene.\n",
    "\n",
    "    Args:\n",
    "      axes: the axes of the new coordinate frame.\n",
    "      center: the center of the new coordinate frame.\n",
    "\n",
    "    Returns:\n",
    "      A new SceneManager with transformed points and cameras.\n",
    "    \"\"\"\n",
    "    transform_mat = np.zeros((3, 4))\n",
    "    transform_mat[:3, :3] = axes.T\n",
    "    transform_mat[:, 3] = -(center @ axes)\n",
    "    return self.transform(transform_mat)\n",
    "\n",
    "  def transform(self, transform_mat):\n",
    "    \"\"\"Transform the scene using a transformation matrix.\n",
    "\n",
    "    Args:\n",
    "      transform_mat: a 3x4 transformation matrix representation a\n",
    "        transformation.\n",
    "\n",
    "    Returns:\n",
    "      A new SceneManager with transformed points and cameras.\n",
    "    \"\"\"\n",
    "    if transform_mat.shape != (3, 4):\n",
    "      raise ValueError('transform_mat should be a 3x4 transformation matrix.')\n",
    "\n",
    "    points = None\n",
    "    if self.points is not None:\n",
    "      points = self.points.copy()\n",
    "      points = points @ transform_mat[:, :3].T + transform_mat[:, 3]\n",
    "\n",
    "    new_cameras = {}\n",
    "    for image_id, camera in self.camera_dict.items():\n",
    "      new_cameras[image_id] = _transform_camera(camera, transform_mat)\n",
    "\n",
    "    return SceneManager(new_cameras, points, self.image_path)\n",
    "\n",
    "  def filter_images(self, image_ids):\n",
    "    num_filtered = 0\n",
    "    for image_id in image_ids:\n",
    "      if self.camera_dict.pop(image_id, None) is not None:\n",
    "        num_filtered += 1\n",
    "\n",
    "    return num_filtered\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "HdAegiHVWdY9"
   },
   "outputs": [],
   "source": [
    "# @title Load COLMAP scene.\n",
    "from plotly.offline import init_notebook_mode, iplot\n",
    "\n",
    "init_notebook_mode(connected=True)     \n",
    "import plotly.graph_objs as go\n",
    "\n",
    "scene_manager = SceneManager.from_pycolmap(\n",
    "    colmap_dir / 'sparse/0', \n",
    "    rgb_dir / f'1x', \n",
    "    min_track_length=5)\n",
    "\n",
    "if colmap_image_scale > 1:\n",
    "  print(f'Scaling COLMAP cameras back to 1x from {colmap_image_scale}x.')\n",
    "  for item_id in scene_manager.image_ids:\n",
    "    camera = scene_manager.camera_dict[item_id]\n",
    "    scene_manager.camera_dict[item_id] = camera.scale(colmap_image_scale)\n",
    "\n",
    "\n",
    "fig = go.Figure()\n",
    "fig.add_trace(go.Scatter3d(\n",
    "    x=scene_manager.points[:, 0],\n",
    "    y=scene_manager.points[:, 1],\n",
    "    z=scene_manager.points[:, 2],\n",
    "    mode='markers',\n",
    "    marker=dict(size=2),\n",
    "))\n",
    "fig.add_trace(go.Scatter3d(\n",
    "    x=scene_manager.camera_positions[:, 0],\n",
    "    y=scene_manager.camera_positions[:, 1],\n",
    "    z=scene_manager.camera_positions[:, 2],\n",
    "    mode='markers',\n",
    "    marker=dict(size=2),\n",
    "))\n",
    "fig.update_layout(scene_dragmode='orbit')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "e92Kcuoa5i9h"
   },
   "outputs": [],
   "source": [
    "# @title Filter blurry frames.\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "def variance_of_laplacian(image: np.ndarray) -> np.ndarray:\n",
    "  \"\"\"Compute the variance of the Laplacian which measure the focus.\"\"\"\n",
    "  gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n",
    "  return cv2.Laplacian(gray, cv2.CV_64F).var()\n",
    "\n",
    "\n",
    "blur_filter_perc = 95.0 # @param {type: 'number'}\n",
    "if blur_filter_perc > 0.0:\n",
    "  image_paths = sorted(rgb_dir.iterdir())\n",
    "  print('Loading images.')\n",
    "  images = list(map(scene_manager.load_image, scene_manager.image_ids))\n",
    "  print(image_paths[0])\n",
    "  print('Computing blur scores.')\n",
    "  blur_scores = np.array([variance_of_laplacian(im) for im in images])\n",
    "  blur_thres = np.percentile(blur_scores, blur_filter_perc)\n",
    "  blur_filter_inds = np.where(blur_scores >= blur_thres)[0]\n",
    "  blur_filter_scores = [blur_scores[i] for i in blur_filter_inds]\n",
    "  blur_filter_inds = blur_filter_inds[np.argsort(blur_filter_scores)]\n",
    "  blur_filter_scores = np.sort(blur_filter_scores)\n",
    "  blur_filter_image_ids = [scene_manager.image_ids[i] for i in blur_filter_inds]\n",
    "  print(f'Filtering {len(blur_filter_image_ids)} IDs: {blur_filter_image_ids}')\n",
    "  num_filtered = scene_manager.filter_images(blur_filter_image_ids)\n",
    "  print(f'Filtered {num_filtered} images')\n",
    "\n",
    "  plt.figure(figsize=(15, 10))\n",
    "  plt.subplot(121)\n",
    "  plt.title('Least blurry')\n",
    "  plt.imshow(images[blur_filter_inds[-1]])\n",
    "  plt.subplot(122)\n",
    "  plt.title('Most blurry')\n",
    "  plt.imshow(images[blur_filter_inds[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xtSV7C5y3Yuv"
   },
   "source": [
    "### Face Processing.\n",
    "\n",
    "This section runs the optional step of computing facial landmarks for the purpose of test camera generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "lDOphUXt5AQ-"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from tensorflow_graphics.geometry.representation.ray import triangulate as ray_triangulate\n",
    "\n",
    "use_face = False  # @param {type: 'boolean'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "hVjyA5sW3AVZ"
   },
   "outputs": [],
   "source": [
    "# @title Compute 2D landmarks.\n",
    "\n",
    "import imageio\n",
    "import mediapipe as mp\n",
    "from PIL import Image\n",
    "\n",
    "if use_face:\n",
    "  mp_face_mesh = mp.solutions.face_mesh\n",
    "  mp_drawing = mp.solutions.drawing_utils \n",
    "  drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)\n",
    "  \n",
    "  # Initialize MediaPipe Face Mesh.\n",
    "  face_mesh = mp_face_mesh.FaceMesh(\n",
    "      static_image_mode=True,\n",
    "      max_num_faces=2,\n",
    "      min_detection_confidence=0.5)\n",
    "  \n",
    "  \n",
    "  def compute_landmarks(image):\n",
    "    height, width = image.shape[:2]\n",
    "    results = face_mesh.process(image)\n",
    "    if results.multi_face_landmarks is None:\n",
    "      return None\n",
    "    # Choose first face found.\n",
    "    landmarks = results.multi_face_landmarks[0].landmark\n",
    "    landmarks = np.array(\n",
    "        [(o.x * width, o.y * height) for o in landmarks],\n",
    "        dtype=np.uint32)\n",
    "    return landmarks\n",
    "\n",
    "  landmarks_dict = {}\n",
    "  for item_id in scene_manager.image_ids:\n",
    "    image = scene_manager.load_image(item_id)\n",
    "    landmarks = compute_landmarks(image)\n",
    "    if landmarks is not None:\n",
    "      landmarks_dict[item_id] = landmarks\n",
    "  \n",
    "  landmark_item_ids = sorted(landmarks_dict)\n",
    "  landmarks_pixels = np.array([landmarks_dict[i] for i in landmark_item_ids])\n",
    "  landmarks_cameras = [scene_manager.camera_dict[i] for i in landmark_item_ids]\n",
    "  \n",
    "  from matplotlib import pyplot as plt\n",
    "  plt.imshow(image)\n",
    "  plt.scatter(x=landmarks[..., 0], y=landmarks[..., 1], s=1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "axRj1ItALAuC"
   },
   "outputs": [],
   "source": [
    "# @title Triangulate landmarks in 3D.\n",
    "\n",
    "if use_face:\n",
    "  def compute_camera_rays(points, camera):\n",
    "    origins = np.broadcast_to(camera.position[None, :], (points.shape[0], 3))\n",
    "    directions = camera.pixels_to_rays(points.astype(jnp.float32))\n",
    "    endpoints = origins + directions\n",
    "    return origins, endpoints\n",
    "  \n",
    "  \n",
    "  def triangulate_landmarks(landmarks, cameras):\n",
    "    all_origins = []\n",
    "    all_endpoints = []\n",
    "    nan_inds = []\n",
    "    for i, (camera_landmarks, camera) in enumerate(zip(landmarks, cameras)):\n",
    "      origins, endpoints = compute_camera_rays(camera_landmarks, camera)\n",
    "      if np.isnan(origins).sum() > 0.0 or np.isnan(endpoints).sum() > 0.0:\n",
    "        continue\n",
    "      all_origins.append(origins)\n",
    "      all_endpoints.append(endpoints)\n",
    "    all_origins = np.stack(all_origins, axis=-2).astype(np.float32)\n",
    "    all_endpoints = np.stack(all_endpoints, axis=-2).astype(np.float32)\n",
    "    weights = np.ones(all_origins.shape[:2], dtype=np.float32)\n",
    "    points = np.array(ray_triangulate(all_origins, all_endpoints, weights))\n",
    "  \n",
    "    return points\n",
    "  \n",
    "\n",
    "  landmark_points = triangulate_landmarks(landmarks_pixels, landmarks_cameras)\n",
    "else:\n",
    "  landmark_points = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "gRU-bJ8NYzR_"
   },
   "outputs": [],
   "source": [
    "# @title Normalize scene based on landmarks.\n",
    "from scipy import linalg\n",
    "\n",
    "DEFAULT_IPD = 0.06\n",
    "NOSE_TIP_IDX = 1\n",
    "FOREHEAD_IDX = 10\n",
    "CHIN_IDX = 152\n",
    "RIGHT_EYE_IDX = 145\n",
    "LEFT_EYE_IDX = 385\n",
    "RIGHT_TEMPLE_IDX = 162\n",
    "LEFT_TEMPLE_IDX = 389\n",
    "\n",
    "\n",
    "def _normalize(x):\n",
    "  return x / linalg.norm(x)\n",
    "\n",
    "\n",
    "def fit_plane_normal(points):\n",
    "  \"\"\"Fit a plane to the points and return the normal.\"\"\"\n",
    "  centroid = points.sum(axis=0) / points.shape[0]\n",
    "  _, _, vh = linalg.svd(points - centroid)\n",
    "  return vh[2, :]\n",
    "\n",
    "\n",
    "def metric_scale_from_ipd(landmark_points, reference_ipd):\n",
    "  \"\"\"Infer the scene-to-metric conversion ratio from facial landmarks.\"\"\"\n",
    "  left_eye = landmark_points[LEFT_EYE_IDX]\n",
    "  right_eye = landmark_points[RIGHT_EYE_IDX]\n",
    "  model_ipd = linalg.norm(left_eye - right_eye)\n",
    "  return reference_ipd / model_ipd\n",
    "\n",
    "\n",
    "def basis_from_landmarks(landmark_points):\n",
    "  \"\"\"Computes an orthonormal basis from facial landmarks.\"\"\"\n",
    "  # Estimate Z by fitting a plane\n",
    "  # This works better than trusting the chin to forehead vector, especially in\n",
    "  # full body captures.\n",
    "  face_axis_z = _normalize(fit_plane_normal(landmark_points))\n",
    "  face_axis_y = _normalize(landmark_points[FOREHEAD_IDX] -\n",
    "                           landmark_points[CHIN_IDX])\n",
    "  face_axis_x = _normalize(landmark_points[LEFT_TEMPLE_IDX] -\n",
    "                           landmark_points[RIGHT_TEMPLE_IDX])\n",
    "\n",
    "  # Fitted plane normal might be flipped. Check using a heuristic and flip it if\n",
    "  # it's flipped.\n",
    "  z_flipped = np.dot(np.cross(face_axis_x, face_axis_y), face_axis_z)\n",
    "  if z_flipped < 0.0:\n",
    "    face_axis_z *= -1\n",
    "\n",
    "  # Ensure axes are orthogonal, with the Z axis being fixed.\n",
    "  face_axis_y = np.cross(face_axis_z, face_axis_x)\n",
    "  face_axis_x = np.cross(face_axis_y, face_axis_z)\n",
    "\n",
    "  return np.stack([face_axis_x, face_axis_y, face_axis_z]).T\n",
    "\n",
    "\n",
    "if use_face:\n",
    "  face_basis = basis_from_landmarks(landmark_points)\n",
    "  new_scene_manager = scene_manager.change_basis(\n",
    "      face_basis, landmark_points[NOSE_TIP_IDX])\n",
    "  new_cameras = [new_scene_manager.camera_dict[i] for i in landmark_item_ids]\n",
    "  new_landmark_points = triangulate_points(landmarks_pixels, new_cameras)\n",
    "  face_basis = basis_from_landmarks(landmark_points)\n",
    "  scene_to_metric = metric_scale_from_ipd(landmark_points, DEFAULT_IPD)\n",
    "  \n",
    "  print(f'Computed basis: {face_basis}')\n",
    "  print(f'Estimated metric scale = {scene_to_metric:.02f}')\n",
    "else:\n",
    "  new_scene_manager = scene_manager"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iPuR5MKk6Ubh"
   },
   "source": [
    "## Compute scene information.\n",
    "\n",
    "This section computes the scene information necessary for NeRF training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "klgXn8BQ8uH9"
   },
   "outputs": [],
   "source": [
    "# @title Compute near/far planes.\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "def estimate_near_far_for_image(scene_manager, image_id):\n",
    "  \"\"\"Estimate near/far plane for a single image based via point cloud.\"\"\"\n",
    "  points = filter_outlier_points(scene_manager.points, 0.95)\n",
    "  points = np.concatenate([\n",
    "      points,\n",
    "      scene_manager.camera_positions,\n",
    "  ], axis=0)\n",
    "  camera = scene_manager.camera_dict[image_id]\n",
    "  pixels = camera.project(points)\n",
    "  depths = camera.points_to_local_points(points)[..., 2]\n",
    "\n",
    "  # in_frustum = camera.ArePixelsInFrustum(pixels)\n",
    "  in_frustum = (\n",
    "      (pixels[..., 0] >= 0.0)\n",
    "      & (pixels[..., 0] <= camera.image_size_x)\n",
    "      & (pixels[..., 1] >= 0.0)\n",
    "      & (pixels[..., 1] <= camera.image_size_y))\n",
    "  depths = depths[in_frustum]\n",
    "\n",
    "  in_front_of_camera = depths > 0\n",
    "  depths = depths[in_front_of_camera]\n",
    "\n",
    "  near = np.quantile(depths, 0.001)\n",
    "  far = np.quantile(depths, 0.999)\n",
    "\n",
    "  return near, far\n",
    "\n",
    "\n",
    "def estimate_near_far(scene_manager):\n",
    "  \"\"\"Estimate near/far plane for a set of randomly-chosen images.\"\"\"\n",
    "  # image_ids = sorted(scene_manager.images.keys())\n",
    "  image_ids = scene_manager.image_ids\n",
    "  rng = np.random.RandomState(0)\n",
    "  image_ids = rng.choice(\n",
    "      image_ids, size=len(scene_manager.camera_list), replace=False)\n",
    "  \n",
    "  result = []\n",
    "  for image_id in image_ids:\n",
    "    near, far = estimate_near_far_for_image(scene_manager, image_id)\n",
    "    result.append({'image_id': image_id, 'near': near, 'far': far})\n",
    "  result = pd.DataFrame.from_records(result)\n",
    "  return result\n",
    "\n",
    "\n",
    "near_far = estimate_near_far(new_scene_manager)\n",
    "print('Statistics for near/far computation:')\n",
    "print(near_far.describe())\n",
    "print()\n",
    "\n",
    "near = near_far['near'].quantile(0.001) / 0.8\n",
    "far = near_far['far'].quantile(0.999) * 1.2\n",
    "print('Selected near/far values:')\n",
    "print(f'Near = {near:.04f}')\n",
    "print(f'Far = {far:.04f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "kOgCoT62ArbD"
   },
   "outputs": [],
   "source": [
    "# @title Compute scene center and scale.\n",
    "\n",
    "def get_bbox_corners(points):\n",
    "  lower = points.min(axis=0)\n",
    "  upper = points.max(axis=0)\n",
    "  return np.stack([lower, upper])\n",
    "\n",
    "\n",
    "points = filter_outlier_points(new_scene_manager.points, 0.95)\n",
    "bbox_corners = get_bbox_corners(\n",
    "    np.concatenate([points, new_scene_manager.camera_positions], axis=0))\n",
    "\n",
    "scene_center = np.mean(bbox_corners, axis=0)\n",
    "scene_scale = 1.0 / np.sqrt(np.sum((bbox_corners[1] - bbox_corners[0]) ** 2))\n",
    "\n",
    "print(f'Scene Center: {scene_center}')\n",
    "print(f'Scene Scale: {scene_scale}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "6Q1KC4xw6Til"
   },
   "outputs": [],
   "source": [
    "# @title Visualize scene.\n",
    "\n",
    "def scatter_points(points, size=2):\n",
    "  return go.Scatter3d(\n",
    "    x=points[:, 0],\n",
    "    y=points[:, 1],\n",
    "    z=points[:, 2],\n",
    "    mode='markers',\n",
    "    marker=dict(size=size),\n",
    "  )\n",
    "\n",
    "camera = new_scene_manager.camera_list[0]\n",
    "near_points = camera.pixels_to_points(\n",
    "    camera.get_pixel_centers()[::8, ::8], jnp.array(near)).reshape((-1, 3))\n",
    "far_points = camera.pixels_to_points(\n",
    "    camera.get_pixel_centers()[::8, ::8], jnp.array(far)).reshape((-1, 3))\n",
    "\n",
    "data = [\n",
    "  scatter_points(new_scene_manager.points),\n",
    "  scatter_points(new_scene_manager.camera_positions),\n",
    "  scatter_points(bbox_corners),\n",
    "  scatter_points(near_points),\n",
    "  scatter_points(far_points),\n",
    "]\n",
    "if use_face:\n",
    "  data.append(scatter_points(new_landmark_points))\n",
    "fig = go.Figure(data=data)\n",
    "fig.update_layout(scene_dragmode='orbit')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KtOTEI_Tbpt_"
   },
   "source": [
    "## Generate test cameras."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "WvvOLabUeJUX"
   },
   "outputs": [],
   "source": [
    "# @title Define Utilities.\n",
    "_EPSILON = 1e-5\n",
    "\n",
    "\n",
    "def points_bound(points):\n",
    "  \"\"\"Computes the min and max dims of the points.\"\"\"\n",
    "  min_dim = np.min(points, axis=0)\n",
    "  max_dim = np.max(points, axis=0)\n",
    "  return np.stack((min_dim, max_dim), axis=1)\n",
    "\n",
    "\n",
    "def points_centroid(points):\n",
    "  \"\"\"Computes the centroid of the points from the bounding box.\"\"\"\n",
    "  return points_bound(points).mean(axis=1)\n",
    "\n",
    "\n",
    "def points_bounding_size(points):\n",
    "  \"\"\"Computes the bounding size of the points from the bounding box.\"\"\"\n",
    "  bounds = points_bound(points)\n",
    "  return np.linalg.norm(bounds[:, 1] - bounds[:, 0])\n",
    "\n",
    "\n",
    "def look_at(camera,\n",
    "            camera_position: np.ndarray,\n",
    "            look_at_position: np.ndarray,\n",
    "            up_vector: np.ndarray):\n",
    "  look_at_camera = camera.copy()\n",
    "  optical_axis = look_at_position - camera_position\n",
    "  norm = np.linalg.norm(optical_axis)\n",
    "  if norm < _EPSILON:\n",
    "    raise ValueError('The camera center and look at position are too close.')\n",
    "  optical_axis /= norm\n",
    "\n",
    "  right_vector = np.cross(optical_axis, up_vector)\n",
    "  norm = np.linalg.norm(right_vector)\n",
    "  if norm < _EPSILON:\n",
    "    raise ValueError('The up-vector is parallel to the optical axis.')\n",
    "  right_vector /= norm\n",
    "\n",
    "  # The three directions here are orthogonal to each other and form a right\n",
    "  # handed coordinate system.\n",
    "  camera_rotation = np.identity(3)\n",
    "  camera_rotation[0, :] = right_vector\n",
    "  camera_rotation[1, :] = np.cross(optical_axis, right_vector)\n",
    "  camera_rotation[2, :] = optical_axis\n",
    "\n",
    "  look_at_camera.position = camera_position\n",
    "  look_at_camera.orientation = camera_rotation\n",
    "  return look_at_camera\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "e5cHTuhP9Dgp"
   },
   "outputs": [],
   "source": [
    "# @title Generate camera trajectory.\n",
    "\n",
    "import math\n",
    "from scipy import interpolate\n",
    "from plotly.offline import iplot\n",
    "import plotly.graph_objs as go\n",
    "\n",
    "\n",
    "def compute_camera_rays(points, camera):\n",
    "  origins = np.broadcast_to(camera.position[None, :], (points.shape[0], 3))\n",
    "  directions = camera.pixels_to_rays(points.astype(jnp.float32))\n",
    "  endpoints = origins + directions\n",
    "  return origins, endpoints\n",
    "\n",
    "\n",
    "def triangulate_rays(origins, directions):\n",
    "  origins = origins[np.newaxis, ...]\n",
    "  directions = directions[np.newaxis, ...]\n",
    "  weights = np.ones(origins.shape[:2], dtype=np.float32)\n",
    "  points = np.array(ray_triangulate(origins, origins + directions, weights))\n",
    "  return points.squeeze()\n",
    "\n",
    "\n",
    "ref_cameras = [c for c in new_scene_manager.camera_list]\n",
    "origins = np.array([c.position for c in ref_cameras])\n",
    "directions = np.array([c.optical_axis for c in ref_cameras])\n",
    "look_at = triangulate_rays(origins, directions)\n",
    "print('look_at', look_at)\n",
    "\n",
    "avg_position = np.mean(origins, axis=0)\n",
    "print('avg_position', avg_position)\n",
    "\n",
    "up = -np.mean([c.orientation[..., 1] for c in ref_cameras], axis=0)\n",
    "print('up', up)\n",
    "\n",
    "bounding_size = points_bounding_size(origins) / 2\n",
    "x_scale =   0.75# @param {type: 'number'}\n",
    "y_scale = 0.75  # @param {type: 'number'}\n",
    "xs = x_scale * bounding_size\n",
    "ys = y_scale * bounding_size\n",
    "radius = 0.75  # @param {type: 'number'}\n",
    "num_frames = 100  # @param {type: 'number'}\n",
    "\n",
    "\n",
    "origin = np.zeros(3)\n",
    "\n",
    "ref_camera = ref_cameras[0]\n",
    "print(ref_camera.position)\n",
    "z_offset = -0.1\n",
    "\n",
    "angles = np.linspace(0, 2*math.pi, num=num_frames)\n",
    "positions = []\n",
    "for angle in angles:\n",
    "  x = np.cos(angle) * radius * xs\n",
    "  y = np.sin(angle) * radius * ys\n",
    "  # x = xs * radius * np.cos(angle) / (1 + np.sin(angle) ** 2)\n",
    "  # y = ys * radius * np.sin(angle) * np.cos(angle) / (1 + np.sin(angle) ** 2)\n",
    "\n",
    "  position = np.array([x, y, z_offset])\n",
    "  # Make distance to reference point constant.\n",
    "  position = avg_position + position\n",
    "  positions.append(position)\n",
    "\n",
    "positions = np.stack(positions)\n",
    "\n",
    "orbit_cameras = []\n",
    "for position in positions:\n",
    "  camera = ref_camera.look_at(position, look_at, up)\n",
    "  orbit_cameras.append(camera)\n",
    "\n",
    "camera_paths = {'orbit-mild': orbit_cameras}\n",
    "\n",
    "traces = [\n",
    "  scatter_points(new_scene_manager.points),\n",
    "  scatter_points(new_scene_manager.camera_positions),\n",
    "  scatter_points(bbox_corners),\n",
    "  scatter_points(near_points),\n",
    "  scatter_points(far_points),\n",
    "\n",
    "  scatter_points(positions),\n",
    "  scatter_points(origins),\n",
    "]\n",
    "fig = go.Figure(traces)\n",
    "fig.update_layout(scene_dragmode='orbit')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UYJ6aI45IIwd"
   },
   "source": [
    "## Save data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "aDFYTpGB6_Gl"
   },
   "outputs": [],
   "source": [
    "# @title Save scene information to `scene.json`.\n",
    "from pprint import pprint\n",
    "import json\n",
    "\n",
    "scene_json_path = root_dir /  'scene.json'\n",
    "with scene_json_path.open('w') as f:\n",
    "  json.dump({\n",
    "      'scale': scene_scale,\n",
    "      'center': scene_center.tolist(),\n",
    "      'bbox': bbox_corners.tolist(),\n",
    "      'near': near * scene_scale,\n",
    "      'far': far * scene_scale,\n",
    "  }, f, indent=2)\n",
    "\n",
    "print(f'Saved scene information to {scene_json_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "k_oQ-4MTGFpz"
   },
   "outputs": [],
   "source": [
    "# @title Save dataset split to `dataset.json`.\n",
    "\n",
    "all_ids = scene_manager.image_ids\n",
    "val_ids = all_ids[::20]\n",
    "train_ids = sorted(set(all_ids) - set(val_ids))\n",
    "dataset_json = {\n",
    "    'count': len(scene_manager),\n",
    "    'num_exemplars': len(train_ids),\n",
    "    'ids': scene_manager.image_ids,\n",
    "    'train_ids': train_ids,\n",
    "    'val_ids': val_ids,\n",
    "}\n",
    "\n",
    "dataset_json_path = root_dir / 'dataset.json'\n",
    "with dataset_json_path.open('w') as f:\n",
    "    json.dump(dataset_json, f, indent=2)\n",
    "\n",
    "print(f'Saved dataset information to {dataset_json_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "3PWkPkBVGnSl"
   },
   "outputs": [],
   "source": [
    "# @title Save metadata information to `metadata.json`.\n",
    "import bisect\n",
    "\n",
    "metadata_json = {}\n",
    "for i, image_id in enumerate(train_ids):\n",
    "  metadata_json[image_id] = {\n",
    "      'warp_id': i,\n",
    "      'appearance_id': i,\n",
    "      'camera_id': 0,\n",
    "  }\n",
    "for i, image_id in enumerate(val_ids):\n",
    "  i = bisect.bisect_left(train_ids, image_id)\n",
    "  metadata_json[image_id] = {\n",
    "      'warp_id': i,\n",
    "      'appearance_id': i,\n",
    "      'camera_id': 0,\n",
    "  }\n",
    "\n",
    "metadata_json_path = root_dir / 'metadata.json'\n",
    "with metadata_json_path.open('w') as f:\n",
    "    json.dump(metadata_json, f, indent=2)\n",
    "\n",
    "print(f'Saved metadata information to {metadata_json_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "4Uxu0yKlGs3V"
   },
   "outputs": [],
   "source": [
    "# @title Save cameras.\n",
    "camera_dir = root_dir / 'camera'\n",
    "camera_dir.mkdir(exist_ok=True, parents=True)\n",
    "for item_id, camera in new_scene_manager.camera_dict.items():\n",
    "  camera_path = camera_dir / f'{item_id}.json'\n",
    "  print(f'Saving camera to {camera_path!s}')\n",
    "  with camera_path.open('w') as f:\n",
    "    json.dump(camera.to_json(), f, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WA_Icz5_Ia4h"
   },
   "outputs": [],
   "source": [
    "# @title Save test cameras.\n",
    "\n",
    "import json\n",
    "\n",
    "test_camera_dir = root_dir / 'camera-paths'\n",
    "for test_path_name, test_cameras in camera_paths.items():\n",
    "  out_dir = test_camera_dir / test_path_name\n",
    "  out_dir.mkdir(exist_ok=True, parents=True)\n",
    "  for i, camera in enumerate(test_cameras):\n",
    "    camera_path = out_dir / f'{i:06d}.json'\n",
    "    print(f'Saving camera to {camera_path!s}')\n",
    "    with camera_path.open('w') as f:\n",
    "      json.dump(camera.to_json(), f, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimate depths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import torch\n",
    "import tqdm\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "model_type = \"DPT_Large\"     # MiDaS v3 - Large     (highest accuracy, slowest inference speed)\n",
    "#model_type = \"DPT_Hybrid\"   # MiDaS v3 - Hybrid    (medium accuracy, medium inference speed)\n",
    "#model_type = \"MiDaS_small\"  # MiDaS v2.1 - Small   (lowest accuracy, highest inference speed)\n",
    "\n",
    "midas = torch.hub.load(\"intel-isl/MiDaS\", model_type)\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "midas.to(device)\n",
    "midas.eval()\n",
    "\n",
    "midas_transforms = torch.hub.load(\"intel-isl/MiDaS\", \"transforms\")\n",
    "\n",
    "if model_type == \"DPT_Large\" or model_type == \"DPT_Hybrid\":\n",
    "    transform = midas_transforms.dpt_transform\n",
    "else:\n",
    "    transform = midas_transforms.small_transform\n",
    "    \n",
    "img_files = list((Path(rgb_dir) / \"1x\").rglob(\"*.png\"))\n",
    "\n",
    "img_sizes = [\n",
    "  cv2.imread(list((Path(rgb_dir) / f\"{scale}x\").glob(\"*.png\"))[0].as_posix()).shape[:-1]\n",
    "  for scale in image_scales\n",
    "]\n",
    "print(image_scales, img_sizes)\n",
    "for img_file in tqdm.tqdm(img_files):\n",
    "  img = cv2.imread(img_file.as_posix())\n",
    "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "  input_batch = transform(img).to(device)\n",
    "  with torch.no_grad():\n",
    "      prediction = midas(input_batch)\n",
    "      for size, scale in zip(img_sizes, image_scales):\n",
    "        scaled_pred = torch.nn.functional.interpolate(\n",
    "            prediction.unsqueeze(1),\n",
    "            size=size,\n",
    "            mode=\"bicubic\",\n",
    "            align_corners=False,\n",
    "        ).squeeze()\n",
    "\n",
    "        output = scaled_pred.cpu().numpy()\n",
    "        np.save(\n",
    "          (Path(rgb_dir) / f\"{scale}x\" / (img_file.stem + \"_depth\")).with_suffix(\".npy\"),\n",
    "          output\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data annotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "num_classes = None\n",
    "frame_num = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes_slider = widgets.IntSlider(\n",
    "  min=0,\n",
    "  max=100,\n",
    "  value=0,\n",
    "  description=\"Number of classes:\",\n",
    ")\n",
    "\n",
    "set_num_classes_button = widgets.interact_manual.options(manual_name=\"Set num of classes\")\n",
    "@set_num_classes_button(num_cls=num_classes_slider)\n",
    "def assign_num_classes(num_cls):\n",
    "  global num_classes\n",
    "  num_classes = num_cls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import json\n",
    "\n",
    "available_image_ids = json.loads((root_dir / \"dataset.json\").read_text())[\"ids\"]\n",
    "\n",
    "images = sorted(\n",
    "  [(rgb_raw_dir / an_id).with_suffix(\".png\") for an_id in available_image_ids],\n",
    "  key=lambda x: int(x.with_suffix(\"\").name),\n",
    ")\n",
    "image_index_slider = widgets.IntSlider(\n",
    "  min=0,\n",
    "  max=len(images) - 1,\n",
    "  value=0,\n",
    "  description=\"Index:\",\n",
    ")\n",
    "@widgets.interact(index=image_index_slider)\n",
    "def show_image(index):\n",
    "  global frame_num\n",
    "  frame_num = int(images[index].with_suffix(\"\").name)\n",
    "  print(frame_num)\n",
    "  image = Image.open(images[index].as_posix())\n",
    "  return image.resize( [int(0.1 * s) for s in image.size] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "import json\n",
    "\n",
    "annotations = {}\n",
    "value_sliders = OrderedDict(\n",
    "    {\n",
    "      f\"attr_{i}\": widgets.FloatSlider(\n",
    "        min=-1.0,\n",
    "        max=1.0,\n",
    "        value=0.0,\n",
    "        step=1.0,\n",
    "        description=f\"Attribute #{i}\",\n",
    "      )\n",
    "      for i in range(num_classes)\n",
    "  } \n",
    ")\n",
    "add_annotation_btn = widgets.interact_manual.options(manual_name=\"Add annotation\")\n",
    "@add_annotation_btn(**value_sliders)\n",
    "def add_attribute(**attribute_values):\n",
    "  annotations[frame_num] = [val for val in attribute_values.values()]\n",
    "  print(\"Added annotation: {}\".format(\",\".join([str(val) for val in attribute_values.values()])))\n",
    "  print(\"Current annotations:\")\n",
    "  print(json.dumps(annotations))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "\n",
    "save_annotation_btn = widgets.interact_manual.options(manual_name=\"Save annotation\")\n",
    "@save_annotation_btn()\n",
    "def save_annotation():\n",
    "  data = []\n",
    "  for key, values in annotations.items():\n",
    "    for index, value in enumerate(values):\n",
    "      data.append(\n",
    "        {\n",
    "          \"frame\": key,\n",
    "          \"class\": index,\n",
    "          \"value\": value,\n",
    "        }\n",
    "      )\n",
    "  with open(annotations_file.as_posix(), \"w\") as f:\n",
    "    yaml.safe_dump(data, f)\n",
    "    \n",
    "  !cat \"$annotations_file\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mapping = {}\n",
    "name_sliders = OrderedDict(\n",
    "    {\n",
    "      f\"attr_{i}\": widgets.Text(\n",
    "        description=f\"Attribute #{i} name\",\n",
    "      )\n",
    "      for i in range(num_classes)\n",
    "  } \n",
    ")\n",
    "add_mapping_btn = widgets.interact_manual.options(manual_name=\"Add mapping\")\n",
    "@add_mapping_btn(**name_sliders)\n",
    "def add_mapping(**attribute_names):\n",
    "  global mapping\n",
    "  for i, value in enumerate(attribute_names.values()):\n",
    "    mapping[i] = value\n",
    "  print(\"Current mapping:\")\n",
    "  print(json.dumps(mapping))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "\n",
    "save_mapping_btn = widgets.interact_manual.options(manual_name=\"Save label mapping\")\n",
    "@save_mapping_btn()\n",
    "def save_mapping():\n",
    "  with open(mapping_file.as_posix(), \"w\") as f:\n",
    "    yaml.safe_dump(mapping, f)\n",
    "    \n",
    "  !cat \"$mapping_file\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb_raw_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annotations_dir.mkdir(exist_ok=True)\n",
    "\n",
    "with open(annotations_file) as f:\n",
    "  data = yaml.safe_load(f)\n",
    "  \n",
    "annotations = {}\n",
    "for entry in data:\n",
    "  if entry[\"frame\"] not in annotations:\n",
    "    annotations[entry[\"frame\"]] = [0.0] * len([subentry for subentry in data if subentry[\"frame\"] == entry[\"frame\"]])\n",
    "  annotations[entry[\"frame\"]][entry[\"class\"]] = entry[\"value\"]\n",
    "  \n",
    "if len(list(annotations_dir.rglob(\"*.json\"))) != len(annotations.keys()):\n",
    "  raise RuntimeError(\"Requires mask for each annotation\")\n",
    "print(\"Everything ok!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3iV-YLB_TEMq"
   },
   "source": [
    "## Training\n",
    "\n",
    " * You are now ready to train a Nerfie!\n",
    " * Head over to the [training Colab](https://colab.sandbox.google.com/github/google/nerfies/blob/main/notebooks/Nerfies_Training.ipynb) for a basic demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!which python"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Nerfies Capture Processing v2.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}