Merge pull request #55 from ethanweber/tancik/multi_color_optinos

Add multiple color options

configs/data/default_data.yaml

@@ -3,7 +3,7 @@ dataset:
   data_directory: data/blender/lego
   dataset_type: blender
   downscale_factor: 4
-  white_background: True
+  alpha_color: white
 
 # validation options
 validation_image_indices: [0, 10] # Which images to evaluate on.

mattport/nerf/dataset/image_dataset.py

@@ -2,7 +2,7 @@
 Some dataset code.
 """
 
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -10,6 +10,7 @@ import PIL
 import torch
 from PIL import Image
 from torch.utils.data import default_collate
+from torchtyping import TensorType
 
 
 def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False):
@@ -50,19 +51,21 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False)
 class ImageDataset(torch.utils.data.Dataset):
     """Dataset that returns images."""
 
-    def __init__(self, image_filenames: List[str], downscale_factor: int = 1, white_background: bool = False):
+    def __init__(
+        self, image_filenames: List[str], downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None
+    ):
         """_summary_
 
         Args:
             image_filenames (List[str]): List of image filenames
             downscale_factor (int, optional): How much to downscale the image. Defaults to 1.
-            white_background (bool, optional): Sets transparent regions to white, otherwise black. Defaults to False.
+            alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black.
         """
         super().__init__()
         assert isinstance(downscale_factor, int)
         self.image_filenames = image_filenames
         self.downscale_factor = downscale_factor
-        self.white_background = white_background
+        self.alpha_color = alpha_color
 
     def __len__(self):
         return len(self.image_filenames)
@@ -101,8 +104,10 @@ class ImageDataset(torch.utils.data.Dataset):
         # 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)
 
-        if self.white_background:
-            image = original_image[:, :, :3] * original_image[:, :, -1:] + (1.0 - original_image[:, :, -1:])
+        if self.alpha_color is not None:
+            image = original_image[:, :, :3] * original_image[:, :, -1:] + self.alpha_color * (
+                1.0 - original_image[:, :, -1:]
+            )
         else:
             image = original_image[:, :, :3]
         num_channels = original_image.shape[2]

mattport/nerf/dataset/utils.py

@@ -4,11 +4,14 @@ For loading the blender dataset format.
 
 import os
 from dataclasses import dataclass
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import imageio
 import numpy as np
+from omegaconf import ListConfig
 import torch
+from torchtyping import TensorType
+from mattport.structures.colors import get_color
 
 from mattport.utils.io import get_absolute_path, load_from_json
 
@@ -19,13 +22,13 @@ class DatasetInputs:
 
     image_filenames: List[str]
     downscale_factor: float = 1.0
-    white_background: bool = False
+    alpha_color: Optional[TensorType[3]] = None
     intrinsics: torch.tensor = None
     camera_to_world: torch.tensor = None
 
 
 def load_blender_data(
-    basedir: str, downscale_factor: int = 1, white_background: bool = False, split: str = "train"
+    basedir: str, downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, split: str = "train"
 ) -> DatasetInputs:
     """Processes the a blender dataset directory.
     Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37.
