notebooks/NeRS In-the-wild Demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77a10a73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "import sys\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import pytorch3d as p3d\n",
    "\n",
    "if sys.path[0] != \"..\":\n",
    "    sys.path.insert(0, \"..\")\n",
    "\n",
    "from ners import Ners, load_data_from_dir\n",
    "from ners.pytorch3d import PerspectiveCameras, get_renderers\n",
    "from ners.models import pretrain_template_uv, shape_model_to_mesh, TemplateUV"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b73cd19f",
   "metadata": {},
   "source": [
    "# Demo Notebook for NeRS\n",
    "\n",
    "This notebook generates a NeRS model given images, masks, a cuboid initialization (user specified), and coarse poses (user specified).\n",
    "\n",
    "\n",
    "Compute requirements: ~3-4 GPUs\n",
    "\n",
    "Data requirements:\n",
    "* Images: try to take 8-10 images, roughly in a turntable, under consistent illumination.\n",
    "* Masks: you can use your favorite segmentation tool, such as [PointRend](https://github.com/facebookresearch/detectron2/tree/main/projects/PointRend) or [my interactive interface for GrabCut](https://github.com/jasonyzhang/interactive_grabcut).\n",
    "\n",
    "File structure:\n",
    "```\n",
    "INPUT_DIR\n",
    "|_ images\n",
    "|___ image1.jpg\n",
    "|___ ...\n",
    "|_ masks\n",
    "|___ image1.png  <- should have same filename as the images\n",
    "|___ ...\n",
    "```\n",
    "\n",
    "Manual Input:\n",
    "* Extents of the cuboid initialization\n",
    "* Roughly binned camera poses (azimuth and elevation)\n",
    "\n",
    "Lines that require manual input will have a `# TODO: Set this.`\n",
    "\n",
    "Timings on 4x1080 TIs:\n",
    "* Stage 1: 1.5 min\n",
    "* Stage 2: 1.5 min\n",
    "* Stage 3: 30 min\n",
    "* Stage 4: 10 min\n",
    "* Render a video: 4 min"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec17f4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\")\n",
    "\n",
    "INPUT_DIR = \"../data/espresso\"  # TODO: Set this.\n",
    "OUTPUT_DIR = osp.join(\"../output\", osp.basename(osp.normpath(INPUT_DIR)))\n",
    "\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "renderer_textured, renderer_silhouette = get_renderers(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d58f238",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the images and masks.\n",
    "\n",
    "data = load_data_from_dir(INPUT_DIR)\n",
    "\n",
    "n = len(data[\"images\"])\n",
    "num_rows = (n - 1) // 2 + 1\n",
    "fig, axs = plt.subplots(num_rows, 4, figsize=(8, num_rows * 2))\n",
    "axs = axs.flatten()\n",
    "for i in range(n):\n",
    "    image = data[\"images\"][i]\n",
    "    mask = data[\"masks\"][i]\n",
    "    axs[i * 2].imshow(image)\n",
    "    axs[i * 2].set_title(f\"Image {i}\")\n",
    "    axs[i * 2 + 1].imshow(mask)\n",
    "    axs[i * 2 + 1].set_title(f\"Mask {i}\")\n",
    "for ax in axs:\n",
    "    ax.axis(\"off\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e3947c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Set this.\n",
    "# Try to keep the largest dimension around 1.\n",
    "template_extents = [0.7, 1, 1]  # [W, H, D].\n",
    "\n",
    "\n",
    "# Pre-train the template shape.\n",
    "f_template = pretrain_template_uv(\n",
    "    template_uv=TemplateUV(),\n",
    "    extents=template_extents,\n",
    "    device=device,\n",
    ")\n",
    "template_mesh = shape_model_to_mesh(f_template)\n",
    "\n",
    "# Save template shape.\n",
    "template_path = osp.join(OUTPUT_DIR, \"template.pth\")\n",
    "torch.save(f_template.state_dict(), template_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7efd4b54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Set this.\n",
    "azimuths = [0, 60, 270, 190, 210, 90, 270]  # In Degrees.\n",
    "# TODO: Set this.\n",
    "elevations = [30, 30, 0, 30, 30, 0, 90]  # In Degrees.\n",
    "\n",
    "R, T = p3d.renderer.look_at_view_transform(\n",
    "    dist=2,\n",
    "    elev=elevations,\n",
    "    azim=azimuths,\n",
    "    device=device,\n",
    ")\n",
    "cameras = PerspectiveCameras(\n",
    "    device=device,\n",
    "    R=R,\n",
    "    T=T,\n",
    "    fov=60,\n",
    ")\n",
    "\n",
    "N = len(cameras)\n",
    "\n",
    "rend = renderer_textured(template_mesh.extend(N), cameras=cameras)\n",
    "fig, axs = plt.subplots(N, 2, figsize=(4, N * 2))\n",
    "for i, (im, r) in enumerate(zip(data[\"images\"], rend.detach().cpu()[..., :3])):\n",
    "    axs[i, 0].imshow(im)\n",
    "    axs[i, 0].set_title(f\"Image {i}\")\n",
    "    axs[i, 1].imshow(r)\n",
    "for ax in axs.flatten():\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f7cc0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ners = Ners(\n",
    "    images=data[\"images\"],\n",
    "    masks=data[\"masks\"],\n",
    "    masks_dt=data[\"masks_dt\"],\n",
    "    initial_poses=R.tolist(),\n",
    "    image_center=data[\"image_centers\"],\n",
    "    crop_scale=data[\"crop_scales\"],\n",
    "    f_template=f_template,\n",
    "    symmetrize=True,  # TODO: Set this to False for non-symmetric objects.\n",
    "    num_layers_tex=12,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01496b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ners.visualize_input_views(title=\"Initialization\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83022032",
   "metadata": {},
   "outputs": [],
   "source": [
    "ners.optimize_camera()\n",
    "ners.visualize_input_views(title=\"Stage 1: Optimize Camera\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcd67284",
   "metadata": {},
   "outputs": [],
   "source": [
    "ners.optimize_shape()\n",
    "ners.visualize_input_views(title=\"Stage 2: Optimize $f_{shape}$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d771f95",
   "metadata": {},
   "outputs": [],
   "source": [
    "ners.optimize_texture(3000)\n",
    "ners.visualize_input_views(title=\"Stage 3: Optimize $f_{tex}$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bb9a54a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Output 3D Mesh\n",
    "mesh_path = osp.join(OUTPUT_DIR, \"mesh.obj\")\n",
    "ners.save_obj(mesh_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed2e212",
   "metadata": {},
   "outputs": [],
   "source": [
    "ners.optimize_radiance()\n",
    "ners.visualize_input_views(title=\"Stage 4: Optimize $f_{env}$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b795224e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Output 360 degree video.\n",
    "video_path = osp.join(OUTPUT_DIR, \"video\")\n",
    "ners.make_video(video_path, image_size=512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6142cb66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Video\n",
    "\n",
    "Video(video_path + \".mp4\", embed=True, width=600)"
   ]
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}