extract_mesh.ipynb · lulala/nerf_pl

加入 Gitee

与超过 1200万开发者一起发现、参与优秀开源项目，私有仓库也完全免费：）

克隆/下载

extract_mesh.ipynb 8.16 KB

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "import mcubes\n",
    "import trimesh\n",
    "\n",
    "from models.rendering import *\n",
    "from models.nerf import *\n",
    "\n",
    "from datasets import dataset_dict\n",
    "\n",
    "from utils import load_ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load model and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change here #\n",
    "img_wh = (800, 800) # full resolution of the input images\n",
    "dataset_name = 'blender' # blender or llff (own data)\n",
    "scene_name = 'lego' # whatever you want\n",
    "root_dir = '/home/ubuntu/data/nerf_example_data/nerf_synthetic/lego/' # the folder containing data\n",
    "ckpt_path = 'ckpts/exp2/epoch=05.ckpt' # the model path\n",
    "###############\n",
    "\n",
    "kwargs = {'root_dir': root_dir,\n",
    "          'img_wh': img_wh}\n",
    "if dataset_name == 'llff':\n",
    "    kwargs['spheric_poses'] = True\n",
    "    kwargs['split'] = 'test'\n",
    "else:\n",
    "    kwargs['split'] = 'train'\n",
    "    \n",
    "chunk = 1024*32\n",
    "dataset = dataset_dict[dataset_name](**kwargs)\n",
    "\n",
    "embedding_xyz = Embedding(3, 10)\n",
    "embedding_dir = Embedding(3, 4)\n",
    "\n",
    "nerf_fine = NeRF()\n",
    "load_ckpt(nerf_fine, ckpt_path, model_name='nerf_fine')\n",
    "nerf_fine.cuda().eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Search for tight bounds of the object (trial and error!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Tune these parameters until the whole object lies tightly in range with little noise ###\n",
    "N = 128 # controls the resolution, set this number small here because we're only finding\n",
    "        # good ranges here, not yet for mesh reconstruction; we can set this number high\n",
    "        # when it comes to final reconstruction.\n",
    "xmin, xmax = -1.2, 1.2 # left/right range\n",
    "ymin, ymax = -1.2, 1.2 # forward/backward range\n",
    "zmin, zmax = -1.2, 1.2 # up/down range\n",
    "## Attention! the ranges MUST have the same length!\n",
    "sigma_threshold = 50. # controls the noise (lower=maybe more noise; higher=some mesh might be missing)\n",
    "############################################################################################\n",
    "\n",
    "x = np.linspace(xmin, xmax, N)\n",
    "y = np.linspace(ymin, ymax, N)\n",
    "z = np.linspace(zmin, zmax, N)\n",
    "\n",
    "xyz_ = torch.FloatTensor(np.stack(np.meshgrid(x, y, z), -1).reshape(-1, 3)).cuda()\n",
    "dir_ = torch.zeros_like(xyz_).cuda()\n",
    "\n",
    "with torch.no_grad():\n",
    "    B = xyz_.shape[0]\n",
    "    out_chunks = []\n",
    "    for i in range(0, B, chunk):\n",
    "        xyz_embedded = embedding_xyz(xyz_[i:i+chunk]) # (N, embed_xyz_channels)\n",
    "        dir_embedded = embedding_dir(dir_[i:i+chunk]) # (N, embed_dir_channels)\n",
    "        xyzdir_embedded = torch.cat([xyz_embedded, dir_embedded], 1)\n",
    "        out_chunks += [nerf_fine(xyzdir_embedded)]\n",
    "    rgbsigma = torch.cat(out_chunks, 0)\n",
    "    \n",
    "sigma = rgbsigma[:, -1].cpu().numpy()\n",
    "sigma = np.maximum(sigma, 0)\n",
    "sigma = sigma.reshape(N, N, N)\n",
    "\n",
    "# The below lines are for visualization, COMMENT OUT once you find the best range and increase N!\n",
    "vertices, triangles = mcubes.marching_cubes(sigma, sigma_threshold)\n",
    "mesh = trimesh.Trimesh(vertices/N, triangles)\n",
    "mesh.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # You can already export \"colorless\" mesh if you don't need color\n",
    "# mcubes.export_mesh(vertices, triangles, f\"{scene_name}.dae\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate .vol file for volume rendering in Unity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert N==512, \\\n",
    "    'Please set N to 512 in the two above cell! Remember to comment out the visualization code (last 3 lines)!'\n",
    "\n",
    "a = 1-np.exp(-(xmax-xmin)/N*sigma)\n",
    "a = a.flatten()\n",
    "rgb = (rgbsigma[:, :3].numpy()*255).astype(np.uint32)\n",
    "i = np.where(a>0)[0] # valid indices (alpha>0)\n",
    "\n",
    "rgb = rgb[i]\n",
    "a = a[i]\n",
    "s = rgb.dot(np.array([1<<24, 1<<16, 1<<8])) + (a*255).astype(np.uint32)\n",
    "res = np.stack([i, s], -1).astype(np.uint32).flatten()\n",
    "with open(f'{scene_name}.vol', 'wb') as f:\n",
    "    f.write(res.tobytes())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract colored mesh\n",
    "\n",
    "Once you find the best range, now **RESTART** the notebook, and copy the configs to the following cell\n",
    "and execute it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy the variables you have above here! ####\n",
    "img_wh = (800, 800) # full resolution of the input images\n",
    "dataset_name = 'blender' # blender or llff (own data)\n",
    "scene_name = 'lego' # whatever you want\n",
    "root_dir = '/home/ubuntu/data/nerf_example_data/nerf_synthetic/lego/' # the folder containing data\n",
    "ckpt_path = 'ckpts/exp2/epoch=05.ckpt' # the model path\n",
    "\n",
    "N = 128 # controls the resolution, set this number small here because we're only finding\n",
    "        # good ranges here, not yet for mesh reconstruction; we can set this number high\n",
    "        # when it comes to final reconstruction.\n",
    "xmin, xmax = -1.2, 1.2 # left/right range\n",
    "ymin, ymax = -1.2, 1.2 # forward/backward range\n",
    "zmin, zmax = -1.2, 1.2 # up/down range\n",
    "sigma_threshold = 50. # controls the noise (lower=maybe more noise; higher=some mesh might be missing)\n",
    "###############################################\n",
    "\n",
    "import os\n",
    "os.environ['ROOT_DIR'] = root_dir\n",
    "os.environ['DATASET_NAME'] = dataset_name\n",
    "os.environ['SCENE_NAME'] = scene_name\n",
    "os.environ['IMG_SIZE'] = f\"{img_wh[0]} {img_wh[1]}\"\n",
    "os.environ['CKPT_PATH'] = ckpt_path\n",
    "os.environ['N_GRID'] = \"256\" # final resolution. You can set this number high to preserve more details\n",
    "os.environ['X_RANGE'] = f\"{xmin} {xmax}\"\n",
    "os.environ['Y_RANGE'] = f\"{ymin} {ymax}\"\n",
    "os.environ['Z_RANGE'] = f\"{zmin} {zmax}\"\n",
    "os.environ['SIGMA_THRESHOLD'] = str(sigma_threshold)\n",
    "os.environ['OCC_THRESHOLD'] = \"0.2\" # probably doesn't require tuning. If you find the color is not close\n",
    "                                    # to real, try to set this number smaller (the effect of this number\n",
    "                                    # is explained in my youtube video)\n",
    "\n",
    "!python extract_color_mesh.py \\\n",
    "    --root_dir $ROOT_DIR \\\n",
    "    --dataset_name $DATASET_NAME \\\n",
    "    --scene_name $SCENE_NAME \\\n",
    "    --img_wh $IMG_SIZE \\\n",
    "    --ckpt_path $CKPT_PATH \\\n",
    "    --N_grid $N_GRID \\\n",
    "    --x_range $X_RANGE \\\n",
    "    --y_range $Y_RANGE \\\n",
    "    --z_range $Z_RANGE \\\n",
    "    --sigma_threshold $SIGMA_THRESHOLD \\\n",
    "    --occ_threshold $OCC_THRESHOLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nerf_pl",
   "language": "python",
   "name": "nerf_pl"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}

一键复制编辑原始数据按行查看历史

提交于 2020-05-26 21:44 . add more comments

马建仓 AI 助手

尝试更多

代码解读

代码找茬

代码优化

https://gitee.com/lula121/nerf_pl.git

git@gitee.com:lula121/nerf_pl.git

lula121

nerf_pl

master

lulala/nerf_pl

简介

发行版

贡献者

近期动态

lulala/nerf_pl .gitee-modal { width: 500px !important; }

简介

发行版

贡献者

近期动态

搜索帮助

lulala/nerf_pl