From 146332cfb629c5e8bbee9b125892095e56094d3b Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Thu, 2 Jun 2022 14:58:46 -0400 Subject: [PATCH] Refactoring (data loading, semantic-nerf/nerf-w fixes, ...) (#8) * in progress. camera_indices broken * faster dataloader. starting on pixel sampler logic * imagesampler and pixelsampler * lint and black * continuing merge * pytest working * fixed docs * docs * minor changes --- .vscode/launch.json | 2 +- README.md | 11 +- configs/data/default_data.yaml | 17 +- configs/instant_ngp.yaml | 2 +- configs/instant_ngp_tcnn.yaml | 4 +- configs/logging/default_logging.yaml | 2 +- configs/nerfw.yaml | 12 +- configs/semantic_nerf.yaml | 12 +- configs/tests/test_dryrun.yaml | 5 +- configs/vanilla_nerf.yaml | 9 +- docs/dataset.rst | 7 - docs/pyrad.rst | 18 +- environment/requirements.txt | 7 +- notebooks/notebook_imports.py | 1 - notebooks/visualize_dataset.ipynb | 203 ++++----------- pyproject.toml | 2 +- .../{mattport_cuda.cpp => pyrad_cuda.cpp} | 0 ...rt_cuda_kernel.cu => pyrad_cuda_kernel.cu} | 0 pyrad/nerf/colliders.py | 12 +- pyrad/nerf/dataset/collate.py | 67 ----- pyrad/nerf/dataset/format/friends.py | 17 +- pyrad/nerf/dataset/image_dataset.py | 171 ++++++------- pyrad/nerf/dataset/structs.py | 22 +- pyrad/nerf/dataset/utils.py | 43 +++- pyrad/nerf/field_modules/field_heads.py | 13 +- pyrad/nerf/fields/nerf_field.py | 2 +- pyrad/nerf/fields/nerfw_field.py | 108 ++++++++ pyrad/nerf/graph/base.py | 29 +-- pyrad/nerf/graph/instant_ngp.py | 6 +- pyrad/nerf/graph/nerfw.py | 236 ++++++++---------- pyrad/nerf/graph/semantic_nerf.py | 221 ++++++++-------- pyrad/nerf/graph/vanilla_nerf.py | 14 +- pyrad/nerf/image_sampler.py | 75 ++++++ pyrad/nerf/pixel_sampler.py | 68 +++++ pyrad/nerf/{sampler.py => ray_sampler.py} | 0 pyrad/nerf/trainer.py | 87 +++---- pyrad/structures/cameras.py | 12 +- pyrad/structures/rays.py | 11 +- {radiance => pyrad}/utils/README.md | 0 pyrad/utils/io.py | 33 ++- pyrad/utils/misc.py | 25 ++ pyrad/utils/visualization.py | 1 + scripts/run_data_preprocessor.py | 34 +++ scripts/run_dataset_processor.py | 6 - scripts/run_train.py | 4 +- tests/nerf/__init__.py | 0 .../nerf/dataloader/configs/blender_lego.yaml | 21 -- tests/nerf/dataloader/test_ray_generator.py | 88 ------- tests/nerf/test_sampler.py | 2 +- tests/test_run_train.py | 3 + 50 files changed, 914 insertions(+), 831 deletions(-) rename pyrad/cuda/{mattport_cuda.cpp => pyrad_cuda.cpp} (100%) rename pyrad/cuda/{mattport_cuda_kernel.cu => pyrad_cuda_kernel.cu} (100%) delete mode 100644 pyrad/nerf/dataset/collate.py create mode 100644 pyrad/nerf/fields/nerfw_field.py create mode 100644 pyrad/nerf/image_sampler.py create mode 100644 pyrad/nerf/pixel_sampler.py rename pyrad/nerf/{sampler.py => ray_sampler.py} (100%) rename {radiance => pyrad}/utils/README.md (100%) create mode 100644 scripts/run_data_preprocessor.py delete mode 100644 scripts/run_dataset_processor.py delete mode 100644 tests/nerf/__init__.py delete mode 100644 tests/nerf/dataloader/configs/blender_lego.yaml delete mode 100644 tests/nerf/dataloader/test_ray_generator.py diff --git a/.vscode/launch.json b/.vscode/launch.json index feda559d3..6ee58033a 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -46,7 +46,7 @@ "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "--config-name=nerfw.yaml" + "--config-name=nerfw.yaml", "++data.dataset.use_cache=true" ] }, { diff --git a/README.md b/README.md index 1427b4825..2fcac7c95 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,11 @@ python scripts/run_train.py data/dataset=friends_TBBT-big_living_room # Run with different datasets and config changes python scripts/run_train.py data/dataset=friends_TBBT-big_living_room graph.network.far_plane=14 + +# Speed up the dataloading pipeline. +python scripts/run_data_preprocessor.py +# Then, specify using the cache. +python scripts/run_train.py ++data.dataset.use_cache=true ``` # Getting around the codebase @@ -66,10 +71,10 @@ The NeRF graph definitions can be found in `nerf/graph/`. Each implementation of To implement any pre-existing NeRF that we have not yet implemented under `nerf/graph/`, create a new graph structure by using provided modules or any new module you define. Then create an associated config making sure `__target__` points to your NeRF class (see [here](./configs/README.md) for more info on how to create the config). Then run training as described above. - # Feature Documentation + ### 1. [Hydra config structure](./configs/README.md) -### 2. [Logging, debugging utilities](./radiance/utils/README.md) +### 2. [Logging, debugging utilities](./pyrad/utils/README.md) ### 3. [Benchmarking, other tooling](./scripts/README.md) ### 4. Running other repos with our data @@ -94,11 +99,11 @@ pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/to ``` To run instant-ngp with tcnn, you can do the following. This is with the fox dataset. + ``` python scripts/run_train.py --config-name=instant_ngp_tcnn.yaml data/dataset=instant_ngp_fox ``` - ### 6. Setting up Jupyter ``` diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index 468770169..839c8b24e 100644 --- a/configs/data/default_data.yaml +++ b/configs/data/default_data.yaml @@ -4,13 +4,16 @@ defaults: # validation options val_image_indices: [0, 10, 20] # which images to evaluate on -val_downscale_factor: 8 +val_downscale_factor: 4 val_num_rays_per_chunk: 4096 -# dataloader options -dataloader: - num_workers: 1 - shuffle: true - num_images_to_sample_from: 8 +image_dataset: + _target_: pyrad.nerf.dataset.image_dataset.ImageDataset + +image_sampler: + num_images_to_sample_from: 8 # 0 + num_times_to_repeat_images: 100 # -1 + move_to_graph_device: true # false + +pixel_sampler: num_rays_per_batch: 1024 - num_times_to_repeat_images: 100 diff --git a/configs/instant_ngp.yaml b/configs/instant_ngp.yaml index 6e3fe38ae..43c62888d 100644 --- a/configs/instant_ngp.yaml +++ b/configs/instant_ngp.yaml @@ -3,7 +3,7 @@ defaults: - _self_ method_name: instant_ngp -data.dataloader.num_rays_per_batch: 4096 +data.pixel_sampler.num_rays_per_batch: 4096 data.val_num_rays_per_chunk: 16384 graph: diff --git a/configs/instant_ngp_tcnn.yaml b/configs/instant_ngp_tcnn.yaml index 56f501bc5..bf16c7d40 100644 --- a/configs/instant_ngp_tcnn.yaml +++ b/configs/instant_ngp_tcnn.yaml @@ -4,9 +4,11 @@ defaults: method_name: instant_ngp data: - dataloader: + pixel_sampler: num_rays_per_batch: 4096 + val_downscale_factor: 1 val_num_rays_per_chunk: 16384 graph: network: field_implementation: tcnn + steps_per_test: 10000 diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 9dfaf40e1..d1873abec 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -7,7 +7,7 @@ writer: LocalWriter: save_dir: './' stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] - max_log_size: 20 # if 0, logs everything with no erasing + max_log_size: 10 # if 0, logs everything with no erasing # profiler logs run times of functions and prints at end of training enable_profiler: True \ No newline at end of file diff --git a/configs/nerfw.yaml b/configs/nerfw.yaml index 4d275b953..f91c37911 100644 --- a/configs/nerfw.yaml +++ b/configs/nerfw.yaml @@ -6,16 +6,20 @@ defaults: experiment_name: friends_TBBT-big_living_room method_name: nerfw +data: + image_dataset: + # "image", "mask", "semantics" + _target_: pyrad.nerf.dataset.image_dataset.PanopticImageDataset + graph: network: _target_: pyrad.nerf.graph.nerfw.NerfWGraph - near_plane: 0.1 - far_plane: 10.0 + collider_config: + _target_: pyrad.nerf.colliders.AABBBoxCollider num_coarse_samples: 64 num_importance_samples: 64 - loss_coefficients: # loss coefficients rgb_loss_coarse: 1.0 rgb_loss_fine: 1.0 uncertainty_loss: 1.0 - density_loss: 0.01 \ No newline at end of file + density_loss: 0.01 diff --git a/configs/semantic_nerf.yaml b/configs/semantic_nerf.yaml index c735fd04f..4e24131f7 100644 --- a/configs/semantic_nerf.yaml +++ b/configs/semantic_nerf.yaml @@ -6,15 +6,19 @@ defaults: experiment_name: friends_TBBT-big_living_room method_name: semantic_nerf +data: + image_dataset: + # "image", "mask", "semantics" + _target_: pyrad.nerf.dataset.image_dataset.PanopticImageDataset + graph: network: _target_: pyrad.nerf.graph.semantic_nerf.SemanticNerfGraph - near_plane: 0.1 - far_plane: 10.0 + collider_config: + _target_: pyrad.nerf.colliders.AABBBoxCollider num_coarse_samples: 64 num_importance_samples: 64 - loss_coefficients: # loss coefficients rgb_loss_coarse: 1.0 rgb_loss_fine: 1.0 - semantic_loss_fine: 0.05 \ No newline at end of file + semantic_loss_fine: 0.05 diff --git a/configs/tests/test_dryrun.yaml b/configs/tests/test_dryrun.yaml index d627fbbd9..8d3ced9f4 100644 --- a/configs/tests/test_dryrun.yaml +++ b/configs/tests/test_dryrun.yaml @@ -12,9 +12,10 @@ data: data_directory: tests/data/lego_test downscale_factor: 8 val_image_indices: [0] - dataloader: + image_sampler: num_images_to_sample_from: 1 - num_rays_per_batch: 8 + pixel_sampler: + num_rays_per_batch: 4 graph: max_num_iterations: 5 diff --git a/configs/vanilla_nerf.yaml b/configs/vanilla_nerf.yaml index b1d763619..d98181e2a 100644 --- a/configs/vanilla_nerf.yaml +++ b/configs/vanilla_nerf.yaml @@ -9,7 +9,7 @@ graph: # saving model specifications model_dir: /tmp/pyrad_models/ steps_per_save: null - steps_per_test: 100 + steps_per_test: 1000 max_num_iterations: 1000000 # pretrained loading options @@ -19,11 +19,12 @@ graph: network: _target_: pyrad.nerf.graph.vanilla_nerf.NeRFGraph - near_plane: 2.0 - far_plane: 6.0 + collider_config: + _target_: pyrad.nerf.colliders.NearFarCollider + near_plane: 2.0 + far_plane: 6.0 num_coarse_samples: 64 num_importance_samples: 128 - loss_coefficients: # loss coefficients rgb_loss_coarse: 1.0 rgb_loss_fine: 1.0 diff --git a/docs/dataset.rst b/docs/dataset.rst index 24eb99a59..5c6957806 100644 --- a/docs/dataset.rst +++ b/docs/dataset.rst @@ -2,13 +2,6 @@ ============ -Collate ----------------- - -.. automodule:: pyrad.nerf.dataset.collate - :members: - :show-inheritance: - Colmap Utils ---------------- diff --git a/docs/pyrad.rst b/docs/pyrad.rst index 1edbb4ead..fa3eadd94 100644 --- a/docs/pyrad.rst +++ b/docs/pyrad.rst @@ -11,6 +11,20 @@ Dataset dataset +Image Samplers +------------------------------- + +.. automodule:: pyrad.nerf.image_sampler + :members: + :show-inheritance: + +Pixel Samplers +------------------------------- + +.. automodule:: pyrad.nerf.pixel_sampler + :members: + :show-inheritance: + Field Modules ------------------------------- @@ -42,9 +56,9 @@ Renderers :members: :show-inheritance: -Samplers +Ray Samplers ------------------------------- -.. automodule:: pyrad.nerf.sampler +.. automodule:: pyrad.nerf.ray_sampler :members: :show-inheritance: \ No newline at end of file diff --git a/environment/requirements.txt b/environment/requirements.txt index 023c14641..1522648c7 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,7 +1,8 @@ black[jupyter]==22.3.0 +h5py>=2.9.0 hydra-colorlog hydra-core -imageio==2.16.1 +imageio ipywidgets>=7.6 jupyterlab==3.3.4 kaleido @@ -11,9 +12,11 @@ plotly==5.7.0 protobuf==3.20.0 pylint==2.13.4 pytest -tensorboard==2.8.0 +tensorboard==2.9.0 +--extra-index-url https://download.pytorch.org/whl/cu113 torch==1.11.0 torchmetrics[image]==0.8.2 torchtyping +--extra-index-url https://download.pytorch.org/whl/cu113 torchvision==0.12.0 tqdm \ No newline at end of file diff --git a/notebooks/notebook_imports.py b/notebooks/notebook_imports.py index b71f57234..d11e1aab0 100644 --- a/notebooks/notebook_imports.py +++ b/notebooks/notebook_imports.py @@ -1,5 +1,4 @@ # image stuff -import cv2 import imageio import mediapy as media diff --git a/notebooks/visualize_dataset.ipynb b/notebooks/visualize_dataset.ipynb index b89563d65..f91633a63 100644 --- a/notebooks/visualize_dataset.ipynb +++ b/notebooks/visualize_dataset.ipynb @@ -20,15 +20,17 @@ "import torch\n", "from torch.utils.data import DataLoader\n", "\n", - "from pyrad.nerf.dataset.image_dataset import ImageDataset, collate_batch\n", - "from pyrad.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one\n", + "from pyrad.nerf.dataset.image_dataset import ImageDataset, PanopticImageDataset\n", + "from pyrad.nerf.image_sampler import CacheImageSampler\n", + "from pyrad.nerf.pixel_sampler import PixelSampler\n", "from pyrad.nerf.dataset.utils import get_dataset_inputs\n", - "from pyrad.nerf.field_modules.ray_generator import RayGenerator\n", - "from pyrad.nerf.field_modules.scene_bounds_collider import SceneBoundsCollider\n", + "from pyrad.nerf.ray_generator import RayGenerator\n", + "from pyrad.nerf.colliders import SceneBoundsCollider, AABBBoxCollider\n", "from pyrad.structures.rays import RayBundle\n", "from pyrad.utils.io import get_absolute_path\n", "from pyrad.viewer.plotly import get_line_segments_from_lines\n", "from pyrad.structures.cameras import get_camera_model\n", + "from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config\n", "\n", "\n", "from hydra import compose, initialize\n", @@ -38,28 +40,6 @@ "import random" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "012bf657", - "metadata": {}, - "outputs": [], - "source": [ - "import goat\n", - "\n", - "vis = goat.vis_utils.get_vis()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "71f46c79", - "metadata": {}, - "outputs": [], - "source": [ - "goat.vis_utils.show_box_test(vis)" - ] - }, { "cell_type": "code", "execution_count": null, @@ -67,70 +47,9 @@ "metadata": {}, "outputs": [], "source": [ - "with initialize(config_path=\"../configs\"):\n", - " cfg = compose(config_name=\"default.yaml\")\n", - " # cfg = compose(config_name=\"default.yaml\", overrides=[\"data/dataset=friends_TBBT-big_living_room\"])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0b010534", - "metadata": {}, - "outputs": [], - "source": [ - "dataset_inputs = get_dataset_inputs(**cfg.data.dataset, split=\"train\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "e95073a9", - "metadata": {}, - "outputs": [], - "source": [ - "# import goat.threedvis.meshcat.geometry as g\n", - "# import numpy as np\n", - "# point_cloud = dataset_inputs.point_cloud\n", - "# size = 0.1\n", - "# vis[\"point_cloud\"].set_object(\n", - "# g.PointCloud(position=np.transpose(np.array(point_cloud.xyz)),\n", - "# color=np.transpose(np.array(point_cloud.rgb)) / 255.0, size=size)\n", - "# )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "aac97d4a", - "metadata": {}, - "outputs": [], - "source": [ - "# show dataset inputs with three.js\n", - "# goat.vis_utils.show_box_test(vis)\n", - "num_images_to_show = 200\n", - "for idx in tqdm(list(range(len(dataset_inputs.image_filenames)))[:num_images_to_show]):\n", - " image = imageio.imread(dataset_inputs.image_filenames[idx])[:, :, :3]\n", - "\n", - " num_intrinsics_params = len(dataset_inputs.intrinsics[idx])\n", - " camera_class = get_camera_model(num_intrinsics_params)\n", - " camera_to_world = dataset_inputs.camera_to_world[idx]\n", - " camera = camera_class(*dataset_inputs.intrinsics[idx].tolist(), camera_to_world=camera_to_world)\n", - " K = camera.get_intrinsics_matrix()\n", - "\n", - " name = \"train/{:06d}\".format(idx)\n", - " # height = 0.5\n", - " displayed_focal_length = 0.5\n", - " goat.vis_utils.draw_camera_frustum(\n", - " vis,\n", - " name=name,\n", - " image=image,\n", - " K=K.numpy(),\n", - " pose=camera_to_world.double().numpy(),\n", - " height=3.0,\n", - " realistic=False,\n", - " displayed_focal_length=0.0,\n", - " )" + "with initialize(version_base=None, config_path=\"../configs\"):\n", + " config = compose(config_name=\"default_setup.yaml\")\n", + "dataset_inputs = get_dataset_inputs(**config.data.dataset, split=\"train\")" ] }, { @@ -140,31 +59,34 @@ "metadata": {}, "outputs": [], "source": [ - "image_dataset = ImageDataset(\n", - " image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n", - ")\n", - "iter_dataset = CollateIterDataset(\n", - " image_dataset,\n", - " collate_fn=lambda batch_list: collate_batch(\n", - " batch_list, cfg.data.dataloader.num_rays_per_batch, keep_full_image=True\n", - " ),\n", - " num_samples_to_collate=cfg.data.dataloader.num_images_to_sample_from,\n", - " num_times_to_repeat=cfg.data.dataloader.num_times_to_repeat_images,\n", - ")\n", - "dataloader = DataLoader(\n", - " iter_dataset,\n", - " batch_size=1,\n", - " num_workers=cfg.data.dataloader.num_workers,\n", - " collate_fn=collate_batch_size_one,\n", - " pin_memory=True,\n", - ")\n", - "dataloader_iter = iter(dataloader)\n", + "device = \"cpu\"\n", + "\n", + "train_image_dataset = instantiate_from_dict_config(\n", + " config.data.image_dataset,\n", + " image_filenames=dataset_inputs.image_filenames,\n", + " downscale_factor=dataset_inputs.downscale_factor,\n", + " semantics=dataset_inputs.semantics,\n", + " alpha_color=dataset_inputs.alpha_color,\n", + ") # ImageDataset\n", + "train_image_sampler = CacheImageSampler(\n", + " train_image_dataset,\n", + " num_samples_to_collate=len(train_image_dataset)\n", + " if config.data.image_sampler.num_images_to_sample_from == 0\n", + " else config.data.image_sampler.num_images_to_sample_from,\n", + " num_times_to_repeat_images=config.data.image_sampler.num_times_to_repeat_images,\n", + " device=device if config.data.image_sampler.move_to_graph_device else \"cpu\",\n", + ") # ImageSampler\n", + "train_pixel_sampler = PixelSampler(\n", + " num_rays_per_batch=config.data.pixel_sampler.num_rays_per_batch, keep_full_image=True\n", + ") # PixelSampler\n", "\n", "ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)\n", "\n", + "iter_train_image_sampler = iter(train_image_sampler)\n", "num_batches = 10\n", "for _ in tqdm(range(num_batches)):\n", - " batch = next(dataloader_iter)\n", + " image_batch = next(iter_train_image_sampler)\n", + " batch = train_pixel_sampler.sample(image_batch)\n", " ray_bundle = ray_generator.forward(batch[\"indices\"])\n", " break" ] @@ -176,7 +98,7 @@ "metadata": {}, "outputs": [], "source": [ - "media.show_image(image_dataset.get_image(10))" + "media.show_image(train_image_dataset.get_image(10), height=200)" ] }, { @@ -188,20 +110,22 @@ "source": [ "def show_batch(batch):\n", " # set the color of the sampled rays\n", + " print(batch.keys())\n", " c, y, x = [i.flatten() for i in torch.split(batch[\"local_indices\"], 1, dim=-1)]\n", - " batch[\"image\"][c, y, x] = 1.0\n", + " batch[\"image\"][c, y, x] = 0.0\n", "\n", " # batch[\"image\"] is num_images, h, w, 3\n", " images = torch.split(batch[\"image\"], 1, dim=0)\n", " image_list = [image[0] for image in images]\n", " image = torch.cat(image_list, dim=1) # cat along the width dimension\n", "\n", - " # the white pixels are rays\n", + " # the black pixels are rays\n", " media.show_image((image * 255).to(torch.uint8))\n", "\n", "\n", "def sample_and_show_batch():\n", - " batch = next(iter(dataloader))\n", + " image_batch = next(iter_train_image_sampler)\n", + " batch = train_pixel_sampler.sample(image_batch)\n", " show_batch(batch)\n", " return batch" ] @@ -256,20 +180,10 @@ "metadata": {}, "outputs": [], "source": [ - "scene_bounds_collider = SceneBoundsCollider(dataset_inputs.scene_bounds)\n", + "scene_bounds_collider = AABBBoxCollider(dataset_inputs.scene_bounds)\n", "intersected_ray_bundle = scene_bounds_collider(ray_bundle)" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "9ca1b3c3", - "metadata": {}, - "outputs": [], - "source": [ - "# intersected_ray_bundle.nears" - ] - }, { "cell_type": "code", "execution_count": null, @@ -277,12 +191,10 @@ "metadata": {}, "outputs": [], "source": [ - "# near_plane = cfg.graph.network.near_plane\n", - "# far_plane = cfg.graph.network.far_plane\n", "lines = torch.stack(\n", " [\n", - " intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.nears,\n", - " intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.fars,\n", + " intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.nears[..., None],\n", + " intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.fars[..., None],\n", " ],\n", " dim=1,\n", ").tolist() # (num_rays, 2, 3)\n", @@ -323,34 +235,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7aa506fe-66a1-4284-ade4-8c90fde5476a", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f259edb0-4126-4930-99f4-b5574dbb2eea", - "metadata": {}, - "outputs": [], - "source": [ - "fig.write_image(\"temp.png\")\n", - "plotly_image = imageio.imread(\"temp.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d1a11772-4692-4800-b446-cfdbe4ab92fa", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "722f5370-3b0f-4ad1-87a1-df479ae85b60", + "id": "fd04703a-aa50-45e4-95fd-4d881ff18d61", "metadata": {}, "outputs": [], "source": [] @@ -361,9 +246,9 @@ "hash": "4a749acc7d255f078aee52e0584cc77b3cb5aaed1b3a7407ec4262c1bf6cb526" }, "kernelspec": { - "display_name": "pyrad", + "display_name": "Python 3 (ipykernel)", "language": "python", - "name": "pyrad" + "name": "python3" }, "language_info": { "codemirror_mode": { diff --git a/pyproject.toml b/pyproject.toml index 33bc605bd..09cb138ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ line-length = 120 max-line-length = 120 generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"] good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$" -ignore = ["pyrad/cuda", "tests/cuda", "nerfw.py", "semantic_nerf.py"] +ignore = ["pyrad/cuda", "tests/cuda"] disable = [ "too-many-arguments", diff --git a/pyrad/cuda/mattport_cuda.cpp b/pyrad/cuda/pyrad_cuda.cpp similarity index 100% rename from pyrad/cuda/mattport_cuda.cpp rename to pyrad/cuda/pyrad_cuda.cpp diff --git a/pyrad/cuda/mattport_cuda_kernel.cu b/pyrad/cuda/pyrad_cuda_kernel.cu similarity index 100% rename from pyrad/cuda/mattport_cuda_kernel.cu rename to pyrad/cuda/pyrad_cuda_kernel.cu diff --git a/pyrad/nerf/colliders.py b/pyrad/nerf/colliders.py index 15cb3e737..818291e21 100644 --- a/pyrad/nerf/colliders.py +++ b/pyrad/nerf/colliders.py @@ -11,6 +11,10 @@ from pyrad.structures.rays import RayBundle class SceneBoundsCollider(nn.Module): """Module for setting near and far values for rays.""" + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + super().__init__() + def forward(self, ray_bundle: RayBundle) -> RayBundle: """To be implemented.""" raise NotImplementedError @@ -19,8 +23,8 @@ class SceneBoundsCollider(nn.Module): class AABBBoxCollider(SceneBoundsCollider): """Module for colliding rays with the scene bounds to compute near and far values.""" - def __init__(self, scene_bounds: SceneBounds) -> None: - super().__init__() + def __init__(self, scene_bounds: SceneBounds, **kwargs) -> None: + super().__init__(**kwargs) self.scene_bounds = scene_bounds @classmethod @@ -77,10 +81,10 @@ class AABBBoxCollider(SceneBoundsCollider): class NearFarCollider(SceneBoundsCollider): """Sets the nears and fars with fixed values.""" - def __init__(self, near_plane, far_plane) -> None: + def __init__(self, near_plane, far_plane, **kwargs) -> None: self.near_plane = near_plane self.far_plane = far_plane - super().__init__() + super().__init__(**kwargs) def forward(self, ray_bundle: RayBundle) -> RayBundle: ones = torch.ones_like(ray_bundle.origins[:, 0]) diff --git a/pyrad/nerf/dataset/collate.py b/pyrad/nerf/dataset/collate.py deleted file mode 100644 index d200dbeb4..000000000 --- a/pyrad/nerf/dataset/collate.py +++ /dev/null @@ -1,67 +0,0 @@ -""" -Classes to help with caching data while training. -""" - -import random -from typing import Callable - -import torch -from torch.utils.data import default_collate - - -def collate_batch_size_one(batch_list): - """Use the default collate function but then squeze to avoid the batch dimension being added.""" - assert len(batch_list) == 1 - collated_batch = default_collate(batch_list) - for key in collated_batch: - assert collated_batch[key].shape[0] == 1 - collated_batch[key] = collated_batch[key].squeeze(0) - return collated_batch - - -class CollateIterDataset(torch.utils.data.IterableDataset): - """Collated image dataset that implements caching of images. - - Args: - torch (_type_): _description_ - """ - - def __init__( - self, - dataset: torch.utils.data.Dataset, - collate_fn: Callable, - num_samples_to_collate: int, - num_times_to_repeat_images: int = 0, - ): - super().__init__() - self.dataset = dataset - self.collate_fn = collate_fn - self.num_samples_to_collate = num_samples_to_collate - self.num_times_to_repeat_images = num_times_to_repeat_images - - self.num_repeated = self.num_times_to_repeat_images # starting value - self.cached_batch_list = None - - def __getitem__(self, idx): - return self.dataset.__getitem__(idx) - - def get_batch_list(self): - """Returns a list of batches from the dataset attribute.""" - # sampling without replacement. TODO(ethan): don't use 'random' - indices = random.sample(range(len(self.dataset)), k=self.num_samples_to_collate) - batch_list = [self.dataset.__getitem__(idx) for idx in indices] - return batch_list - - def __iter__(self): - while True: - if self.num_repeated >= self.num_times_to_repeat_images: - # trigger a reset - self.num_repeated = 0 - batch_list = self.get_batch_list() - # possibly save a cached item - self.cached_batch_list = batch_list if self.num_times_to_repeat_images != 0 else None - else: - batch_list = self.cached_batch_list - self.num_repeated += 1 - collated_batch = self.collate_fn(batch_list) - yield collated_batch diff --git a/pyrad/nerf/dataset/format/friends.py b/pyrad/nerf/dataset/format/friends.py index 4d3140832..344947f5b 100644 --- a/pyrad/nerf/dataset/format/friends.py +++ b/pyrad/nerf/dataset/format/friends.py @@ -3,8 +3,8 @@ Code to handle loading friends datasets. """ import os -import numpy as np +import numpy as np import torch from pyrad.nerf.dataset.colmap_utils import read_cameras_binary, read_images_binary, read_pointsTD_binary @@ -46,6 +46,7 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant include_point_cloud (bool): whether or not to include the point cloud """ # pylint: disable=unused-argument + # pylint: disable=too-many-statements # --- image filenames ---- images_data = read_images_binary(os.path.join(basedir, "colmap", "images.bin")) @@ -62,7 +63,11 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant # -- set the bounding box --- aabb, transposed_point_cloud_transform = get_aabb_and_transform(basedir) - scene_bounds = SceneBounds(aabb=aabb) + scene_bounds_original = SceneBounds(aabb=aabb) + # for shifting and rescale accoding to scene bounds + box_center = scene_bounds_original.get_center() + box_scale_factor = 5.0 / scene_bounds_original.get_diagonal_length() # the target diagonal length + scene_bounds = scene_bounds_original.get_centered_and_scaled_scene_bounds(box_scale_factor) # --- intrinsics --- cameras_data = read_cameras_binary(os.path.join(basedir, "colmap", "cameras.bin")) @@ -90,8 +95,7 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant camera_to_world[..., 1:3] *= -1 camera_to_world = transposed_point_cloud_transform @ camera_to_world camera_to_world = camera_to_world[:, :3] - - # --- masks to mask out things (e.g., people) --- + camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor # center and rescale # --- semantics --- semantics = Semantics() @@ -106,11 +110,15 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant ] panoptic_classes = load_from_json(os.path.join(basedir, "panoptic_classes.json")) stuff_classes = panoptic_classes["stuff"] + stuff_colors = torch.tensor(panoptic_classes["stuff_colors"], dtype=torch.float32) / 255.0 thing_classes = panoptic_classes["thing"] + thing_colors = torch.tensor(panoptic_classes["thing_colors"], dtype=torch.float32) / 255.0 semantics = Semantics( stuff_classes=stuff_classes, + stuff_colors=stuff_colors, stuff_filenames=stuff_filenames, thing_classes=thing_classes, + thing_colors=thing_colors, thing_filenames=thing_filenames, ) @@ -123,6 +131,7 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float() xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1) xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3] + xyz = (xyz - box_center) * box_scale_factor # center and rescale point_cloud.xyz = xyz point_cloud.rgb = rgb diff --git a/pyrad/nerf/dataset/image_dataset.py b/pyrad/nerf/dataset/image_dataset.py index 9a28330cf..fcdbf56e9 100644 --- a/pyrad/nerf/dataset/image_dataset.py +++ b/pyrad/nerf/dataset/image_dataset.py @@ -2,65 +2,19 @@ Some dataset code. """ -from typing import List, Optional +from abc import abstractmethod +from typing import Dict, List, Optional, Union import numpy as np import numpy.typing as npt import PIL import torch from PIL import Image -from torch.utils.data import default_collate from torchtyping import TensorType from pyrad.nerf.dataset.structs import Semantics - -def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False): - """_summary_ - - Args: - batch (_type_): _description_ - """ - # TODO(ethan): handle sampling even when in different processes - # we don't want the same samples for all spawned processed when - # using distributed training - - batch = default_collate(batch_list) - num_images, image_height, image_width, _ = batch["image"].shape - - # # checking where mask is not zero - # # TODO(ethan): add this to the code with a config - # nonzero_indices = torch.nonzero(batch["mask"], as_tuple=False) - # chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) - # indices = nonzero_indices[chosen_indices] - - indices = torch.floor( - torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width]) - ).long() - c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)] - pixels = batch["image"][c, y, x] - mask = batch["mask"][c, y, x] - if "stuff_image" in batch: - stuff_image = batch["stuff_image"][c, y, x] - assert pixels.shape == (num_rays_per_batch, 3), pixels.shape - - # Needed to correct the random indices to their actual camera idx locations. - local_indices = indices.clone() - indices[:, 0] = batch["image_idx"][c] - collated_batch = { - "local_indices": local_indices, # local to the batch returned - "indices": indices, # with the abs camera indices - "pixels": pixels, - "mask": mask, - } - - if "stuff_image" in batch: - collated_batch["stuff_image"] = stuff_image - - if keep_full_image: - collated_batch["image"] = batch["image"] - - return collated_batch +from pyrad.utils.misc import is_not_none class ImageDataset(torch.utils.data.Dataset): @@ -71,7 +25,6 @@ class ImageDataset(torch.utils.data.Dataset): image_filenames: List[str], downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, - semantics: Semantics = Semantics(), **kwargs, ): """_summary_ @@ -86,13 +39,12 @@ class ImageDataset(torch.utils.data.Dataset): self.image_filenames = image_filenames self.downscale_factor = downscale_factor self.alpha_color = alpha_color - self.semantics = semantics self.kwargs = kwargs def __len__(self): return len(self.image_filenames) - def get_image(self, image_idx: int) -> npt.NDArray[np.uint8]: + def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]: """Returns the image. Args: @@ -122,47 +74,84 @@ class ImageDataset(torch.utils.data.Dataset): assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct." return image - def __getitem__(self, image_idx): - # the image might be RGB or RGBA, so we separate it - original_image = torch.from_numpy(self.get_image(image_idx).astype("float32") / 255.0) - + def get_image(self, image_idx: int): + """Returns a 3 channel image.""" + image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0) if self.alpha_color is not None: - image = original_image[:, :, :3] * original_image[:, :, -1:] + self.alpha_color * ( - 1.0 - original_image[:, :, -1:] - ) + assert image.shape[-1] == 4 + image = image[:, :, :3] * image[:, :, -1:] + self.alpha_color * (1.0 - image[:, :, -1:]) else: - image = original_image[:, :, :3] - num_channels = original_image.shape[2] - if num_channels == 4: - mask = original_image[:, :, 3] - elif num_channels == 3: - mask = np.ones_like(original_image[:, :, 0]) - else: - raise ValueError(f"Image shape of {image.shape} is in correct.") + image = image[:, :, :3] + return image - data = { - "image_idx": image_idx, - "image": image, # the pixels - "mask": mask, - } + @abstractmethod + def get_mask(self, image_idx: int) -> Union[TensorType["image_height", "image_width", 1], None]: + """Returns a mask, which indicates which pixels are valid to use with nerf.""" + return None - if self.semantics.stuff_filenames: - stuff_image_filename = self.semantics.stuff_filenames[image_idx] - pil_image = Image.open(stuff_image_filename) - if self.downscale_factor != 1.0: - image_width, image_height = pil_image.size - if image_width % self.downscale_factor != 0: - raise ValueError( - f"Image width {image_width} is not divisible by downscale_factor {self.downscale_factor}" - ) - if image_height % self.downscale_factor != 0: - raise ValueError( - f"Image height {image_height} is not divisible by downscale_factor {self.downscale_factor}" - ) - # NOTE(ethan): the use of NEAREST is important for semantic classes - pil_image = pil_image.resize( - (image_width // self.downscale_factor, image_height // self.downscale_factor), PIL.Image.NEAREST - ) - stuff_image = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] - data["stuff_image"] = stuff_image + @abstractmethod + def get_semantics(self, image_idx: int) -> Union[TensorType["image_height", "image_width", "num_classes"], None]: + """Returns an image with semantic class values.""" + return None + + def get_data(self, image_idx) -> Dict: + """Returns the ImageDataset data as a dictionary.""" + image = self.get_image(image_idx) + mask = self.get_mask(image_idx) + semantics = self.get_semantics(image_idx) + data = {"image_idx": image_idx} + assert is_not_none(image) + data["image"] = image + if is_not_none(mask): + assert mask.shape[:2] == image.shape[:2] + data["mask"] = mask + if is_not_none(semantics): + assert semantics.shape[:2] == image.shape[:2] + data["semantics"] = semantics return data + + def __getitem__(self, image_idx): + data = self.get_data(image_idx) + return data + + +class PanopticImageDataset(ImageDataset): + """Panoptic image dataset that masks out people.""" + + def __init__( + self, + semantics: Semantics, + image_filenames: List[str], + downscale_factor: int = 1, + alpha_color: Optional[TensorType[3]] = None, + **kwargs, + ): + self.semantics = semantics + self.person_index = self.semantics.thing_classes.index("person") + super().__init__(image_filenames, downscale_factor, alpha_color, **kwargs) + + def get_mask(self, image_idx): + """Mask out the people. Valid only where there aren't people.""" + thing_image_filename = self.semantics.thing_filenames[image_idx] + pil_image = Image.open(thing_image_filename) + if self.downscale_factor != 1.0: + image_width, image_height = pil_image.size + # the use of NEAREST is important for semantic classes + pil_image = pil_image.resize( + (image_width // self.downscale_factor, image_height // self.downscale_factor), Image.NEAREST + ) + thing_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + mask = (thing_semantics != self.person_index).to(torch.float32) # 1 where valid + return mask + + def get_semantics(self, image_idx): + stuff_image_filename = self.semantics.stuff_filenames[image_idx] + pil_image = Image.open(stuff_image_filename) + if self.downscale_factor != 1.0: + image_width, image_height = pil_image.size + # the use of NEAREST is important for semantic classes + pil_image = pil_image.resize( + (image_width // self.downscale_factor, image_height // self.downscale_factor), Image.NEAREST + ) + stuff_semantics = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + return stuff_semantics diff --git a/pyrad/nerf/dataset/structs.py b/pyrad/nerf/dataset/structs.py index a3f3d3427..74518c0ee 100644 --- a/pyrad/nerf/dataset/structs.py +++ b/pyrad/nerf/dataset/structs.py @@ -13,8 +13,8 @@ from torchtyping import TensorType class PointCloud: """_summary_""" - xyz = None - rgb = None + xyz: TensorType["num_points", 3] = None + rgb: TensorType["num_points", 3] = None @dataclass @@ -22,8 +22,10 @@ class Semantics: """_summary_""" stuff_classes: List[str] = None + stuff_colors: List[List[int]] = None stuff_filenames: List[str] = None thing_classes: List[str] = None + thing_colors: List[List[int]] = None thing_filenames: List[str] = None @@ -37,6 +39,22 @@ class SceneBounds: aabb: TensorType[2, 3] = None + def get_diagonal_length(self): + """Returns the longest diagonal length.""" + diff = self.aabb[1] - self.aabb[0] + length = torch.sqrt((diff**2).sum() + 1e-20) + return length + + def get_center(self): + """Reterns the center of the box.""" + diff = self.aabb[1] - self.aabb[0] + return self.aabb[0] + diff / 2.0 + + def get_centered_and_scaled_scene_bounds(self, scale_factor=1.0): + """Returns a new box that has been shifted and rescaled to be centered + about the origin.""" + return SceneBounds(aabb=(self.aabb - self.get_center()) * scale_factor) + @dataclass class DatasetInputs: diff --git a/pyrad/nerf/dataset/utils.py b/pyrad/nerf/dataset/utils.py index aa8ad92cb..c6f00bc28 100644 --- a/pyrad/nerf/dataset/utils.py +++ b/pyrad/nerf/dataset/utils.py @@ -1,7 +1,8 @@ """ -For loading the blender dataset format. +Code for loading the dataset inputs and for caching the data for fast loading. """ +import logging from typing import Optional, Union from omegaconf import ListConfig @@ -11,7 +12,30 @@ from pyrad.nerf.dataset.format.friends import load_friends_data from pyrad.nerf.dataset.format.instant_ngp import load_instant_ngp_data from pyrad.nerf.dataset.structs import DatasetInputs from pyrad.structures.colors import get_color -from pyrad.utils.io import get_absolute_path +from pyrad.utils.io import get_absolute_path, load_from_pkl, make_dir, write_to_pkl +from pyrad.utils.misc import get_hash_str_from_dict + + +def get_cache_filename_from_kwargs(kwargs: dict, split: str): + """Creates a cache filename from the dataset inputs arguments.""" + dataset_config_hash = get_hash_str_from_dict(kwargs) + dataset_config_hash_filename = make_dir( + get_absolute_path(f"cache/dataset_inputs/{dataset_config_hash}-{split}.pkl") + ) + return dataset_config_hash_filename + + +def save_dataset_inputs_kwargs_to_cache(kwargs: dict, split: str): + """Saves the dataset inputs to cache.""" + dataset_inputs = get_dataset_inputs(**kwargs, split=split) + dataset_config_hash_filename = get_cache_filename_from_kwargs(kwargs, split) + write_to_pkl(dataset_config_hash_filename, dataset_inputs) + + +def get_dataset_inputs_from_cache(kwargs: dict, split: str): + """Loads the dataset inputs from cache.""" + dataset_config_hash_filename = get_cache_filename_from_kwargs(kwargs, split) + return load_from_pkl(dataset_config_hash_filename) def get_dataset_inputs( @@ -20,10 +44,8 @@ def get_dataset_inputs( split: str, downscale_factor: int = 1, alpha_color: Optional[Union[str, list, ListConfig]] = None, - load_dataset_inputs_from_cache: bool = False, ) -> DatasetInputs: """Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator. - # TODO: implement the `test` split, which will have depths and normals, etc. Args: data_directory (str): Location of data @@ -35,10 +57,6 @@ def get_dataset_inputs( DatasetInputs: The inputs needed for generating rays. """ - if load_dataset_inputs_from_cache: - print("Loading from cache.") - return None - if alpha_color is not None: alpha_color = get_color(alpha_color) @@ -68,3 +86,12 @@ def get_dataset_inputs( raise NotImplementedError(f"{dataset_format} is not a valid dataset type") return dataset_inputs + + +def get_dataset_inputs_from_dataset_config(*, split: str, use_cache: bool = False, **kwargs): + """This may optionally use the cache to return the dataset inputs.""" + if use_cache: + logging.info("Loading from cache! Be careful with using this when making changes!") + print(kwargs) + return get_dataset_inputs_from_cache(kwargs, split) + return get_dataset_inputs(**kwargs, split=split) diff --git a/pyrad/nerf/field_modules/field_heads.py b/pyrad/nerf/field_modules/field_heads.py index 99b18bff7..43cee6f91 100644 --- a/pyrad/nerf/field_modules/field_heads.py +++ b/pyrad/nerf/field_modules/field_heads.py @@ -19,8 +19,7 @@ class FieldHeadNames(Enum): UNCERTAINTY = "uncertainty" TRANSIENT_RGB = "transient_rgb" TRANSIENT_DENSITY = "transient_density" - SEMANTICS_STUFF = "semantics_stuff" - SEMANTICS_THING = "semantics_thing" + SEMANTICS = "semantics" class FieldHead(FieldModule): @@ -113,14 +112,14 @@ class UncertaintyFieldHead(FieldHead): super().__init__(in_dim=in_dim, out_dim=1, field_head_name=FieldHeadNames.UNCERTAINTY, activation=activation) -class TransientRGBHead(FieldHead): +class TransientRGBFieldHead(FieldHead): """Transient RGB output""" def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.TRANSIENT_RGB, activation=activation) -class TransientDensityHead(FieldHead): +class TransientDensityFieldHead(FieldHead): """Transient density output""" def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: @@ -129,7 +128,7 @@ class TransientDensityHead(FieldHead): ) -class SemanticStuffHead(FieldHead): +class SemanticFieldHead(FieldHead): """Semantic stuff output""" def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None: @@ -138,6 +137,4 @@ class SemanticStuffHead(FieldHead): num_classes (int): Number of semantic classes in_dim (int): Input dimension """ - super().__init__( - in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS_STUFF, activation=None - ) + super().__init__(in_dim=in_dim, out_dim=num_classes, field_head_name=FieldHeadNames.SEMANTICS, activation=None) diff --git a/pyrad/nerf/fields/nerf_field.py b/pyrad/nerf/fields/nerf_field.py index 9a4b7589a..3327c9abb 100644 --- a/pyrad/nerf/fields/nerf_field.py +++ b/pyrad/nerf/fields/nerf_field.py @@ -1,4 +1,4 @@ -""" Classic NeRF field""" +"""Classic NeRF field""" from typing import Dict, Optional, Tuple diff --git a/pyrad/nerf/fields/nerfw_field.py b/pyrad/nerf/fields/nerfw_field.py new file mode 100644 index 000000000..c9ca451af --- /dev/null +++ b/pyrad/nerf/fields/nerfw_field.py @@ -0,0 +1,108 @@ +"""Fields for nerf-w""" + +from typing import Dict, Optional, Tuple + +import torch +from torch import nn +from torchtyping import TensorType + +from pyrad.nerf.field_modules.embedding import Embedding +from pyrad.nerf.field_modules.encoding import Encoding, Identity +from pyrad.nerf.field_modules.field_heads import ( + DensityFieldHead, + FieldHeadNames, + RGBFieldHead, + TransientDensityFieldHead, + TransientRGBFieldHead, + UncertaintyFieldHead, +) +from pyrad.nerf.field_modules.mlp import MLP +from pyrad.nerf.fields.base import Field +from pyrad.structures.rays import PointSamples + + +class VanillaNerfWField(Field): + """The NeRF-W field which has appearance and transient conditioning.""" + + def __init__( + self, + num_images: int, + position_encoding: Encoding = Identity(in_dim=3), + direction_encoding: Encoding = Identity(in_dim=3), + base_mlp_num_layers: int = 8, + base_mlp_layer_width: int = 256, + head_mlp_num_layers: int = 2, + head_mlp_layer_width: int = 128, + appearance_embedding_dim: int = 48, + transient_embedding_dim: int = 16, + skip_connections: Tuple[int] = (4,), + ) -> None: + super().__init__() + self.num_images = num_images + self.position_encoding = position_encoding + self.direction_encoding = direction_encoding + self.base_mlp_num_layers = base_mlp_num_layers + self.base_mlp_layer_width = base_mlp_layer_width + self.head_mlp_num_layers = head_mlp_num_layers + self.head_mlp_layer_width = head_mlp_layer_width + self.appearance_embedding_dim = appearance_embedding_dim + self.transient_embedding_dim = transient_embedding_dim + + self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim) + self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim) + + self.mlp_base = MLP( + in_dim=self.position_encoding.get_out_dim(), + num_layers=base_mlp_num_layers, + layer_width=base_mlp_layer_width, + skip_connections=skip_connections, + ) + self.mlp_transient = MLP( + in_dim=self.mlp_base.get_out_dim() + self.embedding_transient.get_out_dim(), + out_dim=base_mlp_layer_width // 2, + num_layers=4, + layer_width=base_mlp_layer_width, + activation=nn.ReLU(), + ) + self.mlp_head = MLP( + in_dim=self.mlp_base.get_out_dim() + + self.direction_encoding.get_out_dim() + + self.embedding_appearance.get_out_dim(), + num_layers=head_mlp_num_layers, + layer_width=head_mlp_layer_width, + ) + + self.field_head_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) + self.field_head_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim()) + + self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.get_out_dim()) + self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.get_out_dim()) + self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.get_out_dim()) + + def get_density(self, point_samples: PointSamples): + """Computes and returns the densities.""" + encoded_xyz = self.position_encoding(point_samples.positions) + base_mlp_out = self.mlp_base(encoded_xyz) + density = self.field_head_density(base_mlp_out) + return density, base_mlp_out + + def get_outputs( + self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + ) -> Dict[FieldHeadNames, TensorType]: + outputs = {} + encoded_dir = self.direction_encoding(point_samples.directions) + embedded_appearance = self.embedding_appearance(point_samples.camera_indices) + mlp_head_out = self.mlp_head(torch.cat([density_embedding, encoded_dir, embedded_appearance], dim=-1)) + outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_head_out) # static rgb + embedded_transient = self.embedding_transient(point_samples.camera_indices) + transient_mlp_out = self.mlp_transient(torch.cat([density_embedding, embedded_transient], dim=-1)) + outputs[self.field_head_transient_uncertainty.field_head_name] = self.field_head_transient_uncertainty( + transient_mlp_out + ) # uncertainty + outputs[self.field_head_transient_rgb.field_head_name] = self.field_head_transient_rgb( + transient_mlp_out + ) # transient rgb + outputs[self.field_head_transient_density.field_head_name] = self.field_head_transient_density( + transient_mlp_out + ) # transient density + return outputs diff --git a/pyrad/nerf/graph/base.py b/pyrad/nerf/graph/base.py index 8e05bdd8e..f18f9b700 100644 --- a/pyrad/nerf/graph/base.py +++ b/pyrad/nerf/graph/base.py @@ -13,11 +13,9 @@ from torchtyping import TensorType from pyrad.nerf.dataset.structs import SceneBounds from pyrad.nerf.ray_generator import RayGenerator -from pyrad.nerf.colliders import AABBBoxCollider from pyrad.structures.cameras import get_camera_model from pyrad.structures.rays import RayBundle -from pyrad.utils.misc import get_masked_dict -from pyrad.utils.misc import is_not_none +from pyrad.utils.misc import get_masked_dict, instantiate_from_dict_config, is_not_none class AbstractGraph(nn.Module): @@ -52,6 +50,7 @@ class Graph(AbstractGraph): loss_coefficients: DictConfig = None, steps_per_occupancy_grid_update=16, scene_bounds: SceneBounds = None, + collider_config: DictConfig = None, **kwargs, ) -> None: super().__init__() @@ -59,6 +58,7 @@ class Graph(AbstractGraph): self.intrinsics = intrinsics self.camera_to_world = camera_to_world self.scene_bounds = scene_bounds + self.collider_config = collider_config self.loss_coefficients = loss_coefficients self.steps_per_occupancy_grid_update = steps_per_occupancy_grid_update self.kwargs = kwargs @@ -70,7 +70,7 @@ class Graph(AbstractGraph): def populate_collider(self): """Set the scene bounds collider to use.""" - self.collider = AABBBoxCollider(self.scene_bounds) + self.collider = instantiate_from_dict_config(self.collider_config, scene_bounds=self.scene_bounds) @abstractmethod def populate_misc_modules(self): @@ -110,39 +110,36 @@ class Graph(AbstractGraph): masked_batch = get_masked_dict(batch, valid_mask) # NOTE(ethan): this is really slow if on CPU! outputs = self.get_outputs(masked_intersected_ray_bundle) loss_dict = self.get_loss_dict(outputs=outputs, batch=masked_batch) - return outputs, loss_dict + aggregated_loss_dict = self.get_aggregated_loss_dict(loss_dict) + return outputs, aggregated_loss_dict @abstractmethod def get_loss_dict(self, outputs, batch) -> Dict[str, torch.tensor]: """Computes and returns the losses.""" - def get_aggregated_loss_from_loss_dict(self, loss_dict): + def get_aggregated_loss_dict(self, loss_dict): """Computes the aggregated loss from the loss_dict and the coefficients specified.""" - aggregated_loss = 0.0 + aggregated_loss_dict = {} for loss_name, loss_value in loss_dict.items(): assert loss_name in self.loss_coefficients, f"{loss_name} no in self.loss_coefficients" loss_coefficient = self.loss_coefficients[loss_name] - aggregated_loss += loss_coefficient * loss_value - return aggregated_loss + aggregated_loss_dict[loss_name] = loss_coefficient * loss_value + aggregated_loss_dict["aggregated_loss"] = sum(loss_dict.values()) + return aggregated_loss_dict @torch.no_grad() def get_outputs_for_camera(self, intrinsics, camera_to_world, chunk_size=1024, training_camera_index=0): """Takes in camera parameters and computes the output of the graph.""" - # NOTE(ethan): this function has a spike in CPU usage - - device = self.get_device() assert len(intrinsics.shape) == 1 num_intrinsics_params = len(intrinsics) camera_class = get_camera_model(num_intrinsics_params) camera = camera_class(*intrinsics.tolist(), camera_to_world=camera_to_world) - camera_ray_bundle = camera.generate_camera_rays() + camera_ray_bundle = camera.generate_camera_rays(device=self.get_device()) # TODO(ethan): decide how to properly handle the image indices for validation images camera_ray_bundle.set_camera_indices(camera_index=training_camera_index) - camera_ray_bundle.move_to_device(device) image_height, image_width = camera_ray_bundle.origins.shape[:2] num_rays = len(camera_ray_bundle) - outputs = {} outputs_lists = defaultdict(list) for i in range(0, num_rays, chunk_size): @@ -158,5 +155,5 @@ class Graph(AbstractGraph): return outputs @abstractmethod - def log_test_image_outputs(self, image_idx, step, image, mask, outputs): + def log_test_image_outputs(self, image_idx, step, batch, outputs): """Writes the test image outputs.""" diff --git a/pyrad/nerf/graph/instant_ngp.py b/pyrad/nerf/graph/instant_ngp.py index 5e2a4ee65..84a6d4fb6 100644 --- a/pyrad/nerf/graph/instant_ngp.py +++ b/pyrad/nerf/graph/instant_ngp.py @@ -16,7 +16,7 @@ from pyrad.nerf.graph.base import Graph from pyrad.nerf.loss import MSELoss from pyrad.nerf.occupancy_grid import OccupancyGrid from pyrad.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer -from pyrad.nerf.sampler import PDFSampler, UniformSampler +from pyrad.nerf.ray_sampler import PDFSampler, UniformSampler from pyrad.structures import colors from pyrad.structures.rays import RayBundle from pyrad.utils import visualization, writer @@ -115,10 +115,10 @@ class NGPGraph(Graph): pixels = batch["pixels"].to(device) rgb_loss = self.rgb_loss(pixels, outputs["rgb"]) loss_dict = {"rgb_loss": rgb_loss} - loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, mask, outputs): + def log_test_image_outputs(self, image_idx, step, batch, outputs): + image = batch["image"] rgb = outputs["rgb"] acc = visualization.apply_colormap(outputs["accumulation"]) depth = visualization.apply_depth_colormap( diff --git a/pyrad/nerf/graph/nerfw.py b/pyrad/nerf/graph/nerfw.py index 43acd7554..13be193b8 100644 --- a/pyrad/nerf/graph/nerfw.py +++ b/pyrad/nerf/graph/nerfw.py @@ -1,163 +1,142 @@ """ NeRF-W (NeRF in the wild) implementation. -TODO: """ -from typing import Tuple - import torch -from torch import nn +from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity -from pyrad.nerf.field_modules.embedding import Embedding -from pyrad.nerf.field_modules.field_heads import ( - FieldHeadNames, - TransientDensityHead, - TransientRGBHead, - UncertaintyFieldHead, -) -from pyrad.nerf.field_modules.mlp import MLP -from pyrad.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph -from pyrad.nerf.renderers import UncertaintyRenderer -from pyrad.structures.rays import RayBundle, RaySamples +from pyrad.nerf.field_modules.encoding import NeRFEncoding +from pyrad.nerf.field_modules.field_heads import FieldHeadNames +from pyrad.nerf.fields.nerf_field import NeRFField +from pyrad.nerf.fields.nerfw_field import VanillaNerfWField +from pyrad.nerf.graph.base import Graph +from pyrad.nerf.loss import MSELoss +from pyrad.nerf.ray_sampler import PDFSampler, UniformSampler +from pyrad.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer, UncertaintyRenderer +from pyrad.structures import colors +from pyrad.structures.rays import RayBundle from pyrad.utils import visualization, writer -class NerfWField(Field): - """The NeRF-W field which has appearance and transient conditioning.""" +class NerfWGraph(Graph): + """NeRF-W graph""" def __init__( self, - position_encoding: Encoding = Identity(in_dim=3), - direction_encoding: Encoding = Identity(in_dim=3), - base_mlp_num_layers: int = 8, - base_mlp_layer_width: int = 256, - head_mlp_num_layers: int = 2, - head_mlp_layer_width: int = 128, - skip_connections: Tuple[int] = (4,), - field_heads: Tuple[FieldHead] = (RGBFieldHead(),), + intrinsics=None, + camera_to_world=None, + near_plane=2.0, + far_plane=6.0, + num_coarse_samples=64, + num_importance_samples=128, + **kwargs, ) -> None: - super().__init__() - self.num_images = num_images - self.appearance_embedding_dim = 48 - self.transient_embedding_dim = 16 - - self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim) - self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim) - self.mlp_transient = MLP( - in_dim=self.mlp_base.get_out_dim() + self.embedding_transient.get_out_dim(), - out_dim=self.layer_width // 2, - num_layers=4, - layer_width=self.layer_width, - activation=nn.ReLU(), - ) - - self.mlp_head = MLP( - in_dim=self.mlp_base.get_out_dim() - + self.encoding_dir.get_out_dim() - + self.embedding_appearance.get_out_dim(), - out_dim=self.layer_width // 2, - num_layers=1, - activation=nn.ReLU(), - ) - - def build_heads(self): - super().build_heads() - self.field_output_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.get_out_dim()) - self.field_output_transient_rgb = TransientRGBHead(in_dim=self.mlp_transient.get_out_dim()) - self.field_output_transient_density = TransientDensityHead(in_dim=self.mlp_transient.get_out_dim()) - - def forward(self, ray_samples: RaySamples): - positions = ray_samples.positions - directions = ray_samples.directions - camera_indices = ray_samples.get_camera_indices() - encoded_xyz = self.encoding_xyz(positions) - encoded_dir = self.encoding_dir(directions) - embedded_appearance = self.embedding_appearance(camera_indices) - embedded_transient = self.embedding_transient(camera_indices) - - base_mlp_out = self.mlp_base(encoded_xyz) - rgb_mlp_out = self.mlp_rgb(torch.cat([base_mlp_out, encoded_dir, embedded_appearance], dim=-1)) - transient_mlp_out = self.mlp_transient(torch.cat([base_mlp_out, embedded_transient], dim=-1)) - - field_rgb_output = self.field_output_rgb(rgb_mlp_out) - field_density_out = self.field_output_density(base_mlp_out) - field_transient_uncertainty_out = self.field_output_transient_uncertainty(transient_mlp_out) - field_transient_rgb_out = self.field_output_transient_rgb(transient_mlp_out) - field_transient_density_out = self.field_output_transient_density(transient_mlp_out) - - field_outputs = {} - field_outputs.update(field_rgb_output) - field_outputs.update(field_density_out) - field_outputs.update(field_transient_uncertainty_out) - field_outputs.update(field_transient_rgb_out) - field_outputs.update(field_transient_density_out) - return field_outputs - - -class NerfWGraph(NeRFGraph): - """NeRF-W graph""" - - def __init__(self, intrinsics=None, camera_to_world=None, **kwargs) -> None: + self.near_plane = near_plane + self.far_plane = far_plane + self.num_coarse_samples = num_coarse_samples + self.num_importance_samples = num_importance_samples + self.field_coarse = None + self.field_fine = None self.num_images = len(intrinsics) self.appearance_embedding_dim = 48 self.transient_embedding_dim = 16 super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) - self.renderer_uncertainty = UncertaintyRenderer() def populate_fields(self): """Set the fields.""" - self.field_coarse = NeRFField() - self.field_fine = NerfWField( + + position_encoding = NeRFEncoding( + in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ) + direction_encoding = NeRFEncoding( + in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True + ) + + self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) + self.field_fine = VanillaNerfWField( num_images=self.num_images, + position_encoding=position_encoding, + direction_encoding=direction_encoding, appearance_embedding_dim=self.appearance_embedding_dim, transient_embedding_dim=self.transient_embedding_dim, ) + def populate_misc_modules(self): + # samplers + self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples) + self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=colors.BLACK) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + self.renderer_uncertainty = UncertaintyRenderer() + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = StructuralSimilarityIndexMeasure() + self.lpips = LearnedPerceptualImagePatchSimilarity() + + def get_param_groups(self): + param_groups = {} + param_groups["fields"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters()) + return param_groups + def get_outputs(self, ray_bundle: RayBundle): - # coarse network - uniform_ray_samples = self.sampler_uniform(ray_bundle) - coarse_field_outputs = self.field_coarse(uniform_ray_samples) - coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) + # uniform sampling + ray_samples_uniform = self.sampler_uniform(ray_bundle) + + # coarse field + field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples()) + weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) rgb_coarse = self.renderer_rgb( - rgb=coarse_field_outputs[FieldHeadNames.RGB], - weights=coarse_weights, - ).rgb - depth_coarse = self.renderer_depth(coarse_weights, uniform_ray_samples.ts).depth - - # fine network - pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) - fine_field_outputs = self.field_fine(pdf_ray_samples) - fine_weights = pdf_ray_samples.get_weights( - fine_field_outputs[FieldHeadNames.DENSITY] + fine_field_outputs[FieldHeadNames.TRANSIENT_DENSITY] + rgb=field_outputs_coarse[FieldHeadNames.RGB], + weights=weights_coarse, ) - fine_weights_static = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY]) - fine_weights_transient = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.TRANSIENT_DENSITY]) + depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.ts) + # pdf sampling + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) + + # fine field + field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) + + # fine weights + weights_fine = ray_samples_pdf.get_weights( + field_outputs_fine[FieldHeadNames.DENSITY] + field_outputs_fine[FieldHeadNames.TRANSIENT_DENSITY] + ) + weights_fine_static = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) + weights_fine_transient = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.TRANSIENT_DENSITY]) + + # rgb rgb_fine_static_component = self.renderer_rgb( - rgb=fine_field_outputs[FieldHeadNames.RGB], - weights=fine_weights, - ).rgb + rgb=field_outputs_fine[FieldHeadNames.RGB], + weights=weights_fine, + ) rgb_fine_transient_component = self.renderer_rgb( - rgb=fine_field_outputs[FieldHeadNames.TRANSIENT_RGB], - weights=fine_weights, - ).rgb + rgb=field_outputs_fine[FieldHeadNames.TRANSIENT_RGB], + weights=weights_fine, + ) rgb_fine = rgb_fine_static_component + rgb_fine_transient_component - rgb_fine_static = self.renderer_rgb( - rgb=fine_field_outputs[FieldHeadNames.RGB], - weights=fine_weights_static, - ).rgb + rgb=field_outputs_fine[FieldHeadNames.RGB], + weights=weights_fine_static, + ) - # fine_renderer_accumulation_static = self.renderer_accumulation(fine_weights_static) - depth_fine = self.renderer_depth(fine_weights, pdf_ray_samples.ts).depth - depth_fine_static = self.renderer_depth(fine_weights_static, pdf_ray_samples.ts).depth - uncertainty = self.renderer_uncertainty( - fine_field_outputs[FieldHeadNames.UNCERTAINTY], fine_weights_transient - ).uncertainty + # density + density_transient = field_outputs_fine[FieldHeadNames.TRANSIENT_DENSITY] - density_transient = fine_field_outputs[FieldHeadNames.TRANSIENT_DENSITY] + # depth + depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.ts) + depth_fine_static = self.renderer_depth(weights_fine_static, ray_samples_pdf.ts) + + # uncertainty + uncertainty = self.renderer_uncertainty(field_outputs_fine[FieldHeadNames.UNCERTAINTY], weights_fine_transient) - # outputs: outputs = { "rgb_coarse": rgb_coarse, # (num_rays, 3) "rgb_fine": rgb_fine, @@ -188,10 +167,10 @@ class NerfWGraph(NeRFGraph): "uncertainty_loss": uncertainty_loss, "density_loss": density_loss, } - loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, mask, outputs): + def log_test_image_outputs(self, image_idx, step, batch, outputs): + image = batch["image"] rgb_coarse = outputs["rgb_coarse"] rgb_fine = outputs["rgb_fine"] rgb_fine_static = outputs["rgb_fine_static"] @@ -210,4 +189,7 @@ class NerfWGraph(NeRFGraph): row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2) combined_image = torch.cat([row0, row1, row2], dim=-3) - writer.put_image(name="img/image_idx_{image_idx}-nerfw", image=combined_image, step=step) + writer.put_image(name=f"img/image_idx_{image_idx}-nerfw", image=combined_image, step=step) + + mask = batch["mask"].repeat(1, 1, 3) + writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step) diff --git a/pyrad/nerf/graph/semantic_nerf.py b/pyrad/nerf/graph/semantic_nerf.py index ead8e0b4a..73a2eb611 100644 --- a/pyrad/nerf/graph/semantic_nerf.py +++ b/pyrad/nerf/graph/semantic_nerf.py @@ -2,157 +2,172 @@ Semantic NeRF implementation. """ -from typing import Tuple +from typing import Dict, Optional, Tuple import torch from torch import nn +from torchtyping import TensorType -from pyrad.nerf.field_modules.field_heads import FieldHeadNames, SemanticStuffHead +from pyrad.nerf.field_modules.encoding import Encoding, Identity, NeRFEncoding +from pyrad.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadNames, RGBFieldHead, SemanticFieldHead from pyrad.nerf.field_modules.mlp import MLP -from pyrad.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph +from pyrad.nerf.fields.base import Field +from pyrad.nerf.fields.nerf_field import NeRFField +from pyrad.nerf.graph.vanilla_nerf import NeRFGraph from pyrad.nerf.renderers import SemanticRenderer -from pyrad.structures.rays import RayBundle, RaySamples -from pyrad.utils import visualization, writer +from pyrad.structures.rays import PointSamples, RayBundle +from pyrad.utils import writer -class SemanticNerfField(NeRFField): +class SemanticNerfField(Field): """Semantic-NeRF field""" def __init__( - self, num_layers=8, layer_width=256, skip_connections: Tuple = (4,), num_stuff_classes: int = None + self, + num_semantic_classes: int, + position_encoding: Encoding = Identity(in_dim=3), + direction_encoding: Encoding = Identity(in_dim=3), + base_mlp_num_layers: int = 8, + base_mlp_layer_width: int = 256, + head_mlp_num_layers: int = 2, + head_mlp_layer_width: int = 128, + skip_connections: Tuple[int] = (4,), ) -> None: - assert num_stuff_classes is not None - self.num_stuff_classes = num_stuff_classes - super().__init__(num_layers=num_layers, layer_width=layer_width, skip_connections=skip_connections) - self.semantic_mlp = MLP( - in_dim=self.mlp_base.get_out_dim(), - out_dim=self.layer_width // 2, + super().__init__() + self.num_semantic_classes = num_semantic_classes + self.position_encoding = position_encoding + self.direction_encoding = direction_encoding + self.mlp_base = MLP( + in_dim=self.position_encoding.get_out_dim(), + num_layers=base_mlp_num_layers, + layer_width=base_mlp_layer_width, + skip_connections=skip_connections, + ) + self.mlp_head = MLP( + in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), + num_layers=head_mlp_num_layers, + layer_width=head_mlp_layer_width, + ) + self.mlp_semantic = MLP( + in_dim=self.mlp_head.get_out_dim(), + layer_width=self.mlp_head.layer_width // 2, num_layers=1, activation=nn.ReLU(), ) - self.field_output_semantic = SemanticStuffHead( - in_dim=self.semantic_mlp.get_out_dim(), num_classes=num_stuff_classes + self.field_head_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) + self.field_head_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim()) + self.field_head_semantic = SemanticFieldHead( + in_dim=self.mlp_semantic.get_out_dim(), num_classes=self.num_semantic_classes ) - self.semantic_sequential = torch.nn.Sequential(self.semantic_mlp, self.field_output_semantic) - def forward(self, ray_samples: RaySamples): - """Evaluates the field at points along the ray.""" - positions = ray_samples.positions - directions = ray_samples.directions - encoded_xyz = self.encoding_xyz(positions) - encoded_dir = self.encoding_dir(directions) + def get_density(self, point_samples: PointSamples): + encoded_xyz = self.position_encoding(point_samples.positions) base_mlp_out = self.mlp_base(encoded_xyz) - rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1)) + density = self.field_head_density(base_mlp_out) + return density, base_mlp_out - field_rgb_output = self.field_output_rgb(rgb_mlp_out) - field_density_out = self.field_output_density(base_mlp_out) - field_semantic_output = self.semantic_sequential(base_mlp_out) - - field_outputs = {} - field_outputs.update(field_rgb_output) - field_outputs.update(field_density_out) - field_outputs.update(field_semantic_output) - return field_outputs + def get_outputs( + self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None + ) -> Dict[FieldHeadNames, TensorType]: + encoded_dir = self.direction_encoding(point_samples.directions) + mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) + outputs = {} + # rgb + outputs[self.field_head_rgb.field_head_name] = self.field_head_rgb(mlp_out) + # semantic + mlp_out_sem = self.mlp_semantic(mlp_out) + outputs[self.field_head_semantic.field_head_name] = self.field_head_semantic(mlp_out_sem) + return outputs class SemanticNerfGraph(NeRFGraph): """Semantic-NeRF graph""" - def __init__(self, intrinsics=None, camera_to_world=None, stuff_classes: int = None, **kwargs) -> None: + def __init__(self, intrinsics=None, camera_to_world=None, stuff_classes=None, stuff_colors=None, **kwargs) -> None: assert stuff_classes is not None self.stuff_classes = stuff_classes + self.stuff_colors = stuff_colors super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) - self.renderer_semantic = SemanticRenderer() - self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="mean") def populate_fields(self): """Set the fields.""" - self.field_coarse = NeRFField() - self.field_fine = SemanticNerfField(num_stuff_classes=len(self.stuff_classes)) + position_encoding = NeRFEncoding( + in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ) + direction_encoding = NeRFEncoding( + in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True + ) + self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding) + + num_semantic_classes = len(self.stuff_classes) + self.field_fine = SemanticNerfField( + num_semantic_classes, position_encoding=position_encoding, direction_encoding=direction_encoding + ) + + def populate_misc_modules(self): + super().populate_misc_modules() + self.renderer_semantic = SemanticRenderer() + self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="mean") def get_outputs(self, ray_bundle: RayBundle): - # coarse network: - uniform_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples + # uniform sampling + ray_samples_uniform = self.sampler_uniform(ray_bundle) - coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs - - coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) - - coarse_renderer_outputs = self.renderer_rgb( - rgb=coarse_field_outputs[FieldHeadNames.RGB], - weights=coarse_weights, - ) # RendererOutputs - coarse_renderer_accumulation = self.renderer_accumulation(coarse_weights) # RendererOutputs - coarse_renderer_depth = self.renderer_depth(coarse_weights, uniform_ray_samples.ts) - - # fine network: - pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples - fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs - - fine_weights = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY]) - - fine_renderer_outputs = self.renderer_rgb( - rgb=fine_field_outputs[FieldHeadNames.RGB], - weights=fine_weights, - ) # RendererOutputs - fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs - fine_renderer_depth = self.renderer_depth(fine_weights, pdf_ray_samples.ts) - - # TODO refactor this into "vis" section. Doesn't need to be run during training. - coarse_renderer_depth = visualization.apply_depth_colormap( - coarse_renderer_depth.depth, - accumulation=coarse_renderer_accumulation.accumulation, - near_plane=self.near_plane, - far_plane=self.far_plane, + # coarse field + field_outputs_coarse = self.field_coarse.forward(ray_samples_uniform.to_point_samples()) + weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) + rgb_coarse = self.renderer_rgb( + rgb=field_outputs_coarse[FieldHeadNames.RGB], + weights=weights_coarse, ) - fine_renderer_depth = visualization.apply_depth_colormap( - fine_renderer_depth.depth, - accumulation=fine_renderer_accumulation.accumulation, - near_plane=self.near_plane, - far_plane=self.far_plane, + accumulation_coarse = self.renderer_accumulation(weights_coarse) + depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.ts) + + # pdf sampling + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) + + # fine field + field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) + weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) + rgb_fine = self.renderer_rgb( + rgb=field_outputs_fine[FieldHeadNames.RGB], + weights=weights_fine, ) - coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation) - fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation) + accumulation_fine = self.renderer_accumulation(weights_fine) + depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.ts) + semantic_fine = self.renderer_semantic(field_outputs_fine[FieldHeadNames.SEMANTICS], weights=weights_fine) - semantic_logits = self.renderer_semantic( - fine_field_outputs[FieldHeadNames.SEMANTICS_STUFF], weights=fine_weights - ).semantics - - # outputs: outputs = { - "rgb_coarse": coarse_renderer_outputs.rgb, - "rgb_fine": fine_renderer_outputs.rgb, - "accumulation_coarse": coarse_renderer_accumulation, - "accumulation_fine": fine_renderer_accumulation, - "depth_coarse": coarse_renderer_depth, - "depth_fine": fine_renderer_depth, - "semantic_logits": semantic_logits, + "rgb_coarse": rgb_coarse, + "rgb_fine": rgb_fine, + "accumulation_coarse": accumulation_coarse, + "accumulation_fine": accumulation_fine, + "depth_coarse": depth_coarse, + "depth_fine": depth_fine, + "semantic_fine": semantic_fine, } return outputs def get_loss_dict(self, outputs, batch): - # TODO(ethan): batch has "stuff_image". use it - device = outputs["rgb_coarse"].device - pixels = batch["pixels"].to(device) + pixels = batch["pixels"] rgb_loss_coarse = self.rgb_loss(pixels, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(pixels, outputs["rgb_fine"]) - - semantic_logits = outputs["semantic_logits"] - semantic_classes = batch["stuff_image"][..., 0].to(device).long() + semantic_logits = outputs["semantic_fine"] + semantic_classes = batch["semantics"][..., 0].long() semantic_loss_fine = self.cross_entropy_loss(semantic_logits, semantic_classes) loss_dict = { "rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine, "semantic_loss_fine": semantic_loss_fine, } - - loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, mask, outputs): - super().log_test_image_outputs(image_idx, step, image, mask, outputs) - semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantic_logits"], dim=-1), dim=-1) / len( - self.stuff_classes - ) - semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None]) - writer.put_image(name="semantics/image_idx_{image_idx}", image=semantic_labels_image, step=step) + def log_test_image_outputs(self, image_idx, step, batch, outputs): + super().log_test_image_outputs(image_idx, step, batch, outputs) + semantic_logits = outputs["semantic_fine"] + semantic_labels = torch.argmax(torch.nn.functional.softmax(semantic_logits, dim=-1), dim=-1) + semantic_labels_image = self.stuff_colors[semantic_labels] + writer.put_image(name=f"semantics/image_idx_{image_idx}", image=semantic_labels_image, step=step) + + mask = batch["mask"].repeat(1, 1, 3) + writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step) diff --git a/pyrad/nerf/graph/vanilla_nerf.py b/pyrad/nerf/graph/vanilla_nerf.py index c8421f4b9..79a04141b 100644 --- a/pyrad/nerf/graph/vanilla_nerf.py +++ b/pyrad/nerf/graph/vanilla_nerf.py @@ -9,17 +9,16 @@ import torch from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity + from pyrad.nerf.field_modules.encoding import NeRFEncoding from pyrad.nerf.field_modules.field_heads import FieldHeadNames - from pyrad.nerf.fields.nerf_field import NeRFField from pyrad.nerf.graph.base import Graph from pyrad.nerf.loss import MSELoss +from pyrad.nerf.ray_sampler import PDFSampler, UniformSampler from pyrad.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer -from pyrad.nerf.sampler import PDFSampler, UniformSampler from pyrad.structures import colors from pyrad.structures.rays import RayBundle -from pyrad.nerf.colliders import NearFarCollider from pyrad.utils import visualization, writer @@ -44,9 +43,6 @@ class NeRFGraph(Graph): self.field_fine = None super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) - def populate_collider(self): - self.collider = NearFarCollider(self.near_plane, self.far_plane) - def populate_fields(self): """Set the fields.""" @@ -126,10 +122,10 @@ class NeRFGraph(Graph): rgb_loss_coarse = self.rgb_loss(pixels, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(pixels, outputs["rgb_fine"]) loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} - loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, mask, outputs): + def log_test_image_outputs(self, image_idx, step, batch, outputs): + image = batch["image"] rgb_coarse = outputs["rgb_coarse"] rgb_fine = outputs["rgb_fine"] acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"]) @@ -150,12 +146,10 @@ class NeRFGraph(Graph): combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) combined_acc = torch.cat([acc_coarse, acc_fine], dim=1) combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) - mask = visualization.apply_depth_colormap(mask[..., None]) writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step) writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step) writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step) - writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] diff --git a/pyrad/nerf/image_sampler.py b/pyrad/nerf/image_sampler.py new file mode 100644 index 000000000..18eadf3e4 --- /dev/null +++ b/pyrad/nerf/image_sampler.py @@ -0,0 +1,75 @@ +""" +Code for sampling images. +""" + +import random + +import torch +from torch.utils.data import default_collate + +from pyrad.utils.misc import get_dict_to_torch + + +class ImageSampler(torch.utils.data.IterableDataset): + """Samples 'image_batch's.""" + + def __getitem__(self, idx): + raise NotImplementedError() + + def __iter__(self): + raise NotImplementedError() + + +class CacheImageSampler(ImageSampler): + """Collated image dataset that implements caching of images. + Creates batches of the ImageDataset return type. + """ + + def __init__( + self, + dataset: torch.utils.data.Dataset, + num_samples_to_collate: int, + num_times_to_repeat_images: int = 0, + device="cpu", + ): + super().__init__() + self.dataset = dataset + self.num_samples_to_collate = num_samples_to_collate + self.num_times_to_repeat_images = num_times_to_repeat_images + self.device = device + + self.num_repeated = self.num_times_to_repeat_images # starting value + self.cached_collated_batch = None + self.first_time = True + + def __getitem__(self, idx): + return self.dataset.__getitem__(idx) + + def get_batch_list(self): + """Returns a list of batches from the dataset attribute.""" + indices = random.sample(range(len(self.dataset)), k=self.num_samples_to_collate) + batch_list = [self.dataset.__getitem__(idx) for idx in indices] + return batch_list + + def get_collated_batch(self): + """Returns a collated batch.""" + batch_list = self.get_batch_list() + collated_batch = default_collate(batch_list) + collated_batch = get_dict_to_torch(collated_batch, device=self.device) + return collated_batch + + def __iter__(self): + while True: + if self.first_time or ( + self.num_times_to_repeat_images != -1 and self.num_repeated >= self.num_times_to_repeat_images + ): + # trigger a reset + self.num_repeated = 0 + collated_batch = self.get_collated_batch() + # possibly save a cached item + self.cached_collated_batch = collated_batch if self.num_times_to_repeat_images != 0 else None + self.first_time = False + else: + collated_batch = self.cached_collated_batch + self.num_repeated += 1 + yield collated_batch diff --git a/pyrad/nerf/pixel_sampler.py b/pyrad/nerf/pixel_sampler.py new file mode 100644 index 000000000..fc053c43c --- /dev/null +++ b/pyrad/nerf/pixel_sampler.py @@ -0,0 +1,68 @@ +""" +Code for sampling pixels. +""" + +import random + +import torch + + +def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool = False): + """ + Operates on a batch of images and samples pixels to use for generating rays. + Returns a collated batch which is input to the Graph. + It will sample only within the valid 'mask' if it's specified. + """ + + num_images, image_height, image_width, _ = batch["image"].shape + + # only sample within the mask, if the mask is in the batch + if "mask" in batch: + nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) + chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) + indices = nonzero_indices[chosen_indices] + else: + indices = torch.floor( + torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width]) + ).long() + + c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)] + pixels = batch["image"][c, y, x] + if "mask" in batch: + mask = batch["mask"][c, y, x] + if "semantics" in batch: + semantics = batch["semantics"][c, y, x] + assert pixels.shape == (num_rays_per_batch, 3), pixels.shape + + # Needed to correct the random indices to their actual camera idx locations. + local_indices = indices.clone() + indices[:, 0] = batch["image_idx"][c] + collated_batch = { + "local_indices": local_indices, # local to the batch returned + "indices": indices, # with the abs camera indices + "pixels": pixels, + } + if "mask" in batch: + collated_batch["mask"] = mask + if "semantics" in batch: + collated_batch["semantics"] = semantics + + if keep_full_image: + collated_batch["image"] = batch["image"] + + return collated_batch + + +class PixelSampler: # pylint: disable=too-few-public-methods + """Samples 'pixel_batch's from 'image_batch's.""" + + def __init__(self, num_rays_per_batch, keep_full_image=False) -> None: + self.num_rays_per_batch = num_rays_per_batch + self.keep_full_image = keep_full_image + + def sample(self, image_batch): + """Sample an image batch and return a pixel batch.""" + pixel_batch = collate_image_dataset_batch( + image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image + ) + return pixel_batch diff --git a/pyrad/nerf/sampler.py b/pyrad/nerf/ray_sampler.py similarity index 100% rename from pyrad/nerf/sampler.py rename to pyrad/nerf/ray_sampler.py diff --git a/pyrad/nerf/trainer.py b/pyrad/nerf/trainer.py index 1c9bc1ebc..56507b5bd 100644 --- a/pyrad/nerf/trainer.py +++ b/pyrad/nerf/trainer.py @@ -4,7 +4,6 @@ Code to train model. import copy import logging import os -from pydoc import locate from time import time from typing import Callable, Dict, List @@ -12,16 +11,15 @@ import torch import torch.distributed as dist from omegaconf import DictConfig from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data import DataLoader -from pyrad.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one -from pyrad.nerf.dataset.image_dataset import ImageDataset, collate_batch -from pyrad.nerf.dataset.utils import DatasetInputs, get_dataset_inputs +from pyrad.nerf.dataset.utils import DatasetInputs, get_dataset_inputs_from_dataset_config +from pyrad.nerf.image_sampler import CacheImageSampler from pyrad.nerf.optimizers import Optimizers +from pyrad.nerf.pixel_sampler import PixelSampler from pyrad.utils import profiler, writer from pyrad.utils.callbacks import update_occupancy from pyrad.utils.decorators import check_main_thread -from pyrad.utils.misc import get_dict_to_torch +from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config logging.getLogger("PIL").setLevel(logging.WARNING) @@ -43,8 +41,8 @@ class Trainer: self.world_size = world_size # dataset variables self.train_image_dataset = None - self.train_dataset = None - self.train_dataloader = None + self.train_image_sampler = None + self.train_pixel_sampler = None self.val_image_dataset = None self.val_image_intrinsics = None self.val_image_camera_to_world = None @@ -59,29 +57,25 @@ class Trainer: def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions.""" - dataset_inputs_train = get_dataset_inputs(**self.config.data.dataset, split="train") + dataset_inputs_train = get_dataset_inputs_from_dataset_config(**self.config.data.dataset, split="train") if test_mode: - dataset_inputs_eval = get_dataset_inputs(**self.config.data.dataset, split="test") + dataset_inputs_eval = get_dataset_inputs_from_dataset_config(**self.config.data.dataset, split="test") else: config_data_dataset_val = copy.deepcopy(self.config.data.dataset) config_data_dataset_val.downscale_factor = self.config.data.val_downscale_factor - dataset_inputs_eval = get_dataset_inputs(**config_data_dataset_val, split="val") + dataset_inputs_eval = get_dataset_inputs_from_dataset_config(**config_data_dataset_val, split="val") self.setup_dataset_train(dataset_inputs_train) self.setup_dataset_eval(dataset_inputs_eval) self.setup_graph(dataset_inputs_train) - def collate_fn(self, batch_list): - """TODO(ethan): I need to replace this. - I'm only using this for multiprocess pickle issues for now. - """ - return collate_batch(batch_list, self.config.data.dataloader.num_rays_per_batch, keep_full_image=False) - @profiler.time_function def setup_dataset_eval(self, dataset_inputs: DatasetInputs): """Helper method to load test or val dataset based on test/train mode""" - self.val_image_dataset = ImageDataset( + self.val_image_dataset = instantiate_from_dict_config( + self.config.data.image_dataset, image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor, + semantics=dataset_inputs.semantics, alpha_color=dataset_inputs.alpha_color, ) self.val_image_intrinsics = dataset_inputs.intrinsics @@ -90,25 +84,24 @@ class Trainer: @profiler.time_function def setup_dataset_train(self, dataset_inputs: DatasetInputs): """_summary_""" - self.train_image_dataset = ImageDataset( + self.train_image_dataset = instantiate_from_dict_config( + self.config.data.image_dataset, image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor, semantics=dataset_inputs.semantics, alpha_color=dataset_inputs.alpha_color, - ) - self.train_dataset = CollateIterDataset( + ) # ImageDataset + self.train_image_sampler = CacheImageSampler( self.train_image_dataset, - collate_fn=self.collate_fn, - num_samples_to_collate=self.config.data.dataloader.num_images_to_sample_from, - num_times_to_repeat_images=self.config.data.dataloader.num_times_to_repeat_images, - ) - self.train_dataloader = DataLoader( - self.train_dataset, - batch_size=1, - num_workers=self.config.data.dataloader.num_workers, - collate_fn=collate_batch_size_one, - pin_memory=True, - ) + num_samples_to_collate=len(self.train_image_dataset) + if self.config.data.image_sampler.num_images_to_sample_from == 0 + else self.config.data.image_sampler.num_images_to_sample_from, + num_times_to_repeat_images=self.config.data.image_sampler.num_times_to_repeat_images, + device=self.device if self.config.data.image_sampler.move_to_graph_device else "cpu", + ) # ImageSampler + self.train_pixel_sampler = PixelSampler( + num_rays_per_batch=self.config.data.pixel_sampler.num_rays_per_batch, keep_full_image=False + ) # PixelSampler @profiler.time_function def setup_graph(self, dataset_inputs: DatasetInputs): @@ -117,15 +110,13 @@ class Trainer: Args: dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters. """ - # hydra instantiate fails because of scene_bounds, so we are doing the following instead - kwargs = {k: v for k, v in self.config.graph.network.items() if k != "_target_"} - graph_class = locate(self.config.graph.network._target_) # pylint: disable=protected-access - self.graph = graph_class( - **kwargs, + self.graph = instantiate_from_dict_config( + self.config.graph.network, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world, scene_bounds=dataset_inputs.scene_bounds, stuff_classes=dataset_inputs.semantics.stuff_classes, + stuff_colors=dataset_inputs.semantics.stuff_colors, ) self.graph.to(self.device) @@ -197,11 +188,12 @@ class Trainer: """_summary_""" train_start = time() num_iterations = self.config.graph.max_num_iterations - iter_dataset = iter(self.train_dataloader) + iter_train_image_sampler = iter(self.train_image_sampler) for step in range(self.start_step, self.start_step + num_iterations): data_start = time() - batch = next(iter_dataset) - batch = get_dict_to_torch(batch, device=self.device) + image_batch = next(iter_train_image_sampler) + image_batch = get_dict_to_torch(image_batch, device=self.device) + batch = self.train_pixel_sampler.sample(image_batch) writer.put_time( name=writer.EventName.ITER_LOAD_TIME, start_time=data_start, @@ -209,7 +201,6 @@ class Trainer: step=step, avg_over_iters=True, ) - iter_start = time() loss_dict = self.train_iteration(batch, step, _callback=[update_occupancy]) writer.put_time( @@ -231,10 +222,9 @@ class Trainer: if step != 0 and step % self.config.logging.steps_per_log == 0: writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step) - # TODO: add the learning rates to tensorboard/logging if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0: self.save_checkpoint(self.config.graph.model_dir, step) - if step % self.config.graph.steps_per_test == 0: # NOTE(ethan): we should still run this in dry-run mode! + if step % self.config.graph.steps_per_test == 0: for image_idx in self.config.data.val_image_indices: _ = self.test_image(image_idx=image_idx, step=step) self._write_out_storage(step) @@ -273,16 +263,15 @@ class Trainer: def test_image(self, image_idx, step): """Test a specific image.""" self.graph.eval() - intrinsics = self.val_image_intrinsics[image_idx] - camera_to_world = self.val_image_camera_to_world[image_idx] + intrinsics = self.val_image_intrinsics[image_idx].to(self.device) + camera_to_world = self.val_image_camera_to_world[image_idx].to(self.device) chunk_size = self.config.data.val_num_rays_per_chunk training_camera_index = image_idx # TODO(ethan): change this because training and test should not be the same outputs = self.graph.get_outputs_for_camera( intrinsics, camera_to_world, chunk_size=chunk_size, training_camera_index=training_camera_index ) - val_image_data = self.val_image_dataset[image_idx] - image = val_image_data["image"].to(self.device) - mask = val_image_data["mask"].to(self.device) - psnr = self.graph.log_test_image_outputs(image_idx, step, image, mask, outputs) + batch = self.val_image_dataset[image_idx] + batch = get_dict_to_torch(batch, device=self.device) + psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs) self.graph.train() return psnr diff --git a/pyrad/structures/cameras.py b/pyrad/structures/cameras.py index 24ad6ebde..548b6a6a4 100644 --- a/pyrad/structures/cameras.py +++ b/pyrad/structures/cameras.py @@ -98,17 +98,21 @@ class Camera: """ return - def generate_camera_rays(self) -> RayBundle: + def generate_camera_rays(self, device=None) -> RayBundle: """Generate rays for the camera. Returns: Rays: Camera rays of shape [height, width] """ + if device is None: + device = self.camera_to_world.device height = self.get_image_height() width = self.get_image_width() - intrinsics = self.get_intrinsics().unsqueeze(0).repeat(height, width, 1) # (num_rays, num_intrinsics_params) - camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1) # (num_rays, 3, 4) - coords = self.get_image_coords() + intrinsics = ( + self.get_intrinsics().unsqueeze(0).repeat(height, width, 1).to(device) + ) # (num_rays, num_intrinsics_params) + camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1).to(device) # (num_rays, 3, 4) + coords = self.get_image_coords().to(device) ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords) return ray_bundle diff --git a/pyrad/structures/rays.py b/pyrad/structures/rays.py index f2458c527..2052ad2b9 100644 --- a/pyrad/structures/rays.py +++ b/pyrad/structures/rays.py @@ -74,7 +74,12 @@ class RaySamples: def to_point_samples(self) -> PointSamples: """Convert to PointSamples instance and return.""" # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) + return PointSamples( + positions=self.positions, + directions=self.directions, + camera_indices=self.camera_indices, + valid_mask=self.valid_mask, + ) def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: """Return weights based on predicted densities @@ -139,7 +144,7 @@ class RayBundle: """ self.origins = self.origins.to(device) self.directions = self.directions.to(device) - if not isinstance(self.camera_indices, type(None)): + if is_not_none(self.camera_indices): self.camera_indices = self.camera_indices.to(device) def set_camera_indices(self, camera_index: int) -> None: @@ -227,7 +232,7 @@ class RayBundle: dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples] deltas = dists * torch.norm(self.directions[..., None, :], dim=-1) - if self.camera_indices is not None: + if is_not_none(self.camera_indices): camera_indices = self.camera_indices.unsqueeze(1).repeat(1, positions.shape[1]) else: camera_indices = None diff --git a/radiance/utils/README.md b/pyrad/utils/README.md similarity index 100% rename from radiance/utils/README.md rename to pyrad/utils/README.md diff --git a/pyrad/utils/io.py b/pyrad/utils/io.py index 1afd26e48..ce735c9b7 100644 --- a/pyrad/utils/io.py +++ b/pyrad/utils/io.py @@ -2,37 +2,44 @@ Input/output utils. """ import json -import os - import logging +import os +import pickle def load_from_json(filename: str): - """_summary_ - - Args: - filename (str): _description_ - - Returns: - _type_: _description_ - """ + """Load a dictionary from a JSON filename.""" assert filename.endswith(".json") with open(filename, "r", encoding="UTF-8") as file: return json.load(file) def write_to_json(filename: str, content: dict): - """_summary_ + """Write data to a JSON file. Args: - filename (str): _description_ - content (dict): _description_ + filename (str): The filename to write to. + content (dict): The dictionary data to write. """ assert filename.endswith(".json") with open(filename, "w", encoding="UTF-8") as file: json.dump(content, file) +def load_from_pkl(filename: str): + """Load from a pickle file.""" + assert filename.endswith(".pkl") + with open(filename, "rb") as f: + return pickle.load(f) + + +def write_to_pkl(filename: str, content): + "Write to a pickle file." + assert filename.endswith(".pkl") + with open(filename, "wb") as f: + pickle.dump(content, f) + + def get_git_root(path, dirs=(".git",), default=None): """_summary_ "https://stackoverflow.com/questions/22081209/find-the-root-of-the-git-repository-where-the-file-lives diff --git a/pyrad/utils/misc.py b/pyrad/utils/misc.py index e1f90aa5d..2c0868e1a 100644 --- a/pyrad/utils/misc.py +++ b/pyrad/utils/misc.py @@ -2,7 +2,13 @@ Miscellaneous helper code. """ +import hashlib +import json +from pydoc import locate +from typing import Any, Dict + import torch +from omegaconf import DictConfig class DotDict(dict): @@ -51,3 +57,22 @@ def get_masked_dict(d, mask): for key, value in d.items(): masked_dict[key] = value[mask] return masked_dict + + +def instantiate_from_dict_config(dict_config: DictConfig, **kwargs): + """Our version of hydra's instantiate function.""" + dict_config_kwargs = {k: v for k, v in dict_config.items() if k != "_target_"} + uninstantiated_class = locate(dict_config._target_) # pylint: disable=protected-access + all_kwargs = dict_config_kwargs + all_kwargs.update(kwargs) + instantiated_class = uninstantiated_class(**all_kwargs) + return instantiated_class + + +def get_hash_str_from_dict(dictionary: Dict[str, Any]) -> str: + """MD5 hash of a dictionary. Based on + https://www.doc.ic.ac.uk/~nuric/coding/how-to-hash-a-dictionary-in-python.html""" + dhash = hashlib.md5() + encoded = json.dumps(dictionary, sort_keys=True).encode() + dhash.update(encoded) + return dhash.hexdigest() diff --git a/pyrad/utils/visualization.py b/pyrad/utils/visualization.py index aceb8c024..738644d7d 100644 --- a/pyrad/utils/visualization.py +++ b/pyrad/utils/visualization.py @@ -52,6 +52,7 @@ def apply_depth_colormap( depth = (depth - near_plane) / (far_plane - near_plane + 1e-10) depth = torch.clip(depth, 0, 1) + # depth = torch.nan_to_num(depth, nan=0.0) # TODO(ethan): remove this colored_image = apply_colormap(depth, cmap=cmap) diff --git a/scripts/run_data_preprocessor.py b/scripts/run_data_preprocessor.py new file mode 100644 index 000000000..fe20976f7 --- /dev/null +++ b/scripts/run_data_preprocessor.py @@ -0,0 +1,34 @@ +""" +Script to process different dataset types into our format. +""" + +# TODO(ethan): we should write a script here to convert from other datasets to our formats +# e.g., Blender, LLFF, ShapeNet, DTU +from hydra import compose, initialize +import json + +from pyrad.nerf.dataset.utils import get_dataset_inputs +from pyrad.utils.io import get_absolute_path, make_dir, write_to_pkl +from pyrad.utils.misc import get_hash_str_from_dict + + +def save_dataset_inputs_to_cache( + dataset_name="friends_TBBT-big_living_room", splits=["train", "val", "test"], downscale_factors=[1, 2, 4] +): + """Save the datasets input to cache. + TODO(ethan): use a serializable representation that's cleaner instead of pickle.""" + with initialize(version_base=None, config_path="../configs"): + config = compose(config_name="default_setup.yaml", overrides=[f"data/dataset={dataset_name}"]) + for downscale_factor in downscale_factors: + for split in splits: + config.data.dataset.downscale_factor = downscale_factor + dataset_config_hash = get_hash_str_from_dict(dict(config.data.dataset)) + dataset_config_hash_filename = make_dir( + get_absolute_path(f"cache/dataset_inputs/{dataset_config_hash}-{split}.pkl") + ) + dataset_inputs = get_dataset_inputs(**config.data.dataset, split=split) + write_to_pkl(dataset_config_hash_filename, dataset_inputs) + + +if __name__ == "__main__": + save_dataset_inputs_to_cache() diff --git a/scripts/run_dataset_processor.py b/scripts/run_dataset_processor.py deleted file mode 100644 index aa510b4e7..000000000 --- a/scripts/run_dataset_processor.py +++ /dev/null @@ -1,6 +0,0 @@ -""" -Script to process different dataset types into our format. -""" - -# TODO(ethan): we should write a script here to convert from other datasets to our formats -# e.g., Blender, LLFF, ShapeNet, DTU diff --git a/scripts/run_train.py b/scripts/run_train.py index fd6e70d43..b51b78f81 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -1,13 +1,14 @@ """ run_train_nerf.py """ + import datetime import logging import os import random import socket -from datetime import timedelta import traceback +from datetime import timedelta from typing import Any, Callable import hydra @@ -20,7 +21,6 @@ import yaml from pyrad.nerf.trainer import Trainer from pyrad.utils import comms, profiler -from time import time logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) diff --git a/tests/nerf/__init__.py b/tests/nerf/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/nerf/dataloader/configs/blender_lego.yaml b/tests/nerf/dataloader/configs/blender_lego.yaml deleted file mode 100644 index 5fb9a3164..000000000 --- a/tests/nerf/dataloader/configs/blender_lego.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# dataset options -dataset: - data_directory: data/lego - dataset_format: blender - downscale_factor: 1.0 - -# dataloader options -dataloader: - num_workers: 1 - shuffle: true - num_images_to_sample_from: 8 - num_rays_per_batch: 1024 - -# ray generation module and its option -network: - ray_generator: - class_name: RayGenerator - inputs: ["cxy"] # camera, x, and y indices - meta_data: - in_dim: 0 - learning_rate: 0.0 diff --git a/tests/nerf/dataloader/test_ray_generator.py b/tests/nerf/dataloader/test_ray_generator.py deleted file mode 100644 index dad77d383..000000000 --- a/tests/nerf/dataloader/test_ray_generator.py +++ /dev/null @@ -1,88 +0,0 @@ -""" -Code to test ray generation. -""" - -import imageio -import torch -from torch.utils.data import DataLoader -from tqdm import tqdm - -from pyrad.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one -from pyrad.nerf.dataset.image_dataset import ImageDataset, collate_batch -from pyrad.nerf.dataset.utils import get_dataset_inputs -from pyrad.nerf.ray_generator import RayGenerator -from pyrad.utils.io import get_absolute_path -from pyrad.viewer.plotly import visualize_dataset - - -def visualize_batch(batch): - """Returns an image from a batch.""" - assert "image" in batch, "You must call collate_batch with keep_full_image=True for this to work." - # set the color of the sampled rays - c, y, x = [i.flatten() for i in torch.split(batch["local_indices"], 1, dim=-1)] - batch["image"][c, y, x] = 1.0 - - # batch["image"] is num_images, h, w, 3 - images = torch.split(batch["image"], 1, dim=0) - image_list = [image[0] for image in images] - image = torch.cat(image_list, dim=1) # cat along the width dimension - image = (image * 255.0).to(torch.uint8) - return image - - -def test_dataloader(visualize=False): - """Testing for the dataloader from input dataset parameters to rays.""" - - data_directory = "tests/data/lego_test" - dataset_format = "blender" - downscale_factor = 1 - num_images_to_sample_from = 1 - num_times_to_repeat_images = 40 - num_rays_per_batch = 1024 - num_workers = 0 - - data_directory = get_absolute_path(data_directory) - dataset_inputs = get_dataset_inputs( - data_directory=data_directory, dataset_format=dataset_format, downscale_factor=downscale_factor, split="train" - ) - image_dataset = ImageDataset( - image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor - ) - iter_dataset = CollateIterDataset( - image_dataset, - collate_fn=lambda batch_list: collate_batch(batch_list, num_rays_per_batch, keep_full_image=True), - num_samples_to_collate=num_images_to_sample_from, - num_times_to_repeat_images=num_times_to_repeat_images, - ) - dataloader = DataLoader( - iter_dataset, - batch_size=1, - num_workers=num_workers, - collate_fn=collate_batch_size_one, - pin_memory=True, - ) - dataloader_iter = iter(dataloader) - - ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world) - - num_batches = 10 - for _ in tqdm(range(num_batches)): - batch = next(dataloader_iter) - ray_bundle = ray_generator.forward(batch["indices"]) - - # visualize the batch - image = visualize_batch(batch) - if visualize: - imageio.imwrite("temp0.png", image) - - # visualize the RayBundle - small_ray_bundle = ray_bundle.sample(100) - fig = visualize_dataset(camera_origins=dataset_inputs.camera_to_world[:, :3, 3], ray_bundle=small_ray_bundle) - if visualize: - fig.write_image("temp1.png") - fig.write_html("temp1.html") - assert True - - -if __name__ == "__main__": - test_dataloader(visualize=True) diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py index 819246d8f..564d85845 100644 --- a/tests/nerf/test_sampler.py +++ b/tests/nerf/test_sampler.py @@ -2,7 +2,7 @@ Test samplers """ import torch -from pyrad.nerf.sampler import PDFSampler, UniformSampler +from pyrad.nerf.ray_sampler import PDFSampler, UniformSampler from pyrad.structures.rays import RayBundle from pyrad.nerf.colliders import NearFarCollider diff --git a/tests/test_run_train.py b/tests/test_run_train.py index b8a82d7a7..652fdfc38 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -13,3 +13,6 @@ def test_run_train(): initialize(config_path="../configs/tests/") conf = compose("test_dryrun.yaml") main(conf) + +if __name__ == "__main__": + test_run_train()