nerf_quick_start.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NeRF Quick Start Notebook\n",
    "\n",
    "Plank-ing Hyundong 3D Reconstruction Project\n",
    "Created 2022.07.05 <br>\n",
    "\n",
    "There are 4 Steps in this notebook.<br>\n",
    "1. Sampling images from video\n",
    "2. Get camera poses with COLMAP\n",
    "3. Run NeRF\n",
    "4. Get mesh file\n",
    "\n",
    "\n",
    "__If you clone this repository in your local PC, set the directory path at first__\n",
    "\n",
    "__You must use the GPU runtime.__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📚 Directory Structure\n",
    "\n",
    "- PlankHyundong\n",
    "\n",
    "    - __nerf_quick_start.ipynb__\n",
    "\n",
    "    - notebook\n",
    "        - nerf_colab.ipynb\n",
    "        - nerf_wandb_colab.ipynb\n",
    "        - colmap_colab.ipynb\n",
    "        - extract_mesh_colab.ipynb\n",
    "        - sampling_colab.ipynb\n",
    "\n",
    "    - data\n",
    "        - video\n",
    "            - video.MOV\n",
    "        - (images)\n",
    "            - ..\n",
    "        - (logs)\n",
    "            - ..\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__We use tensorflow ver.1__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    %tensorflow_version 1.x\n",
    "except ValueError:\n",
    "    # 만약 %tensorflow_version 1.x magic 명령어가 작동하지 않는 경우\n",
    "    !pip uninstall --yes tensorflow\n",
    "    !pip install tensorflow==1.15\n",
    "    import tensorflow\n",
    "    print(tensorflow.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting directory path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# path to data folder\n",
    "video_path = '/content/data/video/video.MOV'\n",
    "img_path = '/content/data/images'\n",
    "logs_path = '/content/data/logs'\n",
    "path = '/content/data'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1\n",
    "### Video Sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- load data folder from PlankHyundong repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git init\n",
    "!git remote add origin https://github.com/ProtossDragoon/PlankHyundong\n",
    "!git config core.sparsecheckout true\n",
    "!echo data/* >> ./.git/info/sparse-checkout\n",
    "!git pull origin main\n",
    "!rm -rf ./git\n",
    "\n",
    "import os\n",
    "\n",
    "os.mkdir(img_path)\n",
    "os.mkdir(logs_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "# Set the number of frame\n",
    "frame = 50\n",
    "\n",
    "vidcap = cv2.VideoCapture(video_path)\n",
    "                \n",
    "cnt, num = 0, 1 # cnt -> Input frame #, num -> output Frame #.\n",
    "\n",
    "total_length = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "cycle = int(total_length / frame) # calculate cycle\n",
    "\n",
    "while vidcap.isOpened():\n",
    "    ret,image = vidcap.read()\n",
    "    if num > frame:\n",
    "        break\n",
    "    if ret and cnt % cycle == 0:  \n",
    "        \n",
    "        try:\n",
    "            cv2.imwrite(f\"{img_path}/image{num}.jpg\", image)\n",
    "            num+=1\n",
    "        except:\n",
    "            print(\"fail\")\n",
    "            \n",
    "    cnt += 1\n",
    "    \n",
    "vidcap.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2\n",
    "### Run COLMAP to get camera pose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Install dependent packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "!sudo apt-get install \\\n",
    "    git \\\n",
    "    cmake \\\n",
    "    build-essential \\\n",
    "    libboost-program-options-dev \\\n",
    "    libboost-filesystem-dev \\\n",
    "    libboost-graph-dev \\\n",
    "    libboost-regex-dev \\\n",
    "    libboost-system-dev \\\n",
    "    libboost-test-dev \\\n",
    "    libeigen3-dev \\\n",
    "    libsuitesparse-dev \\\n",
    "    libfreeimage-dev \\\n",
    "    libgoogle-glog-dev \\\n",
    "    libgflags-dev \\\n",
    "    libglew-dev \\\n",
    "    qtbase5-dev \\\n",
    "    libqt5opengl5-dev \\\n",
    "    libcgal-dev \\\n",
    "    libcgal-qt5-dev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Install Ceres-solver\n",
    "<br><br>It takes 10 ~ 20 minutes.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo apt-get install libatlas-base-dev libsuitesparse-dev\n",
    "!git clone https://ceres-solver.googlesource.com/ceres-solver\n",
    "%cd ceres-solver\n",
    "!git checkout $(git describe --tags)\n",
    "%mkdir build\n",
    "%cd build\n",
    "!cmake .. -DBUILD_TESTING=OFF -DBUILD_EXAMPLES=OFF\n",
    "!make\n",
    "!sudo make install"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Install COLMAP\n",
    "<br><br>\n",
    "It takes 10 ~ 20 minutes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo apt-get install libmetis-dev # https://github.com/colmap/colmap/issues/1469"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/colmap/colmap\n",
    "%cd colmap\n",
    "!git checkout $(git describe --tags)\n",
    "%mkdir build\n",
    "%cd build\n",
    "!cmake ..\n",
    "!make\n",
    "!sudo make install\n",
    "!CC=/usr/bin/gcc-6 CXX=/usr/bin/g++-6 cmake .."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Run LLFF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content\n",
    "!git clone https://github.com/Fyusion/LLFF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/LLFF\n",
    "\n",
    "!python imgs2poses.py {path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3\n",
    "### Run NeRF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- install dependent-packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt-get update\n",
    "!sudo apt -qq install imagemagick\n",
    "!pip install ConfigArgParse -qqq\n",
    "!pip install imageio-ffmpeg -qqq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- clone NeRF source code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd ..\n",
    "!git clone https://github.com/yyongjae/nerf.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd nerf\n",
    "!ls -al"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Setting the Hyperparameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "now = datetime.now().strftime('%y%m%d_%H%M%S')\n",
    "\n",
    "max_iter = 30000\n",
    "dataset_name = f'hyundong360_{frame}'\n",
    "downsample_factor = 64 #@param {type:\"slider\", min:1, max:64, step:1}\n",
    "netdepth = 4 #@param {type:\"slider\", min:4, max:16, step:2}\n",
    "netwidth = 64 #@param {type:\"slider\", min:64, max:256, step:4}\n",
    "experiment_name = f'{dataset_name}_{downsample_factor}_downsampled_{now}'\n",
    "max_iter = 30000 #@param\n",
    "learning_rate = 0.01 #@param\n",
    "video_saving_cnt = 3 #@param {type:\"slider\", min:1, max:10, step:1}\n",
    "n_samples = 32 #@param {type:\"slider\", min:32, max:256, step:32}\n",
    "\n",
    "# fine 모델에서 사용되는 샘플 개수는 coarse 모델의 sampling 개수의 2배로 설정한다.\n",
    "# 공식 논문에서 제안하는 대로, 64이면 128.\n",
    "n_importance = n_samples * 2\n",
    "\n",
    "# Reproduce 를 위해 고정 random_seed 를 사용\n",
    "random_seed = 777 #@param\n",
    "\n",
    "# tradeoff: memory <-> speed (training 에는 속도와 성능 모두에 영향을 미치지 않음. 학습 도중 동영상을 만들 때 OOM 이 난다면 충분히 낮출 것)\n",
    "rendering_speed = 2048 #@param {type:\"slider\", min:1024, max:16384, step:1024}\n",
    "\n",
    "# tradeoff: memory <-> result\n",
    "n_points_per_ray = 65536 #@param {type:\"slider\", min:2048, max:262144, step:1024}\n",
    "\n",
    "\n",
    "print(f'experiment: {experiment_name}')\n",
    "\n",
    "\n",
    "# Make config.txt file.\n",
    "f = open(f\"{logs_path}/config.txt\", 'w')\n",
    "def make_config(a, b):\n",
    "  \n",
    "  if b is None:\n",
    "    f.write(a + '\\n')\n",
    "  else:\n",
    "    data = a + ' = ' + b + '\\n'\n",
    "    f.write(data)\n",
    "\n",
    "\n",
    "make_config(\"maxiter\", str(max_iter))\n",
    "make_config(\"datadir\", img_path)\n",
    "make_config(\"basedir\", logs_path)\n",
    "make_config(\"dataset_type\", 'llff')\n",
    "make_config('factor', str(downsample_factor))\n",
    "make_config('netdepth', str(netdepth))\n",
    "make_config('netwidth', str(netwidth))\n",
    "make_config('netdepth_fine', str(netdepth))\n",
    "make_config('netwidth_fine', str(netwidth))\n",
    "make_config('chunk', str(rendering_speed))\n",
    "make_config('netchunk', str(n_points_per_ray))\n",
    "make_config('lrate', str(learning_rate))\n",
    "make_config('i_video', str(max_iter // video_saving_cnt))\n",
    "make_config('expname', experiment_name)\n",
    "make_config('N_samples', str(n_samples))\n",
    "make_config('N_importance', str(n_importance))\n",
    "make_config('random_seed', str(random_seed))\n",
    "make_config('raw_noise_std', '1.0')\n",
    "make_config('use_viewdirs', None)\n",
    "make_config('no_ndc', None)\n",
    "make_config('spherify', None)\n",
    "make_config('lindisp', None)\n",
    "\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python run_nerf.py \\\n",
    "    --datadir {path} \\\n",
    "    --basedir {logs_path} \\\n",
    "    --dataset_type llff \\\n",
    "    --factor {downsample_factor} \\\n",
    "    --netdepth {netdepth} \\\n",
    "    --netwidth {netwidth} \\\n",
    "    --netdepth_fine {netdepth} \\\n",
    "    --netwidth_fine {netwidth} \\\n",
    "    --chunk {rendering_speed} \\\n",
    "    --netchunk {n_points_per_ray} \\\n",
    "    --lrate {learning_rate} \\\n",
    "    --i_video {max_iter // video_saving_cnt} \\\n",
    "    --expname {experiment_name} \\\n",
    "    --N_samples {n_samples} \\\n",
    "    --N_importance {n_importance} \\\n",
    "    --random_seed {random_seed} \\\n",
    "    --raw_noise_std 1.0 \\\n",
    "    --use_viewdirs \\\n",
    "    --no_ndc \\\n",
    "    --spherify \\\n",
    "    --lindisp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4\n",
    "### Get Mesh file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- install dependent-packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install imageio\n",
    "!pip install imageio-ffmpeg\n",
    "!pip install configargparse\n",
    "!apt install imagemagick\n",
    "!pip install PyMCubes\n",
    "!pip install trimesh\n",
    "!pip install pyrender\n",
    "!sudo apt -qq install imagemagick\n",
    "!pip install ConfigArgParse -qqq\n",
    "!pip install imageio-ffmpeg -qqq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- clone official code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/bmild/nerf.git\n",
    "\n",
    "%cd nerf\n",
    "!ls -al\n",
    "\n",
    "import run_nerf\n",
    "import run_nerf_helpers\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import pprint\n",
    "import imageio\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expname = experiment_name\n",
    "config = f'{logs_path}/config.txt'\n",
    "\n",
    "print('Args:')\n",
    "print(open(config, 'r').read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = run_nerf.config_parser()\n",
    "ft_str = '--ft_path {}'.format(os.path.join(logs_path, 'model_030000.npy'))\n",
    "args = parser.parse_args('--config {} '.format(config) + ft_str)\n",
    "print(args)\n",
    "\n",
    "# create nerf model\n",
    "_, render_kwargs_test, start, grad_vars, models = run_nerf.create_nerf(args)\n",
    "\n",
    "bds_dict = {\n",
    "    'near' : tf.cast(2., tf.float32),\n",
    "    'far' : tf.cast(6., tf.float32),\n",
    "}\n",
    "render_kwargs_test.update(bds_dict)\n",
    "\n",
    "print('Render kwargs:')\n",
    "pprint.pprint(render_kwargs_test)\n",
    "\n",
    "net_fn = render_kwargs_test['network_query_fn']\n",
    "print(net_fn)\n",
    "\n",
    "c2w = np.eye(4)[:3,:4].astype(np.float32) # identity pose matrix\n",
    "c2w[2,-1] = 4.\n",
    "H, W, focal = 800, 800, 1200.\n",
    "down = 8\n",
    "test = run_nerf.render(H//down, W//down, focal/down, c2w=c2w, **render_kwargs_test)\n",
    "img = np.clip(test[0],0,1)\n",
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Query network on dense 3d grid of points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 256\n",
    "t = np.linspace(-1.2, 1.2, N+1)\n",
    "\n",
    "query_pts = np.stack(np.meshgrid(t, t, t), -1).astype(np.float32)\n",
    "print(query_pts.shape)\n",
    "sh = query_pts.shape\n",
    "flat = query_pts.reshape([-1,3])\n",
    "\n",
    "\n",
    "def batchify(fn, chunk):\n",
    "    if chunk is None:\n",
    "        return fn\n",
    "    def ret(inputs):\n",
    "        return tf.concat([fn(inputs[i:i+chunk]) for i in range(0, inputs.shape[0], chunk)], 0)\n",
    "    return ret\n",
    "    \n",
    "    \n",
    "fn = lambda i0, i1 : net_fn(flat[i0:i1,None,:], viewdirs=np.zeros_like(flat[i0:i1]), network_fn=render_kwargs_test['network_fine'])\n",
    "chunk = 1024*64\n",
    "raw = np.concatenate([fn(i, i+chunk).numpy() for i in range(0, flat.shape[0], chunk)], 0)\n",
    "raw = np.reshape(raw, list(sh[:-1]) + [-1])\n",
    "sigma = np.maximum(raw[...,-1], 0.)\n",
    "\n",
    "print(raw.shape)\n",
    "plt.hist(np.maximum(0,sigma.ravel()), log=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Marching cubes with PyMCubes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mcubes\n",
    "\n",
    "threshold = 50.\n",
    "print('fraction occupied', np.mean(sigma > threshold))\n",
    "vertices, triangles = mcubes.marching_cubes(sigma, threshold)\n",
    "print('done', vertices.shape, triangles.shape)\n",
    "\n",
    "mcubes.export_obj(vertices, triangles, '{logs_path}/extract_mesh.obj')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Save out video with pyrender"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import trimesh\n",
    "\n",
    "mesh = trimesh.Trimesh(vertices / N - .5, triangles)\n",
    "mesh.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PYOPENGL_PLATFORM\"] = \"egl\"\n",
    "import pyrender\n",
    "from load_blender import pose_spherical\n",
    "\n",
    "scene = pyrender.Scene()\n",
    "scene.add(pyrender.Mesh.from_trimesh(mesh, smooth=False))\n",
    "\n",
    "# Set up the camera -- z-axis away from the scene, x-axis right, y-axis up\n",
    "camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)\n",
    "\n",
    "camera_pose = pose_spherical(-20., -40., 1.).numpy()\n",
    "nc = pyrender.Node(camera=camera, matrix=camera_pose)\n",
    "scene.add_node(nc)\n",
    "\n",
    "# Set up the light -- a point light in the same spot as the camera\n",
    "light = pyrender.PointLight(color=np.ones(3), intensity=4.0)\n",
    "nl = pyrender.Node(light=light, matrix=camera_pose)\n",
    "scene.add_node(nl)\n",
    "\n",
    "# Render the scene\n",
    "r = pyrender.OffscreenRenderer(640, 480)\n",
    "color, depth = r.render(scene)\n",
    "\n",
    "plt.imshow(color)\n",
    "plt.show()\n",
    "plt.imshow(depth)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from base64 import b64encode\n",
    "\n",
    "imgs = []\n",
    "for th in np.linspace(0, 360., 120+1)[:-1]:\n",
    "    camera_pose = pose_spherical(th, -40., 1.).numpy()\n",
    "    scene.set_pose(nc, pose=camera_pose)\n",
    "    imgs.append(r.render(scene)[0])\n",
    "f = '{logs_path}/mesh_turntable.mp4'\n",
    "imageio.mimwrite(f, imgs, fps=30)\n",
    "print('done')\n",
    "\n",
    "mp4 = open(f,'rb').read()\n",
    "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "HTML(\"\"\"\n",
    "<video width=400 controls autoplay loop>\n",
    "      <source src=\"%s\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\" % data_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.1 64-bit",
   "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.9.1"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}