From ba28d0cdb2cfda1d0d88775c8a107a35fe972673 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 4 May 2022 13:45:53 -0700 Subject: [PATCH] Add training with white backgrounds --- configs/data/default_data.yaml | 1 + configs/logging/default_logging.yaml | 2 +- mattport/nerf/dataset/image_dataset.py | 16 ++++++++++-- mattport/nerf/dataset/utils.py | 35 ++++++++++++++++++-------- mattport/nerf/graph/vanilla_nerf.py | 2 +- mattport/nerf/renderers.py | 14 +++++++++-- mattport/nerf/trainer.py | 2 ++ 7 files changed, 55 insertions(+), 17 deletions(-) diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index a42faf749..76be11297 100644 --- a/configs/data/default_data.yaml +++ b/configs/data/default_data.yaml @@ -3,6 +3,7 @@ dataset: data_directory: data/blender/lego dataset_type: blender downscale_factor: 4 + white_background: True # validation options validation_image_indices: [0, 10] # Which images to evaluate on. diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index fd5e88371..1e74b13c4 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -3,7 +3,7 @@ steps_per_log: 10 enable_stats: True stats_tracker: max_history: 10 - stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options + stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] # see mattport/utils/stats_tracker.py for options # writer logs losses and images per iteration writer: type: TensorboardWriter diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index 791546bda..8945b82eb 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -50,11 +50,19 @@ 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): + def __init__(self, image_filenames: List[str], downscale_factor: int = 1, white_background: bool = False): + """_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. + """ super().__init__() assert isinstance(downscale_factor, int) self.image_filenames = image_filenames self.downscale_factor = downscale_factor + self.white_background = white_background def __len__(self): return len(self.image_filenames) @@ -92,7 +100,11 @@ class ImageDataset(torch.utils.data.Dataset): 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) - image = original_image[:, :, :3] + + if self.white_background: + image = original_image[:, :, :3] * original_image[:, :, -1:] + (1.0 - original_image[:, :, -1:]) + else: + image = original_image[:, :, :3] num_channels = original_image.shape[2] if num_channels == 4: mask = original_image[:, :, 3] diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index 50f5216ea..8076a7ed7 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -19,21 +19,25 @@ class DatasetInputs: image_filenames: List[str] downscale_factor: float = 1.0 + white_background: bool = False intrinsics: torch.tensor = None camera_to_world: torch.tensor = None -def load_blender_data(basedir, downscale_factor=1.0, split="train"): +def load_blender_data( + basedir: str, downscale_factor: int = 1, white_background: bool = False, 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. Args: - basedir (_type_): _description_ - downscale_factor (bool, optional): _description_. Defaults to False. - testskip (int, optional): _description_. Defaults to 1. + 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. + split (str, optional): Which dataset split to generate. Returns: - _type_: _description_ + DatasetInputs """ # pylint: disable=unused-argument # pylint: disable-msg=too-many-locals @@ -62,6 +66,7 @@ def load_blender_data(basedir, downscale_factor=1.0, split="train"): dataset_inputs = DatasetInputs( image_filenames=image_filenames, downscale_factor=downscale_factor, + white_background=white_background, intrinsics=intrinsics * 1.0 / downscale_factor, # downscaling the intrinsics here camera_to_world=camera_to_world, ) @@ -70,25 +75,33 @@ def load_blender_data(basedir, downscale_factor=1.0, split="train"): def get_dataset_inputs_dict( - data_directory: str, dataset_type: str, downscale_factor: float = 1.0, splits: Tuple[str] = ("train", "val") + data_directory: str, + dataset_type: str, + downscale_factor: int = 1, + white_background: bool = False, + splits: Tuple[str] = ("train", "val"), ) -> Dict[str, 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): _description_ - dataset_type (str): _description_ - downscale_factor (float, optional): _description_. Defaults to 1.0. + 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. Returns: Dict[str, DatasetInputs]: The inputs needed for generating rays. """ dataset_inputs_dict = {} + if dataset_type == "blender": for split in splits: - dataset_inputs = load_blender_data(get_absolute_path(data_directory), downscale_factor=downscale_factor) + dataset_inputs = load_blender_data( + get_absolute_path(data_directory), downscale_factor=downscale_factor, white_background=white_background + ) dataset_inputs_dict[split] = dataset_inputs else: - raise NotImplementedError() + raise NotImplementedError(f"{dataset_type} is not a valid dataset type") return dataset_inputs_dict diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index df648d2df..2ea4cd941 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -110,7 +110,7 @@ class NeRFGraph(Graph): self.field_fine = NeRFField() # renderers - self.renderer_rgb = RGBRenderer() + self.renderer_rgb = RGBRenderer(white_background=True) # losses self.rgb_loss = MSELoss() diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index 081b2c4d2..5fde47d1a 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -19,9 +19,16 @@ class RendererOutputs: class RGBRenderer(nn.Module): """Standard volumetic rendering.""" - @classmethod + def __init__(self, white_background: bool = False) -> None: + """ + Args: + white_background (bool, optional): Composite onto white. Defaults to True. + """ + super().__init__() + self.white_background = white_background + def forward( - cls, + self, rgb: TensorType[..., "num_samples", 3], weights: TensorType[..., "num_samples"], ) -> RendererOutputs: @@ -37,5 +44,8 @@ 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] + renderer_outputs = RendererOutputs(rgb=rgb) return renderer_outputs diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 9c5ff4c87..0e8a19ae8 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -78,6 +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, ) self.train_dataset = CollateIterDataset( @@ -96,6 +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, ) @profiler.time_function