Add training with white backgrounds

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.

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

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]

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

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

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

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