diff --git a/colab/demo.ipynb b/colab/demo.ipynb index 1e8404622..0ad063dfd 100644 --- a/colab/demo.ipynb +++ b/colab/demo.ipynb @@ -3,8 +3,8 @@ { "cell_type": "markdown", "metadata": { - "id": "view-in-github", - "colab_type": "text" + "colab_type": "text", + "id": "view-in-github" }, "source": [ "\"Open" @@ -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

Install Nerfstudio and Dependencies (~8 min)

\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 12\u001b[0m \u001b[39m# Installing COLMAP via conda\u001b[39;00m\n\u001b[1;32m 13\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---> 14\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mcondacolab\u001b[39;00m\n\u001b[1;32m 15\u001b[0m condacolab\u001b[39m.\u001b[39minstall()\n\u001b[1;32m 17\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 \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

Install Nerfstudio and Dependencies (~8 min)

\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

Install COLMAP: Necessary for custom dataset (~15 min)

\n", + "#@markdown Skip to the next cell if running on demo data.
\n", + "#@markdown Run this cell if you are uploading your own data. Installation takes a while to complete.
\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

Downloading and Processing Data

\n", - "#@markdown

Pick the preset scene or upload your own images/video

\n", + "#@markdown

Downloading and Processing Data

\n", + "#@markdown Pick the preset scene or upload your own images/video\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('

Error: Data processing did not complete

'))\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": { diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index 33476f92e..65f3b30bd 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -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. diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index e170e0961..9227a64a7 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -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 diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 9f8951042..d1e69411e 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -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. diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 9cf86dc86..05d959836 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -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