Skip to content

Instantly share code, notes, and snippets.

@kwea123

kwea123/extract_mesh_colab.ipynb Secret

Created May 8, 2020 07:30
Show Gist options
  • Star 2 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save kwea123/77ed1640f9bc9550136dc13a6a419e88 to your computer and use it in GitHub Desktop.
Save kwea123/77ed1640f9bc9550136dc13a6a419e88 to your computer and use it in GitHub Desktop.
Download ZIP
extract_mesh_colab.ipynb
Display the source blob
Display the rendered blob
Raw
extract_mesh_colab.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "extract_mesh_colab.ipynb",
"provenance": [],
"collapsed_sections": [],
"authorship_tag": "ABX9TyOv+dJymJtTHyx6MzCzBlLS",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/kwea123/77ed1640f9bc9550136dc13a6a419e88/extract_mesh_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ocOP60tuACFX",
"colab_type": "text"
},
"source": [
"# Installation"
]
},
{
"cell_type": "code",
"metadata": {
"id": "djAoxA4c8ku3",
"colab_type": "code",
"colab": {}
},
"source": [
"!git clone --recursive https://github.com/kwea123/nerf_pl\n",
"\n",
"%cd /content/nerf_pl\n",
"!pip install -r requirements.txt\n",
"\n",
"%cd /content/nerf_pl/torchsearchsorted\n",
"!pip install ."
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "UgZy82gl9ieI",
"colab_type": "text"
},
"source": [
"# Mount your drive (to access data)"
]
},
{
"cell_type": "code",
"metadata": {
"id": "o0KzmUyd9Z7s",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 124
},
"outputId": "fdf978ea-a264-4a10-85a6-f0b1bf5aa994"
},
"source": [
"from google.colab import drive\n",
"drive.mount('/content/drive/', force_remount=True)"
],
"execution_count": 2,
"outputs": [
{
"output_type": "stream",
"text": [
"Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
"\n",
"Enter your authorization code:\n",
"··········\n",
"Mounted at /content/drive/\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "iXb2f51D9S3a",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 34
},
"outputId": "32bd5380-5899-4a19-bb0e-a406a4c9874e"
},
"source": [
"%cd /content/nerf_pl\n",
"\n",
"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"
],
"execution_count": 4,
"outputs": [
{
"output_type": "stream",
"text": [
"/content/nerf_pl\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MX-0OSxi-9Uo",
"colab_type": "text"
},
"source": [
"# Load model"
]
},
{
"cell_type": "code",
"metadata": {
"id": "eGDyBnIk9Vy3",
"colab_type": "code",
"colab": {}
},
"source": [
"# Change here #\n",
"img_wh = (4032, 3024) # full resolution of the input images\n",
"dataset_name = 'llff' # blender or llff (own data)\n",
"scene_name = 'silica' # whatever you want\n",
"root_dir = '/content/drive/My Drive/colab/nerf/my/silica' # the folder containing data (images/ and poses_bounds.npy)\n",
"ckpt_path = '/content/silica.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();"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "hy5p5fBp-22w",
"colab_type": "text"
},
"source": [
"# Search for tight bounds (trial and error!)"
]
},
{
"cell_type": "code",
"metadata": {
"id": "MrFyZ2u5-D9i",
"colab_type": "code",
"colab": {}
},
"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, 1 # left/right range\n",
"ymin, ymax = -1, 1 # forward/backward range\n",
"zmin, zmax = -2.64, -0.64 # up/down range\n",
"## Attention! the ranges MUST have the same length!\n",
"sigma_threshold = 20. # 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",
"vertices, triangles = mcubes.marching_cubes(sigma, sigma_threshold)\n",
"\n",
"mesh = trimesh.Trimesh(vertices/N, triangles)\n",
"mesh.show()"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "uAKgExgM-HcP",
"colab_type": "code",
"colab": {}
},
"source": [
"# You can already export \"colorless\" mesh if you don't need color\n",
"mcubes.export_mesh(vertices, triangles, f\"{scene_name}.dae\")"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "SxUKX2SU_AHn",
"colab_type": "text"
},
"source": [
"# Extract colored mesh\n",
"Once you find the best range, now **RESTART** the notebook(Runtime->Restart runtime), and copy the configs to the following cell and execute it."
]
},
{
"cell_type": "code",
"metadata": {
"id": "1esQrk6a-UwR",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 244
},
"outputId": "cefe2160-7918-430e-9e73-41ccc5f4ae16"
},
"source": [
"%cd /content/nerf_pl\n",
"\n",
"# Copy the variables you have above here! ####\n",
"\n",
"img_wh = (4032, 3024) # full resolution of the input images\n",
"dataset_name = 'llff' # blender or llff (own data)\n",
"scene_name = 'silica' # whatever you want\n",
"root_dir = '/content/drive/My Drive/colab/nerf/my/silica' # the folder containing data\n",
"ckpt_path = '/content/silica.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, 1 # left/right range\n",
"ymin, ymax = -1, 1 # forward/backward range\n",
"zmin, zmax = -2.64, -0.64 # up/down range\n",
"## Attention! the ranges MUST have the same length!\n",
"sigma_threshold = 20. # 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'] = \"512\" # 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"
],
"execution_count": 5,
"outputs": [
{
"output_type": "stream",
"text": [
"/content/nerf_pl\n",
"tcmalloc: large alloc 1073741824 bytes == 0x77da4000 @ 0x7fb3e6ccb1e7 0x7fb39b4d65e1 0x7fb39b53b8e0 0x7fb39b530c95 0x7fb39b5c5957 0x50a635 0x50bfb4 0x507d64 0x509a90 0x50a48d 0x50bfb4 0x507d64 0x588d41 0x59fc4e 0x7fb39b5249ad 0x50a2bf 0x50bfb4 0x507d64 0x509a90 0x50a48d 0x50bfb4 0x507d64 0x50ae13 0x634c82 0x634d37 0x6384ef 0x639091 0x4b0d00 0x7fb3e68c8b97 0x5b250a\n",
"tcmalloc: large alloc 3221225472 bytes == 0x138634000 @ 0x7fb3e6ccb1e7 0x7fb39b4d65e1 0x7fb39b539138 0x7fb39b539253 0x7fb39b5d7368 0x7fb39b5d7bc4 0x7fb39b5d7d12 0x566d63 0x59fc4e 0x7fb39b5249ad 0x50a2bf 0x50bfb4 0x507d64 0x509a90 0x50a48d 0x50cd96 0x507d64 0x588d41 0x59fc4e 0x7fb39b5249ad 0x50a2bf 0x50bfb4 0x507d64 0x509a90 0x50a48d 0x50bfb4 0x507d64 0x50ae13 0x634c82 0x634d37 0x6384ef\n",
"tcmalloc: large alloc 1610612736 bytes == 0x77da4000 @ 0x7fb3e6cadb6b 0x7fb3e6ccd379 0x7fb39bc2b04e 0x7fb39bc2cf4a 0x7fb3d4b1b67b 0x7fb3d476a6be 0x7fb3d49d37b5 0x7fb3d49c57c1 0x7fb3d49c4d0e 0x7fb3d49c57c1 0x7fb3d641a93a 0x7fb3d49c57c1 0x7fb3d4765457 0x7fb3d4766b59 0x7fb3d4a84773 0x7fb3d6502684 0x7fb3d4ad881a 0x7fb3e2fd8cf4 0x7fb3e2fd3dd4 0x7fb3e2fd49a9 0x7fb3e2fd6c7c 0x7fb3e2fb85ec 0x551365 0x5a9cbc 0x50a5c3 0x50bfb4 0x507d64 0x50ae13 0x634c82 0x634d37 0x6384ef\n",
"Predicting occupancy ...\n",
"100% 4096/4096 [00:36<00:00, 113.10it/s]\n",
"Extracting mesh ...\n",
"Removing noise ...\n",
"Mesh has 0.58 M vertices and 1.16 M faces.\n",
"Fusing colors ...\n",
"100% 65/65 [10:45<00:00, 9.94s/it]\n",
"Done!\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "5TQ1DnK3_zs4",
"colab_type": "code",
"colab": {}
},
"source": [
""
],
"execution_count": 0,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment