diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index 76be11297..4f0e661e2 100644 --- a/configs/data/default_data.yaml +++ b/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. diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index 8945b82eb..dfa59f299 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/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] diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index 8076a7ed7..69031846b 100644 --- a/mattport/nerf/dataset/utils.py +++ b/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: diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 2ea4cd941..bd93f9158 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/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() diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index 5fde47d1a..929014ce0 100644 --- a/mattport/nerf/renderers.py +++ b/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 diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 0e8a19ae8..5289c796e 100644 --- a/mattport/nerf/trainer.py +++ b/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 diff --git a/mattport/structures/colors.py b/mattport/structures/colors.py new file mode 100644 index 000000000..343342c68 --- /dev/null +++ b/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)}")