From 5099a4f7c5475baac5ba2c232db1e93c73520eb1 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 22 Apr 2022 18:19:39 +0000 Subject: [PATCH] resizing image --- README.md | 5 +- configs/default.yaml | 19 +- environment/requirements.txt | 1 + mattport/nerf/dataset/image_dataset.py | 20 +- mattport/nerf/dataset/utils.py | 2 +- mattport/nerf/field_modules/mlp.py | 46 +++-- mattport/nerf/graph/base.py | 2 +- mattport/nerf/graph/vanilla_nerf.py | 45 ++++- mattport/nerf/trainer.py | 2 +- mattport/utils/writer.py | 11 +- notebooks/visualize_dataset.ipynb | 266 +++++++++++++++++++++++++ 11 files changed, 374 insertions(+), 45 deletions(-) create mode 100644 notebooks/visualize_dataset.ipynb diff --git a/README.md b/README.md index f8c1f7179..62e03284c 100644 --- a/README.md +++ b/README.md @@ -30,8 +30,9 @@ Download the original NeRF dataset and put it in the following format. ``` - data/ - - fern/ - - lego/ + - blender/ + - fern/ + - lego/ ``` # Training a model diff --git a/configs/default.yaml b/configs/default.yaml index ed779b867..71633700d 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -1,4 +1,9 @@ -hydra: +defaults: + # This makes the logging module logs colorful + - override hydra/job_logging: colorlog + - override hydra/hydra_logging: colorlog + +hydra: run: # Configure output dir of each experiment programmatically from the arguments # Example "outputs/mnist/classifier/baseline/2021-03-10-141516" @@ -21,9 +26,9 @@ model_dir: /tmp/mattport_models/ # dataset options dataset: - data_directory: data/lego + data_directory: data/blender/lego dataset_type: blender - downscale_factor: 1.0 + downscale_factor: 4 # dataloader options dataloader: @@ -40,8 +45,8 @@ param_groups: lr: 1e-4 scheduler: type: StepLR - step_size: 30 - gamma: 0.1 + step_size: 30 # Period of learning rate decay. + gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. cameras: optimizer: type: SGD @@ -61,3 +66,7 @@ machine_config: graph: # NOTE(ethan): currently setting to nothing b/c we're still designing the config format... _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph + near_plane: 1.0 + far_plane: 6.0 + num_coarse_samples: 64 + num_importance_samples: 64 diff --git a/environment/requirements.txt b/environment/requirements.txt index 9d0aa8b6e..e5050a742 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,4 +1,5 @@ black +hydra-colorlog hydra-core imageio==2.16.1 ipywidgets>=7.6 diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index 7acafa66a..da8fd233f 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -6,6 +6,7 @@ from typing import List import numpy as np import torch +import PIL from PIL import Image from torch.utils.data import default_collate @@ -15,8 +16,9 @@ from mattport.utils.misc import DotDict class ImageDataset(torch.utils.data.Dataset): """Dataset that returns images.""" - def __init__(self, image_filenames: List[str], downscale_factor: float = 1.0): + def __init__(self, image_filenames: List[str], downscale_factor: int = 1): super().__init__() + assert isinstance(downscale_factor, int) self.image_filenames = image_filenames self.downscale_factor = downscale_factor @@ -33,7 +35,21 @@ class ImageDataset(torch.utils.data.Dataset): np.uint8: an image of shape (H, W, 3 or 4) """ image_filename = self.image_filenames[image_idx] - image = np.array(Image.open(image_filename)) # shape is (h, w, 3 or 4) + pil_image = Image.open(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}" + ) + pil_image = pil_image.resize( + (image_width // self.downscale_factor, image_height // self.downscale_factor), PIL.Image.BILINEAR + ) + image = np.array(pil_image) # shape is (h, w, 3 or 4) assert len(image.shape) == 3 assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct." return image diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index afeff467c..872f67fad 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -62,7 +62,7 @@ def load_blender_data(basedir, downscale_factor=1.0, split="train"): dataset_inputs = DatasetInputs( image_filenames=image_filenames, downscale_factor=downscale_factor, - intrinsics=intrinsics, + intrinsics=intrinsics * 1.0 / downscale_factor, # downscaling the intrinsics here camera_to_world=camera_to_world, ) diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py index ac6092300..7fcce4c7f 100644 --- a/mattport/nerf/field_modules/mlp.py +++ b/mattport/nerf/field_modules/mlp.py @@ -1,10 +1,13 @@ """ Multi Layer Perceptron """ -from typing import Optional +from tkinter import X +from typing import List, Optional +import torch from torch import nn from torchtyping import TensorType + from mattport.nerf.field_modules.base import FieldModule @@ -16,7 +19,8 @@ class MLP(FieldModule): in_dim: int, out_dim: int, num_layers: int, - layer_width: int, + layer_width: int = None, + skip_connections: List[int] = [], activation: Optional[nn.Module] = nn.ReLU(), out_activation: Optional[nn.Module] = None, ) -> None: @@ -36,26 +40,28 @@ class MLP(FieldModule): self.out_dim = out_dim self.num_layers = num_layers self.layer_width = layer_width + self.skip_connections = skip_connections self.activation = activation self.out_activation = out_activation self.net = None self.build_nn_modules() def build_nn_modules(self) -> None: - """Initialize mulilayer perceptron.""" + """Initialize multi-layer perceptron.""" layers = [] - for i in range(self.num_layers): - if i == 0: - layers.append(nn.Linear(self.in_dim, self.layer_width)) - else: - layers.append(nn.Linear(self.layer_width, self.layer_width)) - if self.activation is not None: - layers.append(self.activation) - layers.append(nn.Linear(self.layer_width, self.out_dim)) - if self.out_activation is not None: - layers.append(self.out_activation) - - self.net = nn.Sequential(*layers) + if self.num_layers == 1: + layers.append(nn.Linear(self.in_dim, self.out_dim)) + else: + for i in range(self.num_layers - 1): + if i == 0: + assert i not in self.skip_connections, "Skip connection at layer 0 doesn't make sense." + layers.append(nn.Linear(self.in_dim, self.layer_width)) + elif i in self.skip_connections: + layers.append(nn.Linear(self.layer_width + self.in_dim, self.layer_width)) + else: + layers.append(nn.Linear(self.layer_width, self.layer_width)) + layers.append(nn.Linear(self.layer_width, self.out_dim)) + self.layers = nn.ModuleList(layers) def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]: """Process input with a multilayer perceptron. @@ -66,5 +72,11 @@ class MLP(FieldModule): Returns: TensorType[..., "out_dim"]: Network output """ - - return self.net(in_tensor) + x = in_tensor + for i, layer in enumerate(self.layers): + if i in self.skip_connections: + x = torch.cat([in_tensor, x], -1) + x = layer(x) + if self.activation: + x = self.activation(x) + return x diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index 2bce70e44..041b3e8b1 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -34,7 +34,7 @@ class Node: class Graph(nn.Module): """_summary_""" - def __init__(self, intrinsics=None, camera_to_world=None) -> None: + def __init__(self, intrinsics=None, camera_to_world=None, **kwargs) -> None: super().__init__() self.intrinsics = intrinsics self.camera_to_world = camera_to_world diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index c58f6bb2b..7696eb237 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -24,18 +24,26 @@ from mattport.structures.rays import RaySamples class NeRFField(nn.Module): """NeRF module""" - def __init__(self) -> None: + def __init__(self, num_layers=8, layer_width=256, skip_connections=[4]) -> None: super().__init__() + self.num_layers = num_layers + self.layer_width = layer_width + self.skip_connections = skip_connections + self.encoding_xyz = NeRFEncoding(in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0) - self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=6, min_freq_exp=0.0, max_freq_exp=4.0) + self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0) self.mlp_base = MLP( - in_dim=self.encoding_xyz.get_out_dim(), out_dim=64, num_layers=8, layer_width=64, activation=nn.ReLU() + in_dim=self.encoding_xyz.get_out_dim(), + out_dim=self.layer_width, + num_layers=self.num_layers, + layer_width=self.layer_width, + skip_connections=self.skip_connections, + activation=nn.ReLU(), ) self.mlp_rgb = MLP( in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(), - out_dim=64, - num_layers=2, - layer_width=64, + out_dim=self.layer_width // 2, + num_layers=1, activation=nn.ReLU(), ) self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim()) @@ -67,16 +75,31 @@ class NeRFField(nn.Module): class NeRFGraph(Graph): """Vanilla NeRF graph""" - def __init__(self, intrinsics=None, camera_to_world=None) -> None: - super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world) + def __init__( + self, + intrinsics=None, + camera_to_world=None, + near_plane=1.0, + far_plane=6.0, + num_coarse_samples=64, + num_importance_samples=128, + **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 + super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) def populate_modules(self): # ray generator self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world) # samplers - self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=6.0, num_samples=64) - self.sampler_pdf = PDFSampler(num_samples=64) + self.sampler_uniform = UniformSampler( + near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_coarse_samples + ) + self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples) # field self.field_coarse = NeRFField() @@ -95,7 +118,7 @@ class NeRFGraph(Graph): Dict[str, List[Parameter]]: Mapping of different parameter groups """ param_groups = {} - param_groups["cameras"] = list(self.ray_generator.parameters()) + # param_groups["cameras"] = list(self.ray_generator.parameters()) param_groups["fields"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters()) return param_groups diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 5f7926bad..4a1c9c5c1 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -137,7 +137,7 @@ class Trainer: batch = next(iter(self.train_dataloader)) loss_dict = self.train_iteration(batch, step) if step != 0 and step % self.config.steps_per_log == 0: - self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="train") + self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-") logging.info(f"{step}/{num_iterations} with losses {loss_dict}") # TODO: add the learning rates to tensorboard/logging if self.config.steps_per_save and step != 0 and step % self.config.steps_per_save == 0: diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index d1e8cd520..d1462303c 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -56,7 +56,7 @@ class Writer: """ raise NotImplementedError - def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int, group: str = None) -> None: + def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None) -> None: """Function that writes out all scalars from a given dictionary to the logger Args: @@ -64,7 +64,7 @@ class Writer: """ if self.is_main_thread: for name, scalar in scalar_dict.items(): - self.write_scalar(name, scalar, step, group=group) + self.write_scalar(name, scalar, step, group=group, prefix=prefix) class TensorboardWriter(Writer): @@ -97,7 +97,7 @@ class TensorboardWriter(Writer): self.tb_writer.add_text(name, x) @check_main_thread - def write_scalar(self, name: str, scalar: float, step: float, group: str = None) -> None: + def write_scalar(self, name: str, scalar: float, step: float, group: str = None, prefix: str = None) -> None: """Tensorboard method to write a single scalar value to the logger Args: @@ -106,8 +106,9 @@ class TensorboardWriter(Writer): y (float): y value to write group (str)): a prefix to group tensorboard scalars """ - prefix = f"{group}/" if group else "" - self.tb_writer.add_scalar(f"{prefix}{name}", scalar, step) + group_s = f"{group}/" if group else "" + prefix_s = f"{prefix}" if prefix else "" + self.tb_writer.add_scalar(f"{group_s}{prefix_s}{name}", scalar, step) class LocalWriter(Writer): diff --git a/notebooks/visualize_dataset.ipynb b/notebooks/visualize_dataset.ipynb new file mode 100644 index 000000000..6f1dec38e --- /dev/null +++ b/notebooks/visualize_dataset.ipynb @@ -0,0 +1,266 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3c00934d-0f14-4c1a-a563-94d627b377e5", + "metadata": {}, + "source": [ + "## Notebook to visualize a dataset\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b15dfdc1-5054-40bc-9631-4c7426eecfd9", + "metadata": {}, + "outputs": [], + "source": [ + "from notebook_imports import *\n", + "\n", + "import torch\n", + "from torch.utils.data import DataLoader\n", + "\n", + "from mattport.nerf.dataset.utils import get_dataset_inputs\n", + "from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch\n", + "from mattport.nerf.field_modules.ray_generator import RayGenerator\n", + "from mattport.structures.rays import RayBundle\n", + "from mattport.utils.io import get_absolute_path\n", + "\n", + "from hydra import compose, initialize\n", + "from omegaconf import open_dict\n", + "import pprint\n", + "from tqdm import tqdm\n", + "\n", + "import goat\n", + "import random\n", + "from tqdm import tqdm" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "001bdf3b-60a0-4bb2-9f88-ae36f2036c80", + "metadata": {}, + "outputs": [], + "source": [ + "with initialize(config_path=\"../configs\"):\n", + " cfg = compose(config_name=\"default.yaml\")\n", + " cfg.dataset.downscale_factor = 4\n", + " \n", + "cfg.dataset.data_directory = get_absolute_path(cfg.dataset.data_directory)\n", + "dataset_inputs = get_dataset_inputs(**cfg.dataset)\n", + "dataset = ImageDataset(\n", + " image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n", + ")\n", + "\n", + "dataloader = DataLoader(\n", + " dataset,\n", + " batch_size=cfg.dataloader.num_images_to_sample_from,\n", + " collate_fn=lambda batch: collate_batch(batch, cfg.dataloader.num_rays_per_batch, keep_full_image=True),\n", + " num_workers=cfg.dataloader.num_workers,\n", + " shuffle=True,\n", + ")\n", + "dataloader_iter = iter(dataloader)\n", + "\n", + "ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)\n", + "\n", + "num_batches = 10\n", + "for _ in tqdm(range(num_batches)):\n", + " batch = next(dataloader_iter)\n", + " ray_bundle = ray_generator.forward(batch.indices)\n", + " break" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19b85061-a1b0-458b-9d00-96cde6f37dca", + "metadata": {}, + "outputs": [], + "source": [ + "media.show_image(dataset.get_image(10))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b90f4a78-4b8b-4615-a127-7d7579874765", + "metadata": {}, + "outputs": [], + "source": [ + "def show_batch(batch):\n", + " # set the color of the sampled rays\n", + " c, y, x = [i.flatten() for i in torch.split(batch.indices, 1, dim=-1)]\n", + " batch.image[c,y,x] = 1.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", + " media.show_image((image*255).to(torch.uint8))\n", + "\n", + "def sample_and_show_batch():\n", + " batch = next(iter(dataloader))\n", + " show_batch(batch)\n", + " return batch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "09e625f7-fa1f-4b3f-acd4-bddfcdf7cd9f", + "metadata": {}, + "outputs": [], + "source": [ + "batch = sample_and_show_batch()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f5f2e2e8-8d12-4e12-9b9f-e70a4f25bcee", + "metadata": {}, + "outputs": [], + "source": [ + "# batch.keys()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd8cf48a-823e-4395-99ca-eda4c2026440", + "metadata": {}, + "outputs": [], + "source": [ + "ray_indices = batch.indices\n", + "ray_bundle = ray_generator(ray_indices)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "646107c0-f979-408a-941f-98b2ad3e07f7", + "metadata": {}, + "outputs": [], + "source": [ + "skip = 1\n", + "size = 8\n", + "data = []\n", + "data += [go.Scatter3d(\n", + " x=ray_generator.camera_to_world[::skip,0,3],\n", + " y=ray_generator.camera_to_world[::skip,1,3],\n", + " z=ray_generator.camera_to_world[::skip,2,3],\n", + " mode='markers',\n", + " name=\"origins\",\n", + " marker=dict(color='rgba(0, 0, 0, 1)', size=size)\n", + ")]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c73ba9bd-2902-4c46-b592-bc87ac8c71f6", + "metadata": {}, + "outputs": [], + "source": [ + "near_plane = 1.0\n", + "far_plane = 6.0\n", + "lines = torch.stack([ray_bundle.origins + ray_bundle.directions * near_plane,\n", + " ray_bundle.origins + ray_bundle.directions * far_plane\n", + " ], dim=1).tolist() # (num_rays, 2, 3)\n", + "lines = torch.tensor(random.sample(lines, k=100))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "991df6a2-efe3-4412-ba41-9a244e2e90cc", + "metadata": {}, + "outputs": [], + "source": [ + "data += goat.plotly_utils.get_line_segments_from_lines(lines)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0b2404d8-ca9e-4630-9f67-7c68d19458cb", + "metadata": {}, + "outputs": [], + "source": [ + "layout = go.Layout(\n", + " autosize=False,\n", + " width=1000,\n", + " height=1000,\n", + " margin=go.layout.Margin(\n", + " l=50,\n", + " r=50,\n", + " b=100,\n", + " t=100,\n", + " pad=4\n", + " ),\n", + " scene=go.layout.Scene(\n", + " aspectmode=\"data\",\n", + " camera=dict(\n", + " up=dict(x=0, y=0, z=1),\n", + " center=dict(x=0, y=0, z=0),\n", + " eye=dict(x=1.25, y=1.25, z=1.25)\n", + " )\n", + " )\n", + ")\n", + "fig = go.Figure(data=data, layout=layout)\n", + "fig.show()" + ] + }, + { + "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\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d1a11772-4692-4800-b446-cfdbe4ab92fa", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "mattport", + "language": "python", + "name": "mattport" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}