resizing image

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

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

environment/requirements.txt

@@ -1,4 +1,5 @@
 black
+hydra-colorlog
 hydra-core
 imageio==2.16.1
 ipywidgets>=7.6

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

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,
     )
 

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

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

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
 

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:

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):

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
+}