tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 19:33:46 +08:00
Code Activity

WIP on gaussian-splatting: 8a3e2473 Update nerfstudio/viewer/server/render_state_machine.py

Browse Source
This commit is contained in:
Gina Wu committed 2023-10-20 14:51:39 -07:00
2 parents 8a3e24738b 09d5a2a405
commit 2dc7f955ce
5 files changed
+153 -58
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

colab/demo.ipynb
+111 -43
View File
@@ -3,8 +3,8 @@
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
"colab_type": "text",
"id": "view-in-github"
},
"source": [
"<a href=\"https://colab.research.google.com/github/nerfstudio-project/nerfstudio/blob/alex%2Fcolab310/colab/demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
@@ -32,19 +32,24 @@
"\n",
"This colab shows how to train and view NeRFs from Nerfstudio both on pre-made datasets or from your own videos/images.\n",
"\n",
"\\\\\n",
"\n",
"\n",
"Credit to [NeX](https://nex-mpi.github.io/) for Google Colab format."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Frequently Asked Questions"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Yyx5h6kz5ga7"
},
"source": [
"## Frequently Asked Questions\n",
"\n",
"* **Downloading custom data is stalling (no output):**\n",
" * This is a bug in Colab. The data is processing, but may take a while to complete. You will know processing completed if `data/nerfstudio/custom_data/transforms.json` exists. Terminating the cell early will result in not being able to train.\n",
"* **Processing custom data is taking a long time:**\n",
@@ -62,15 +67,76 @@
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9oyLHl8QfYwP",
"cellView": "form"
},
"outputs": [],
"cell_type": "markdown",
"metadata": {},
"source": [
"#@markdown <h1>Install Nerfstudio and Dependencies (~8 min)</h1>\n",
"## Nerfstudio Demo"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"cellView": "form",
"id": "9oyLHl8QfYwP"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[Errno 2] No such file or directory: '/content/'\n",
"/home/gwu/nerfstudio/colab\n",
"Requirement already satisfied: pip in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (23.2.1)\n",
"Looking in indexes: https://pypi.org/simple, https://download.pytorch.org/whl/cu118\n",
"Requirement already satisfied: torch==2.0.1+cu118 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (2.0.1+cu118)\n",
"Requirement already satisfied: torchvision==0.15.2+cu118 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (0.15.2+cu118)\n",
"Requirement already satisfied: filelock in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torch==2.0.1+cu118) (3.12.4)\n",
"Requirement already satisfied: typing-extensions in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torch==2.0.1+cu118) (4.8.0)\n",
"Requirement already satisfied: sympy in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torch==2.0.1+cu118) (1.12)\n",
"Requirement already satisfied: networkx in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torch==2.0.1+cu118) (3.1)\n",
"Requirement already satisfied: jinja2 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torch==2.0.1+cu118) (3.1.2)\n",
"Requirement already satisfied: triton==2.0.0 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torch==2.0.1+cu118) (2.0.0)\n",
"Requirement already satisfied: numpy in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torchvision==0.15.2+cu118) (1.26.0)\n",
"Requirement already satisfied: requests in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torchvision==0.15.2+cu118) (2.31.0)\n",
"Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from torchvision==0.15.2+cu118) (10.0.1)\n",
"Requirement already satisfied: cmake in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from triton==2.0.0->torch==2.0.1+cu118) (3.27.5)\n",
"Requirement already satisfied: lit in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from triton==2.0.0->torch==2.0.1+cu118) (17.0.1)\n",
"Requirement already satisfied: MarkupSafe>=2.0 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from jinja2->torch==2.0.1+cu118) (2.1.3)\n",
"Requirement already satisfied: charset-normalizer<4,>=2 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu118) (3.2.0)\n",
"Requirement already satisfied: idna<4,>=2.5 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu118) (3.4)\n",
"Requirement already satisfied: urllib3<3,>=1.21.1 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu118) (1.26.16)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu118) (2023.7.22)\n",
"Requirement already satisfied: mpmath>=0.19 in /home/gwu/miniconda3/envs/nerfstudio/lib/python3.10/site-packages (from sympy->torch==2.0.1+cu118) (1.3.0)\n",
"[Errno 2] No such file or directory: '/content/'\n",
"/home/gwu/nerfstudio/colab\n",
"Downloading...\n",
"From: https://drive.google.com/u/1/uc?id=1-7x7qQfB7bIw2zV4Lr6-yhvMpjXC84Q5&confirm=t\n",
"To: /home/gwu/nerfstudio/colab/tinycudann-1.7-cp310-cp310-linux_x86_64.whl\n",
"100%|██████████████████████████████████████| 31.2M/31.2M [00:00<00:00, 50.7MB/s]\n",
"Processing ./tinycudann-1.7-cp310-cp310-linux_x86_64.whl\n",
"tinycudann is already installed with the same version as the provided wheel. Use --force-reinstall to force an installation of the wheel.\n"
]
},
{
"ename": "RuntimeError",
"evalue": "This module must ONLY run as part of a Colab notebook!",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)",
"File \u001b[0;32m~/miniconda3/envs/nerfstudio/lib/python3.10/site-packages/condacolab.py:27\u001b[0m\n\u001b[1;32m 26\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m---> 27\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mgoogle\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcolab\u001b[39;00m\n\u001b[1;32m 28\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mImportError\u001b[39;00m:\n",
"\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'google.colab'",
"\nDuring handling of the above exception, another exception occurred:\n",
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
"\u001b[1;32m/home/gwu/nerfstudio/colab/demo.ipynb Cell 6\u001b[0m line \u001b[0;36m1\n\u001b[1;32m <a href='vscode-notebook-cell://ssh-remote%2Bdecoder.eecs.berkeley.edu/home/gwu/nerfstudio/colab/demo.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=11'>12</a>\u001b[0m \u001b[39m# Installing COLMAP via conda\u001b[39;00m\n\u001b[1;32m <a href='vscode-notebook-cell://ssh-remote%2Bdecoder.eecs.berkeley.edu/home/gwu/nerfstudio/colab/demo.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=12'>13</a>\u001b[0m get_ipython()\u001b[39m.\u001b[39msystem(\u001b[39m'\u001b[39m\u001b[39mpip install -q condacolab\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2Bdecoder.eecs.berkeley.edu/home/gwu/nerfstudio/colab/demo.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=13'>14</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mcondacolab\u001b[39;00m\n\u001b[1;32m <a href='vscode-notebook-cell://ssh-remote%2Bdecoder.eecs.berkeley.edu/home/gwu/nerfstudio/colab/demo.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=14'>15</a>\u001b[0m condacolab\u001b[39m.\u001b[39minstall()\n\u001b[1;32m <a href='vscode-notebook-cell://ssh-remote%2Bdecoder.eecs.berkeley.edu/home/gwu/nerfstudio/colab/demo.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D?line=16'>17</a>\u001b[0m get_ipython()\u001b[39m.\u001b[39msystem(\u001b[39m'\u001b[39m\u001b[39mconda install -c conda-forge colmap\u001b[39m\u001b[39m'\u001b[39m)\n",
"File \u001b[0;32m~/miniconda3/envs/nerfstudio/lib/python3.10/site-packages/condacolab.py:29\u001b[0m\n\u001b[1;32m 27\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mgoogle\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcolab\u001b[39;00m\n\u001b[1;32m 28\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mImportError\u001b[39;00m:\n\u001b[0;32m---> 29\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39mThis module must ONLY run as part of a Colab notebook!\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m 32\u001b[0m __version__ \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39m0.1.7\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 33\u001b[0m __author__ \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mJaime Rodríguez-Guerra <jaimergp@users.noreply.github.com>\u001b[39m\u001b[39m\"\u001b[39m\n",
"\u001b[0;31mRuntimeError\u001b[0m: This module must ONLY run as part of a Colab notebook!"
]
}
],
"source": [
"#@markdown <h3>Install Nerfstudio and Dependencies (~8 min)</h3>\n",
"\n",
"%cd /content/\n",
"!pip install --upgrade pip\n",
@@ -81,36 +147,30 @@
"!gdown \"https://drive.google.com/u/1/uc?id=1-7x7qQfB7bIw2zV4Lr6-yhvMpjXC84Q5&confirm=t\"\n",
"!pip install tinycudann-1.7-cp310-cp310-linux_x86_64.whl\n",
"\n",
"# Installing COLMAP\n",
"%cd /content/\n",
"!gdown \"https://drive.google.com/u/0/uc?id=15WngFRNar_b8CaPR5R-hvQ3eAnlyk_SL&confirm=t\"\n",
"!sudo apt-get install \\\n",
" build-essential \\\n",
" libboost-program-options-dev \\\n",
" libboost-filesystem-dev \\\n",
" libboost-graph-dev \\\n",
" libboost-system-dev \\\n",
" libboost-test-dev \\\n",
" libeigen3-dev \\\n",
" libflann-dev \\\n",
" libfreeimage-dev \\\n",
" libmetis-dev \\\n",
" libgoogle-glog-dev \\\n",
" libgflags-dev \\\n",
" libsqlite3-dev \\\n",
" libglew-dev \\\n",
" qtbase5-dev \\\n",
" libqt5opengl5-dev \\\n",
" libcgal-dev \\\n",
" libceres-dev\n",
"!unzip local.zip -d /usr/\n",
"!chmod +x /usr/local/bin/colmap\n",
"# Installing COLMAP via conda\n",
"%cd /content\n",
"!conda install -c conda-forge colmap\n",
"\n",
"# Install nerfstudio\n",
"%cd /content/\n",
"!pip install git+https://github.com/nerfstudio-project/nerfstudio.git"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#@markdown <h3>Install COLMAP: Necessary for custom dataset (~15 min)</h3>\n",
"#@markdown <medium>Skip to the next cell if running on demo data.<br/medium>\n",
"#@markdown <medium>Run this cell if you are uploading your own data. Installation takes a while to complete.<br/medium>\n",
"\n",
"# Installing COLMAP via conda\n",
"%cd /content\n",
"!conda install -c conda-forge colmap"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -120,8 +180,8 @@
},
"outputs": [],
"source": [
"#@markdown <h1> Downloading and Processing Data</h1>\n",
"#@markdown <h3>Pick the preset scene or upload your own images/video</h3>\n",
"#@markdown <h3> Downloading and Processing Data</h3>\n",
"#@markdown <medium>Pick the preset scene or upload your own images/video</medium>\n",
"import os\n",
"import glob\n",
"from google.colab import files\n",
@@ -278,7 +338,7 @@
"\n",
"%cd /content\n",
"if os.path.exists(f\"data/nerfstudio/{scene}/transforms.json\"):\n",
" !ns-train nerfacto --viewer.websocket-port 7007 nerfstudio-data --data data/nerfstudio/$scene --downscale-factor 4\n",
" !ns-train nerfacto --viewer.websocket-port 7007 --vis viewer_beta --viewer.make-share-url True nerfstudio-data --data data/nerfstudio/$scene --downscale-factor 4\n",
"else:\n",
" from IPython.core.display import display, HTML\n",
" display(HTML('<h3 style=\"color:red\">Error: Data processing did not complete</h3>'))\n",
@@ -350,8 +410,8 @@
"metadata": {
"accelerator": "GPU",
"colab": {
"provenance": [],
"include_colab_link": true
"include_colab_link": true,
"provenance": []
},
"gpuClass": "standard",
"kernelspec": {
@@ -360,8 +420,16 @@
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"version": "3.8.13"
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
},
"vscode": {
"interpreter": {
nerfstudio/cameras/rays.py
+2
View File
@@ -210,6 +210,8 @@ class RayBundle(TensorDataclass):
times: Optional[Float[Tensor, "*batch 1"]] = None
"""Times at which rays are sampled"""
# TODO Camera
def set_camera_indices(self, camera_index: int) -> None:
"""Sets all the camera indices to a specific camera index.
nerfstudio/data/datamanagers/full_images_datamanager.py
+10 -6
View File
@@ -238,10 +238,11 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
cached_eval.append(data)
self.eval_dataset.cameras[i].fx = K[0, 0]
self.eval_dataset.cameras[i].fy = K[1, 1]
self.eval_dataset.cameras[i].cx = K[0, 2]
self.eval_dataset.cameras[i].cy = K[1, 2]
self.eval_dataset.cameras.fx[i] = float(K[0, 0])
self.eval_dataset.cameras.fy[i] = float(K[1, 1])
self.eval_dataset.cameras.cx[i] = float(K[0, 2])
self.eval_dataset.cameras.cy[i] = float(K[1, 2])
# print(self.eval_dataset.cameras[i].reshape((1,)))
if cache_images_option == "gpu":
for cache in cached_train:
@@ -342,7 +343,10 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
data = self.eval_dataset.get_data(image_idx)
data["image"] = data["image"].to(self.device)
assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension"
camera = self.eval_dataset.cameras[image_idx].to(self.device)
camera = self.eval_dataset.cameras[image_idx : image_idx + 1].to(self.device)
# TODO: camera shape torch.Size([]) for eval dataset
# should be torch.Size([1])
print(camera.shape, data.keys())
return camera, data
def next_eval_image(self, step: int) -> Tuple[int, Cameras, Dict]:
@@ -358,5 +362,5 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
data = self.eval_dataset.get_data(image_idx)
data["image"] = data["image"].to(self.device)
assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension"
camera = self.eval_dataset.cameras[image_idx].to(self.device)
camera = self.eval_dataset.cameras[image_idx: image_idx + 1].to(self.device)
return image_idx, camera, data
nerfstudio/models/gaussian_splatting.py
+29 -8
View File
@@ -25,7 +25,7 @@ from nerfstudio.data.scene_box import OrientedBox
from copy import deepcopy
from nerfstudio.cameras.rays import RayBundle
from torch.nn import Parameter
from torchmetrics.image import PeakSignalNoiseRatio
from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from nerfstudio.cameras.cameras import Cameras
from gsplat._torch_impl import quat_to_rotmat
@@ -162,13 +162,13 @@ class GaussianSplattingModel(Model):
self.shs_rest = torch.nn.Parameter(torch.zeros((self.num_points, 3, dim_sh - 1)))
self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1)))
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
from torchmetrics.image import StructuralSimilarityIndexMeasure
self.ssim = StructuralSimilarityIndexMeasure()
self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True)
self.step = 0
self.crop_box: Optional[OrientedBox] = None
self.back_color = torch.zeros(3)
@@ -457,7 +457,10 @@ class GaussianSplattingModel(Model):
}
def _get_downscale_factor(self):
return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0)
if self.training:
return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0)
else:
return 1
def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
"""Takes in a Ray Bundle and returns a dictionary of outputs.
@@ -472,12 +475,12 @@ class GaussianSplattingModel(Model):
if not isinstance(camera, Cameras):
print("Called get_outputs with not a camera")
return {}
# TODO: camera shape in eval dataset breaking assert
assert camera.shape[0] == 1, "Only one camera at a time"
# dont mutate the input
camera = deepcopy(camera)
if self.training:
d = self._get_downscale_factor()
camera.rescale_output_resolution(1 / d)
d = self._get_downscale_factor()
camera.rescale_output_resolution(1 / d)
if self.crop_box is not None and not self.training:
crop_ids = self.crop_box.within(self.means).squeeze()
if crop_ids.sum() == 0:
@@ -580,8 +583,26 @@ class GaussianSplattingModel(Model):
outputs: the output to compute loss dict to
batch: ground truth batch corresponding to outputs
"""
# batch["image"]: torch.Size([959, 539, 3])
# outputs["rgb"]: torch.Size([239, 134, 3])
d = self._get_downscale_factor()
if d > 1:
# use torchvision to resize
import torchvision.transforms.functional as TF
return {}
newsize = (batch["image"].shape[0] // d, batch["image"].shape[1] // d)
gt_img = TF.resize(batch["image"].permute(2, 0, 1), newsize).permute(1, 2, 0)
else:
gt_img = batch["image"]
# TODO: standard PSNR, L-PIPS, and SSIM metrics
metrics_dict = {}
gt_rgb = gt_img.to(self.device) # RGB or RGBA image
# gt_rgb = self.renderer_rgb.blend_background(gt_rgb) # Blend if RGBA
predicted_rgb = outputs["rgb"]
metrics_dict["psnr"] = self.psnr(predicted_rgb, gt_rgb)
self.camera_optimizer.get_metrics_dict(metrics_dict)
return metrics_dict
def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
"""Computes and returns the losses dict.
nerfstudio/pipelines/base_pipeline.py
+1 -1
View File
@@ -338,7 +338,7 @@ class VanillaPipeline(Pipeline):
"""
self.eval()
image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step)
outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
outputs = self.model.get_outputs_for_camera_ray_bundle(None, camera=camera_ray_bundle)
metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch)
assert "image_idx" not in metrics_dict
metrics_dict["image_idx"] = image_idx