@@ -33,7 +36,7 @@ def load_blender_data(
     Args:
         data_directory (str): Location of data
         downscale_factor (int, optional): How much to downscale images. Defaults to 1.0.
-        white_background (bool, optional): Set pixel with zero opacity to white. Defaults to False.
+        alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black.
         split (str, optional): Which dataset split to generate.
 
     Returns:
@@ -66,7 +69,7 @@ def load_blender_data(
     dataset_inputs = DatasetInputs(
         image_filenames=image_filenames,
         downscale_factor=downscale_factor,
-        white_background=white_background,
+        alpha_color=alpha_color,
         intrinsics=intrinsics * 1.0 / downscale_factor,  # downscaling the intrinsics here
         camera_to_world=camera_to_world,
     )
@@ -78,7 +81,7 @@ def get_dataset_inputs_dict(
     data_directory: str,
     dataset_type: str,
     downscale_factor: int = 1,
-    white_background: bool = False,
+    alpha_color: Optional[Union[str, list, ListConfig]] = None,
     splits: Tuple[str] = ("train", "val"),
 ) -> Dict[str, DatasetInputs]:
     """Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator.
@@ -88,17 +91,22 @@ def get_dataset_inputs_dict(
         data_directory (str): Location of data
         dataset_type (str): Name of dataset type
         downscale_factor (int, optional): How much to downscale images. Defaults to 1.0.
-        white_background (bool, optional): Set pixel with zero opacity to white. Defaults to False.
+        alpha_color (str, list, optional): Sets transparent regions to specified color, otherwise black.
 
     Returns:
         Dict[str, DatasetInputs]: The inputs needed for generating rays.
     """
     dataset_inputs_dict = {}
 
+    if alpha_color is not None:
+        alpha_color = get_color(alpha_color)
+
     if dataset_type == "blender":
         for split in splits:
             dataset_inputs = load_blender_data(
-                get_absolute_path(data_directory), downscale_factor=downscale_factor, white_background=white_background
+                get_absolute_path(data_directory),
+                downscale_factor=downscale_factor,
+                alpha_color=alpha_color,
             )
             dataset_inputs_dict[split] = dataset_inputs
     else:

mattport/nerf/graph/vanilla_nerf.py

@@ -17,7 +17,8 @@ from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
 from mattport.nerf.renderers import RGBRenderer
-from mattport.nerf.sampler import PDFSampler, UniformSampler  # pylint: disable=unused-import
+from mattport.nerf.sampler import PDFSampler, UniformSampler
+from mattport.structures import colors
 from mattport.structures.rays import RaySamples
 
 
@@ -110,7 +111,7 @@ class NeRFGraph(Graph):
         self.field_fine = NeRFField()
 
         # renderers
-        self.renderer_rgb = RGBRenderer(white_background=True)
+        self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
 
         # losses
         self.rgb_loss = MSELoss()

mattport/nerf/renderers.py

@@ -2,6 +2,7 @@
 Collection of renderers
 """
 from dataclasses import dataclass
+from typing import Optional
 
 import torch
 from torch import nn
@@ -19,13 +20,13 @@ class RendererOutputs:
 class RGBRenderer(nn.Module):
     """Standard volumetic rendering."""
 
-    def __init__(self, white_background: bool = False) -> None:
+    def __init__(self, background_color: Optional[TensorType[3]] = None) -> None:
         """
         Args:
-            white_background (bool, optional): Composite onto white. Defaults to True.
+            background_color (TensorType[3], optional): Background color as RGB. Defaults to black.
         """
         super().__init__()
-        self.white_background = white_background
+        self.background_color = background_color
 
     def forward(
         self,
@@ -44,8 +45,13 @@ class RGBRenderer(nn.Module):
 
         rgb = torch.sum(weights[..., None] * rgb, dim=-2)
 
-        if self.white_background:
-            rgb = rgb + (1.0 - torch.sum(weights, dim=-1))[..., None]
+        if self.background_color is not None:
+            rgb = rgb + self.background_color.to(weights.device)[None, ...] * (
+                1.0 - torch.sum(weights, dim=-1)[..., None]
+            )
+
+        assert torch.max(rgb) <= 1.0
+        assert torch.min(rgb) >= 0.0
 
         renderer_outputs = RendererOutputs(rgb=rgb)
         return renderer_outputs

mattport/nerf/trainer.py

@@ -78,7 +78,7 @@ class Trainer:
         self.train_image_dataset = ImageDataset(
             image_filenames=dataset_inputs_dict["train"].image_filenames,
             downscale_factor=dataset_inputs_dict["train"].downscale_factor,
-            white_background=dataset_inputs_dict["train"].white_background,
+            alpha_color=dataset_inputs_dict["train"].alpha_color,
         )
 
         self.train_dataset = CollateIterDataset(
@@ -97,7 +97,7 @@ class Trainer:
         self.val_image_dataset = ImageDataset(
             image_filenames=dataset_inputs_dict["val"].image_filenames,
             downscale_factor=dataset_inputs_dict["val"].downscale_factor,
-            white_background=dataset_inputs_dict["val"].white_background,
+            alpha_color=dataset_inputs_dict["val"].alpha_color,
         )
 
     @profiler.time_function

mattport/structures/colors.py

@@ -0,0 +1,40 @@
+"""Common Colors"""
+from typing import Union
+from omegaconf import ListConfig
+import torch
+from torchtyping import TensorType
+
+WHITE = torch.tensor([1.0, 1.0, 1.0])
+BLACK = torch.tensor([0.0, 0.0, 0.0])
+RED = torch.tensor([1.0, 0.0, 0.0])
+GREEN = torch.tensor([0.0, 1.0, 0.0])
+BLUE = torch.tensor([0.0, 0.0, 1.0])
+
+COLORS_DICT = {
+    "white": WHITE,
+    "black": BLACK,
+    "red": RED,
+    "green": GREEN,
+    "blue": BLUE,
+}
+
+
+def get_color(color: Union[str, list, ListConfig]) -> TensorType[3]:
+    """
+    Args:
+        color (Union[str, list, ListConfig]): Color as a string or a rgb list
+
+    Returns:
+        TensorType[3]: Parsed color
+    """
+    if isinstance(color, str):
+        color = color.lower()
+        if color not in COLORS_DICT:
+            raise ValueError(f"{color} is not a valid preset color")
+        return COLORS_DICT[color]
+    if isinstance(color, (list, ListConfig)):
+        if len(color) != 3:
+            raise ValueError(f"Color should be 3 values (RGB) instead got {color}")
+        return torch.tensor(color)
+
+    raise ValueError(f"Color should be an RGB list or string, instead got {type(color)}")