notebooks/community/model_garden/model_garden_camp_zipnerf.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-e1HpvsDh34Q"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L5o1Ggr5h34U"
      },
      "source": [
        "# Vertex AI Model Garden - CamP ZipNeRF (Jax) Notebook\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_camp_zipnerf.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_camp_zipnerf.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U-SERmqUh34V"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V6QmW0Doh34W"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates a [jax implementation](https://github.com/jonbarron/camp_zipnerf) of [CamP: Camera Preconditioning\n",
        "for Neural Radiance Fields](https://camp-nerf.github.io/) for training and rendering Neural Radiance Fields (NeRFs) more efficiently. It is primarily aimed at addressing some of the limitations of traditional NeRF techniques, which, while powerful for creating detailed 3D models from 2D images, can be computationally intensive and slow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vkSMThcKh34W"
      },
      "source": [
        "## Objective\n",
        "\n",
        "In this tutorial, you learn how to:\n",
        "\n",
        "- Use [COLMAP](https://colmap.github.io/) to perform Structure from Motion (SfM), a technique that estimates the three-dimensional structure of a scene from a series of two-dimensional images.\n",
        "- Calibrate, train and render NERF scenes using [Vertex AI custom jobs](https://cloud.google.com/vertex-ai/docs/samples/aiplatform-create-custom-job-sample).\n",
        "- Render a video along a custom camera path using a series of keyframe photos.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Training\n",
        "- Vertex AI Custom Job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "myi4N60Xh34W"
      },
      "source": [
        "## Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vofRExleAA8k"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qayv5ifRh34Y"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "import os\n",
        "import uuid\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, please change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "\n",
        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "staging_bucket = os.path.join(BUCKET_URI, \"zipnerf_staging\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=staging_bucket)\n",
        "\n",
        "# The pre-built calibration docker image.\n",
        "CALIBRATION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-calibrate:latest\"\n",
        "# The pre-built training docker image.\n",
        "TRAINING_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-cloudnerf-train:latest\"\n",
        "# The pre-built rendering docker image.\n",
        "RENDERING_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-cloudnerf-render:latest\"\n",
        "\n",
        "import subprocess\n",
        "from typing import Any, List\n",
        "\n",
        "IMAGE_EXTENSIONS = (\".png\", \".jpg\", \".jpeg\", \".gif\", \".bmp\")\n",
        "GCS_API_ENDPOINT = \"https://storage.cloud.google.com/\"\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def get_mp4_video_link(mp4_rendering_path: str) -> str:\n",
        "    # Define the gsutil command.\n",
        "    command = f\"gsutil ls {mp4_rendering_path}\"\n",
        "\n",
        "    # Run the command and capture the output.\n",
        "    try:\n",
        "        result = subprocess.check_output(command, shell=True, text=True)\n",
        "        # Split the result by newlines to get a list of files.\n",
        "        file_list = result.strip().split(\"\\n\")\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        print(f\"An error occurred: {e}\")\n",
        "        file_list = []\n",
        "    mp4_video_link = file_list[0].replace(\"gs://\", GCS_API_ENDPOINT)\n",
        "    return mp4_video_link\n",
        "\n",
        "\n",
        "def write_keyframe_list_to_gcs(\n",
        "    bucket_path: str, output_gcs_file: str, max_files: int = 10\n",
        ") -> List[Any]:\n",
        "    # Get the list of files in the GCS bucket.\n",
        "    cmd = f\"gsutil ls {bucket_path}\"\n",
        "    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
        "\n",
        "    if result.returncode != 0:\n",
        "        print(\"Error listing GCS bucket:\", result.stderr)\n",
        "        return []\n",
        "\n",
        "    # Filter for image files and extract file names.\n",
        "    files = result.stdout.splitlines()\n",
        "    image_files = [\n",
        "        os.path.basename(f) for f in files if f.lower().endswith(IMAGE_EXTENSIONS)\n",
        "    ]\n",
        "\n",
        "    output_file = \"out.txt\"\n",
        "    with open(output_file, \"w\") as file:\n",
        "        for name in image_files[:max_files]:\n",
        "            file.write(name + \"\\n\")\n",
        "\n",
        "    cmd = f\"gsutil cp {output_file} {output_gcs_file}\"\n",
        "    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
        "\n",
        "    if result.returncode != 0:\n",
        "        print(\"Error listing GCS bucket:\", result.stderr)\n",
        "        return []"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OBhvKerXh34a"
      },
      "outputs": [],
      "source": [
        "# @title Prepare dataset\n",
        "# @markdown Mip-NeRF 360 dataset contains the following 9 scenes:\n",
        "# @markdown - `bicycle`\n",
        "# @markdown - `bonsai`\n",
        "# @markdown - `counter`\n",
        "# @markdown - `flowers`\n",
        "# @markdown - `garden`\n",
        "# @markdown - `kitchen`\n",
        "# @markdown - `room`\n",
        "# @markdown - `stump`\n",
        "# @markdown - `treehill`\n",
        "\n",
        "# @markdown Please note that `flowers` and `treehill` require author's permission. Each scene comes preprocessed with COLMAP information so the calibration step in the following section is optional.\n",
        "# @markdown If you need to prepare your dataset and store it on Cloud Storage, then the following example shows how to do this for the [mipnerf360 dataset](https://jonbarron.info/mipnerf360/).\n",
        "\n",
        "\n",
        "mipnerf_dataset_directory = \"mipnerf360_dataset\"  # @param {type:\"string\"}\n",
        "MIPNERF_DATA_GCS_PATH = os.path.join(BUCKET_URI, mipnerf_dataset_directory)\n",
        "\n",
        "# Download the bicycle scene data to a local directory.\n",
        "! rm -rf $mipnerf_dataset_directory\n",
        "! mkdir -p $mipnerf_dataset_directory\n",
        "! wget -P $mipnerf_dataset_directory http://storage.googleapis.com/gresearch/refraw360/garden.zip\n",
        "\n",
        "# Unzip the mipnerf360 garden dataset.\n",
        "! unzip $mipnerf_dataset_directory/garden.zip -d $mipnerf_dataset_directory\n",
        "\n",
        "# Move mipnerf360 data from local directory to Cloud Storage.\n",
        "# This step takes a few minutes to finish.\n",
        "! gsutil -m cp -R $mipnerf_dataset_directory/* $MIPNERF_DATA_GCS_PATH"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WdWHe-RmSEn6"
      },
      "source": [
        "## NERF pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "j1evctm2h34g"
      },
      "outputs": [],
      "source": [
        "# @title Run Camera Pose Estimation Custom Job\n",
        "\n",
        "# @markdown Once data and experiment paths have been configured, run the custom job below.\n",
        "\n",
        "# @markdown The following parameters are required:\n",
        "\n",
        "# @markdown * `use_gpu`: Whether to use GPU or not.\n",
        "# @markdown * `gcs_dataset_path`: Path to image folder in GCS dataset.\n",
        "# @markdown * `gcs_experiment_path`: GCS path for storing experiment outputs.\n",
        "# @markdown * `camera`: Type of camera used. `OPENCV` for perspective, `OPENCV_FISHEYE` for fisheye.\n",
        "\n",
        "# @markdown The custom job will run on the images in the `gcs_dataset_path` folder and store the colmap outputs in the `gcs_experiment_path/data` folder.\n",
        "\n",
        "# @markdown On the scenes in this current dataset, this step takes about 30 minutes.\n",
        "\n",
        "# Folder containing all the images of the garden scene.\n",
        "# e.g. f\"{BUCKET_URI}/{mipnerf_dataset_directory}/garden/images\"\n",
        "INPUT_IMAGES_FOLDER = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Folder for storing experiment outputs for calibration, training and rendering.\n",
        "# e.g. f\"{BUCKET_URI}/{mipnerf_dataset_directory}/exp/garden\"\n",
        "OUTPUT_FOLDER = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# This job will run colmap camera pose estimation.\n",
        "data_calibration_job_name = get_job_name_with_datetime(\"colmap\")\n",
        "\n",
        "# Worker pool spec.\n",
        "machine_type = \"n1-highmem-64\"\n",
        "num_nodes = 1\n",
        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
        "num_gpus = 8\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": gpu_type,\n",
        "            \"accelerator_count\": num_gpus,\n",
        "        },\n",
        "        \"replica_count\": num_nodes,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": CALIBRATION_DOCKER_URI,\n",
        "            \"args\": [\n",
        "                \"-use_gpu\",\n",
        "                \"1\",\n",
        "                \"-gcs_dataset_path\",\n",
        "                INPUT_IMAGES_FOLDER,\n",
        "                \"-gcs_experiment_path\",\n",
        "                OUTPUT_FOLDER,\n",
        "                \"-camera\",\n",
        "                \"OPENCV\",\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "data_calibration_custom_job = aiplatform.CustomJob(\n",
        "    display_name=data_calibration_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=staging_bucket,\n",
        ")\n",
        "\n",
        "data_calibration_custom_job.run()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "q7ZWhSpjh34g"
      },
      "outputs": [],
      "source": [
        "# @title Training the ZipNeRF model\n",
        "\n",
        "# @markdown Once the Colmap pose calibration is completed, we can run training.\n",
        "\n",
        "# @markdown The following parameters are required:\n",
        "\n",
        "# @markdown * `gcs_experiment_path`: GCS path for loading processed dataset and storing experiment outputs.\n",
        "# @markdown * `factor`: A factor of the downsampled images in the preprocessing step that affects the resolution or detail level of the training pixel ground truth and rendered images. A factor of 2 is recommended for indoor scenes and a factor of 4 for outdoor scenes.\n",
        "\n",
        "# @markdown The custom job will run on the images in the `gcs_experiment_path/data` colmap dataset and outputs in the checkpoints in `gcs_experiment_path/checkpoints` folder.\n",
        "\n",
        "# @markdown Depending on the configuration, this step could take up to 3 hours.\n",
        "\n",
        "# This job will run zipnerf training.\n",
        "\n",
        "# This is the nerf training job name. You will use it to load the checkpoints\n",
        "# in the rendering job for the current run.\n",
        "nerf_training_job_name = get_job_name_with_datetime(\"nerf_training\")\n",
        "\n",
        "FACTOR = 0  # @param [0, 2, 4, 8]\n",
        "\n",
        "# Worker pool spec.\n",
        "machine_type = \"n1-highmem-64\"\n",
        "num_nodes = 1\n",
        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
        "num_gpus = 8\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": gpu_type,\n",
        "            \"accelerator_count\": num_gpus,\n",
        "        },\n",
        "        \"replica_count\": num_nodes,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAINING_DOCKER_URI,\n",
        "            \"args\": [\n",
        "                \"-training_job_name\",\n",
        "                nerf_training_job_name,\n",
        "                \"-gcs_experiment_path\",\n",
        "                OUTPUT_FOLDER,\n",
        "                \"-factor\",\n",
        "                str(FACTOR),\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "nerf_training_custom_job = aiplatform.CustomJob(\n",
        "    display_name=nerf_training_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=staging_bucket,\n",
        ")\n",
        "\n",
        "nerf_training_custom_job.run(enable_web_access=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7lufsDWzh34g"
      },
      "outputs": [],
      "source": [
        "# @title Rendering the ZipNeRF model (360)\n",
        "\n",
        "# @markdown Once the training is completed, we can run rendering.\n",
        "\n",
        "# @markdown The following parameters are required:\n",
        "\n",
        "# @markdown * `gcs_experiment_path`: GCS path for loading processed dataset and storing experiment outputs.\n",
        "# @markdown * `render_video_fps`: Frame rate of rendered video.\n",
        "# @markdown * `render_path_frames`: Number of frames to render for a path.\n",
        "# @markdown * `render_resolution`: Standard display resolutions, for example: (VIDEO_WIDTH, VIDEO_HEIGHT).\n",
        "\n",
        "# @markdown The custom job will run on the images in the `gcs_experiment_path/data` colmap dataset and outputs in the checkpoints in `gcs_experiment_path/checkpoints` folder.\n",
        "\n",
        "# This job will run zipnerf rendering.\n",
        "nerf_rendering_job_name = get_job_name_with_datetime(\"nerf_rendering\")\n",
        "VIDEO_WIDTH = 1280  # @param {type:\"integer\"}\n",
        "VIDEO_HEIGHT = 720  # @param {type:\"integer\"}\n",
        "RENDER_PATH_FRAMES = 150  # @param {type:\"integer\"}\n",
        "RENDER_VIDEO_FPS = 30  # @param {type:\"integer\"}\n",
        "VIDEO_RESOLUTION = f\"({VIDEO_WIDTH}, {VIDEO_HEIGHT})\"\n",
        "\n",
        "# Worker pool spec.\n",
        "machine_type = \"n1-highmem-64\"\n",
        "num_nodes = 1\n",
        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
        "num_gpus = 8\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": gpu_type,\n",
        "            \"accelerator_count\": num_gpus,\n",
        "        },\n",
        "        \"replica_count\": num_nodes,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": RENDERING_DOCKER_URI,\n",
        "            \"args\": [\n",
        "                \"-rendering_job_name\",\n",
        "                nerf_rendering_job_name,\n",
        "                \"-training_job_name\",\n",
        "                nerf_training_job_name,\n",
        "                \"-gcs_experiment_path\",\n",
        "                OUTPUT_FOLDER,\n",
        "                \"-render_video_fps\",\n",
        "                str(RENDER_VIDEO_FPS),\n",
        "                \"-render_path_frames\",\n",
        "                str(RENDER_PATH_FRAMES),\n",
        "                \"-render_resolution\",\n",
        "                VIDEO_RESOLUTION,\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "nerf_rendering_custom_job = aiplatform.CustomJob(\n",
        "    display_name=nerf_rendering_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=staging_bucket,\n",
        ")\n",
        "\n",
        "nerf_rendering_custom_job.run(enable_web_access=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ifhDb9xeh34g"
      },
      "outputs": [],
      "source": [
        "# @title Show rendered video from GCS\n",
        "\n",
        "from IPython.display import Video\n",
        "\n",
        "MP4_RENDERING_PATH = (\n",
        "    f\"{OUTPUT_FOLDER}/render/{nerf_rendering_job_name}/path_videos/videos/*color.mp4\"\n",
        ")\n",
        "mp4_video_link = get_mp4_video_link(MP4_RENDERING_PATH)\n",
        "Video(mp4_video_link)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Gai5cc-bh34g"
      },
      "outputs": [],
      "source": [
        "# @title Rendering the ZipNeRF model (custom camera trajectory)\n",
        "\n",
        "# @markdown Create keyframe file list for rendering custom camera trajectories.\n",
        "\n",
        "# @markdown To create a custom camera trajectory in a Neural Radiance Field (NeRF) model using images from the same dataset used for training, you can generate a keyframe file list where each keyframe corresponds to the name of an image file stored in a Google Cloud Storage (GCS) bucket. This section will guide you through creating this keyframe file list.\n",
        "\n",
        "# @markdown Step 1: Identifying keyframe images\n",
        "# @markdown First, identify the images within your dataset that you want to use as keyframes. These images should ideally represent the significant views or angles that you want your camera trajectory to include.\n",
        "\n",
        "# @markdown Step 2: Creating a list of image file names\n",
        "# @markdown Access Your GCS Bucket: Navigate to your GCS bucket where the dataset is stored.\n",
        "\n",
        "# @markdown Select Image Files: Choose the specific image files that you want to use as keyframes. Remember, these should be files used in training the NeRF model, as they will have corresponding camera parameters already defined.\n",
        "\n",
        "# @markdown Compile File Names: Create a list of the file names (not the paths) of these selected images. Ensure that each file name is on a separate line. For example:\n",
        "\n",
        "\n",
        "# This job will run zipnerf rendering.\n",
        "nerf_custom_rendering_job_name = get_job_name_with_datetime(\"nerf_custom_rendering\")\n",
        "\n",
        "# Example usage.\n",
        "KEYFRAME_IMAGE_FILELIST = (\n",
        "    f\"{OUTPUT_FOLDER}/keyframe_list_{nerf_custom_rendering_job_name}.txt\"\n",
        ")\n",
        "max_files = 30  # Set this to the number of files you want\n",
        "write_keyframe_list_to_gcs(\n",
        "    INPUT_IMAGES_FOLDER, KEYFRAME_IMAGE_FILELIST, max_files=max_files\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Hi0oIbyZh34h"
      },
      "outputs": [],
      "source": [
        "# @title Run rendering for custom path\n",
        "\n",
        "# @markdown Once the training is completed, we can run rendering.\n",
        "\n",
        "# @markdown The following parameters are required:\n",
        "\n",
        "# @markdown * `gcs_experiment_path`: GCS path for loading processed dataset and storing experiment outputs.\n",
        "# @markdown * `render_video_fps`: Frame rate of rendered video.\n",
        "# @markdown * `render_resolution`: Standard display resolutions, for example: (VIDEO_WIDTH, VIDEO_HEIGHT).\n",
        "# @markdown * `keyframe_image_list`: List of image filename, one per line, for rendering custom camera path.\n",
        "\n",
        "# @markdown With keyframes, an interpolated path is generated. This path represents a smoothly contoured spline that interconnects the specified keyframe camera poses. The process utilizes a configuration variable, `render_spline_n_interp`, which is preset to a default value of 30. As a result, the finalized interpolated path comprises a total of `render_spline_n_interp` * (n - 1) poses. In the specific scenario under discussion, the config.render_spline_n_interp is configured to 30. **With an input of 30 keyframes, the calculation yields a total of 30 * 29, amounting to 870 poses**.\n",
        "\n",
        "# This job will run zipnerf rendering.\n",
        "# Worker pool spec.\n",
        "machine_type = \"n1-highmem-64\"\n",
        "num_nodes = 1\n",
        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
        "num_gpus = 8\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": gpu_type,\n",
        "            \"accelerator_count\": num_gpus,\n",
        "        },\n",
        "        \"replica_count\": num_nodes,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": RENDERING_DOCKER_URI,\n",
        "            \"args\": [\n",
        "                \"-rendering_job_name\",\n",
        "                nerf_custom_rendering_job_name,\n",
        "                \"-training_job_name\",\n",
        "                nerf_training_job_name,\n",
        "                \"-gcs_experiment_path\",\n",
        "                OUTPUT_FOLDER,\n",
        "                \"-render_resolution\",\n",
        "                VIDEO_RESOLUTION,\n",
        "                \"-gcs_keyframes_file\",\n",
        "                KEYFRAME_IMAGE_FILELIST,\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "nerf_custom_rendering_custom_job = aiplatform.CustomJob(\n",
        "    display_name=nerf_custom_rendering_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=staging_bucket,\n",
        ")\n",
        "\n",
        "nerf_custom_rendering_custom_job.run(enable_web_access=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xeup-oLAh34h"
      },
      "outputs": [],
      "source": [
        "# @title Show rendered video from GCS\n",
        "\n",
        "from IPython.display import Video\n",
        "\n",
        "MP4_RENDERING_PATH = (\n",
        "    f\"{OUTPUT_FOLDER}/render/{nerf_rendering_job_name}/path_videos/videos/*color.mp4\"\n",
        ")\n",
        "mp4_video_link = get_mp4_video_link(MP4_RENDERING_PATH)\n",
        "Video(mp4_video_link)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "K9K-sK6INmDP"
      },
      "outputs": [],
      "source": [
        "# @title Clean up resources\n",
        "# @markdown Delete the experiment finished jobs and bucket to avoid\n",
        "# @markdown unnecessary continouous charges that may incur.\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME\n",
        "\n",
        "# Delete pose estimation, training and rendering custom jobs.\n",
        "if data_calibration_custom_job.list(\n",
        "    filter=f'display_name=\"{data_calibration_job_name}\"'\n",
        "):\n",
        "    data_calibration_custom_job.delete()\n",
        "if nerf_training_custom_job.list(filter=f'display_name=\"{nerf_training_job_name}\"'):\n",
        "    nerf_training_custom_job.delete()\n",
        "if nerf_rendering_custom_job.list(filter=f'display_name=\"{nerf_rendering_job_name}\"'):\n",
        "    nerf_rendering_custom_job.delete()\n",
        "if nerf_custom_rendering_custom_job.list(\n",
        "    filter=f'display_name=\"{nerf_custom_rendering_job_name}\"'\n",
        "):\n",
        "    nerf_custom_rendering_custom_job.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_camp_zipnerf.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}