From 6858a0d2d945da333a59c0dbb8274b257ebfca31 Mon Sep 17 00:00:00 2001 From: Boris Feld Date: Thu, 8 Feb 2024 15:42:55 +0100 Subject: [PATCH] Add NerfStudio examples --- ..._NerfStudio_sitcoms3d_evaluate_model.ipynb | 148 ++++++++ .../Comet-NerfStudio_sitcoms3d.ipynb | 328 ++++++++++++++++++ 2 files changed, 476 insertions(+) create mode 100644 integrations/model-evaluation/nerfstudio/notebooks/Comet_NerfStudio_sitcoms3d_evaluate_model.ipynb create mode 100644 integrations/model-training/nerfstudio/notebooks/Comet-NerfStudio_sitcoms3d.ipynb diff --git a/integrations/model-evaluation/nerfstudio/notebooks/Comet_NerfStudio_sitcoms3d_evaluate_model.ipynb b/integrations/model-evaluation/nerfstudio/notebooks/Comet_NerfStudio_sitcoms3d_evaluate_model.ipynb new file mode 100644 index 00000000..79b0749c --- /dev/null +++ b/integrations/model-evaluation/nerfstudio/notebooks/Comet_NerfStudio_sitcoms3d_evaluate_model.ipynb @@ -0,0 +1,148 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "SiiXJ7K_fePG" + }, + "source": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "This notebook shows to use Comet with Nerfstudio a library that provides a simple API for a simplified end-to-end process of creating, training, and testing NeRFs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9oyLHl8QfYwP" + }, + "outputs": [], + "source": [ + "%cd /content/\n", + "%pip install --upgrade pip\n", + "%pip install -U torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "dqnlofRu10Ya" + }, + "outputs": [], + "source": [ + "# Installing TinyCuda\n", + "%cd /content/\n", + "!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" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "OMHbnUxO11a4" + }, + "outputs": [], + "source": [ + "# Install nerfstudio\n", + "%cd /content/\n", + "!pip install git+https://github.com/nerfstudio-project/nerfstudio" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "zuN8Ya9yuULS" + }, + "outputs": [], + "source": [ + "# Download the dataset which is needed for the viewer\n", + "%cd /content/\n", + "!ns-download-data sitcoms3d" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EJBXMB-fF-rB" + }, + "outputs": [], + "source": [ + "import comet_ml\n", + "\n", + "comet_ml.init()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "XfVVcrnr-unc" + }, + "outputs": [], + "source": [ + "import json\n", + "import comet_ml\n", + "\n", + "api = comet_ml.API()\n", + "\n", + "# Retrieve a model from the training experiment\n", + "EXPERIMENT_KEY = \"e2532a362cbc42d9b8d41c2510619707\"\n", + "MODEL_NAME = \"nerfstudio-model\"\n", + "\n", + "api_experiment = api.get_experiment_by_key(EXPERIMENT_KEY)\n", + "\n", + "# Retrieve the expected nerfstudio directory structure\n", + "model_assets = api_experiment.get_model_asset_list(MODEL_NAME)\n", + "metadata = model_assets[-1][\"metadata\"]\n", + "output_dir = json.loads(metadata)[\"nerfstudio_project_dir\"]\n", + "\n", + "api_experiment.download_model(MODEL_NAME, output_path=output_dir)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "background_save": true + }, + "id": "eJdbgqA4_hn7" + }, + "outputs": [], + "source": [ + "# Start the viewer\n", + "!ns-viewer --load-config \"$output_dir/config.yml\" --vis \"viewer\" --viewer.make-share-url True" + ] + } + ], + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.8.13" + }, + "vscode": { + "interpreter": { + "hash": "c59f626636933ef1dc834fb3684b382f705301c5306cf8436d2da634c2289783" + } + }, + "accelerator": "GPU" + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file diff --git a/integrations/model-training/nerfstudio/notebooks/Comet-NerfStudio_sitcoms3d.ipynb b/integrations/model-training/nerfstudio/notebooks/Comet-NerfStudio_sitcoms3d.ipynb new file mode 100644 index 00000000..5197e942 --- /dev/null +++ b/integrations/model-training/nerfstudio/notebooks/Comet-NerfStudio_sitcoms3d.ipynb @@ -0,0 +1,328 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "SiiXJ7K_fePG" + }, + "source": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "This notebook shows to use [Comet](https://www.comet.com/site/?utm_source=nerf&utm_medium=referral&utm_content=colab) with Nerfstudio: a popular library that provides a simple API for a simplified end-to-end process of creating, training, and testing NeRFs." + ] + }, + { + "cell_type": "markdown", + "source": [ + "# Install dependencies" + ], + "metadata": { + "id": "9gwWVZAdG2vk" + } + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9oyLHl8QfYwP" + }, + "outputs": [], + "source": [ + "# Install pytorch with Cuda\n", + "\n", + "%cd /content/\n", + "%pip install --upgrade pip\n", + "%pip install -U torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118" + ] + }, + { + "cell_type": "code", + "source": [ + "# Installing TinyCuda\n", + "%cd /content/\n", + "!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" + ], + "metadata": { + "id": "RGa9Fwz9Gu2Z" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Install nerfstudio\n", + "%cd /content/\n", + "%pip install \"nerfstudio>=1.0.0\"" + ], + "metadata": { + "id": "hJeIoUaIG0y5" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Install dependency to run the training with logs (see https://github.com/nerfstudio-project/nerfstudio/pull/2870 for more context)\n", + "%pip install colab-xterm" + ], + "metadata": { + "id": "WrIMSAvvItuF" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Download the training dataset" + ], + "metadata": { + "id": "LSWsBhQoG-kx" + } + }, + { + "cell_type": "code", + "source": [ + "%%time\n", + "%cd /content/\n", + "!ns-download-data sitcoms3d" + ], + "metadata": { + "id": "zuN8Ya9yuULS" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "# Setup everything\n", + "\n", + "Create your free account at [Comet.com](https://www.comet.com/site/?utm_source=nerf&utm_medium=referral&utm_content=colab) and grab your API key which can be found under account settings." + ], + "metadata": { + "id": "VXLnZCf5HCxc" + } + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EJBXMB-fF-rB" + }, + "outputs": [], + "source": [ + "# Setup Comet\n", + "\n", + "import comet_ml\n", + "\n", + "comet_ml.init()" + ] + }, + { + "cell_type": "markdown", + "source": [ + "# Training" + ], + "metadata": { + "id": "h_ipb05iHJ-Y" + } + }, + { + "cell_type": "code", + "source": [ + "# First let's import everything that we need and define some helper functions\n", + "import os\n", + "import uuid\n", + "from pathlib import Path\n", + "\n", + "\n", + "def generate_guid():\n", + " # type: () -> str\n", + " \"\"\"Generate a GUID\"\"\"\n", + " return uuid.uuid4().hex" + ], + "metadata": { + "id": "dbAXCC5ZHMf2" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Let's define some parameters\n", + "MAX_ITERATIONS = 30000\n", + "NERF_EXPERIMENT_NAME = \"TBBT-nerfstudio\"" + ], + "metadata": { + "id": "xHfnp_mWHVCO" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Let's generate an unique Comet Experiment Key, this will allows us to save everything to a single experiment\n", + "exp_key = generate_guid()\n", + "os.environ[\"COMET_EXPERIMENT_KEY\"] = exp_key" + ], + "metadata": { + "id": "UlhoWknCHbJ5" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Start the training in a real terminal to see logs\n", + "%load_ext colabxterm\n", + "%env TERM=xterm\n", + "\n", + "from IPython.display import clear_output\n", + "\n", + "clear_output(wait=True)\n", + "\n", + "print(\n", + " \"\\033[1m\"\n", + " + \"Copy and paste the following command into the terminal window that pops up under this cell.\"\n", + " + \"\\033[0m\"\n", + ")\n", + "print(\n", + " f\"ns-train nerfacto --vis viewer+comet --viewer.websocket-port 7007 --viewer.make-share-url True --project-name comet-example-TBBT-nerfstudio --experiment-name {NERF_EXPERIMENT_NAME} --max-num-iterations {MAX_ITERATIONS} --logging.steps-per-log 100 sitcoms3d-data --data data/sitcoms3d/TBBT-big_living_room/ --downscale-factor 4\"\n", + ")\n", + "print()\n", + "\n", + "%xterm" + ], + "metadata": { + "id": "bqbu3E8FLU8Z" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Train the model\n", + "%cd /content\n", + "!env PYTHONUNBUFFERED=true ns-train nerfacto \\\n", + " --vis comet \\\n", + " --logging.local-writer.max-log-size 0 \\\n", + " --project-name comet-example-TBBT-nerfstudio \\\n", + " --experiment-name {NERF_EXPERIMENT_NAME} \\\n", + " --max-num-iterations {MAX_ITERATIONS}\\\n", + " --pipeline.datamanager.images-on-gpu True\\\n", + " --logging.steps-per-log 100\\\n", + " sitcoms3d-data \\\n", + " --data data/sitcoms3d/TBBT-big_living_room/ \\\n", + " --downscale-factor 4" + ], + "metadata": { + "id": "EYaXKAqDHjve" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Log the model to Comet\n", + "base_dir = Path(\"./outputs/\") / NERF_EXPERIMENT_NAME / \"nerfacto\"\n", + "# Get the latest trained model\n", + "training_run_dir = sorted(base_dir.iterdir())[-1]\n", + "print(training_run_dir)" + ], + "metadata": { + "id": "-YhAiT_WHrAH" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "experiment = comet_ml.ExistingExperiment(experiment_key=exp_key)\n", + "experiment.log_model(\n", + " \"nerfstudio-model\",\n", + " str(training_run_dir),\n", + " prepend_folder_name=False,\n", + " metadata={\"nerfstudio_project_dir\": str(training_run_dir)},\n", + ")" + ], + "metadata": { + "id": "WIYvYAPkLoH3" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# Generate validation videos\n", + "# Get the validation camera paths from Comet\n", + "config_filename = training_run_dir / \"config.yml\"\n", + "\n", + "artifact = experiment.get_artifact(\"TBTT-Nerfstudio-Camera-Paths\", workspace=\"examples\")\n", + "output_path = Path(\"./validation/\")\n", + "artifact.download(output_path)\n", + "\n", + "# And for each of them, generate a video that we log to Comet\n", + "for camera_path in (output_path / \"camera_paths\").glob(\"camera_path_*.json\"):\n", + " validation_name = camera_path.stem\n", + " output_video_path = \"./renders/%s.mp4\" % (validation_name)\n", + "\n", + " !ns-render camera-path --load-config $config_filename --camera-path-filename $camera_path --output-path $output_video_path\n", + "\n", + " experiment.log_video(output_video_path)" + ], + "metadata": { + "id": "UHpbgfO7Hs-0" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "# And finally closes the Comet Experiment to save the Notebook code and flush everything to Comet\n", + "experiment.end()" + ], + "metadata": { + "id": "YC9OCEVMH5Tt" + }, + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.8.13" + }, + "vscode": { + "interpreter": { + "hash": "c59f626636933ef1dc834fb3684b382f705301c5306cf8436d2da634c2289783" + } + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file