CPU usage improvement and minor changes (#108)

* cpu fixes, masking logic but commented out, nerfw and semanticnerf in progress

* pylint

configs/data/dataset/friends_TBBT-big_living_room.yaml

@@ -1,3 +1,3 @@
 data_directory: data/friends/TBBT-big_living_room
 dataset_format: friends
-downscale_factor: 8
+downscale_factor: 1

configs/data/default_data.yaml

@@ -9,7 +9,7 @@ val_num_rays_per_chunk: 4096
 
 # dataloader options
 dataloader:
-  num_workers: 4
+  num_workers: 1
   shuffle: true
   num_images_to_sample_from: 8
   num_rays_per_batch: 1024

configs/graph/default_graph.yaml

@@ -30,11 +30,3 @@ param_groups:
       _target_: mattport.nerf.optimizers.ExponentialDecaySchedule
       lr_final: 0.000005
       max_steps: 1000000
-
-  cameras:
-    optimizer:
-      _target_: torch.optim.SGD
-      lr: 1e-4
-    scheduler:
-      _target_: torch.optim.lr_scheduler.ExponentialLR
-      gamma: 0.1

configs/instant_ngp_tcnn.yaml

@@ -3,6 +3,10 @@ defaults:
   - override graph: instant_ngp_graph
   - _self_
 method_name: instant_ngp
-data.dataloader.num_rays_per_batch: 4096
-data.val_num_rays_per_chunk: 16384
-graph.network.field_implementation: tcnn
+data:
+  dataloader:
+    num_rays_per_batch: 4096
+  val_num_rays_per_chunk: 16384
+graph:
+  network:
+    field_implementation: tcnn

mattport/nerf/dataset/format/friends.py

@@ -8,11 +8,33 @@ import numpy as np
 import torch
 
 from mattport.nerf.dataset.colmap_utils import read_cameras_binary, read_images_binary, read_pointsTD_binary
-from mattport.nerf.dataset.structs import DatasetInputs, PointCloud, Semantics
+from mattport.nerf.dataset.structs import DatasetInputs, PointCloud, SceneBounds, Semantics
 from mattport.utils import profiler
 from mattport.utils.io import load_from_json
 
 
+def get_aabb_and_transform(basedir):
+    """Returns the aabb and pointcloud transform from the threejs.json file."""
+    filename = os.path.join(basedir, "threejs.json")
+    assert os.path.exists(filename)
+    data = load_from_json(filename)
+
+    # point cloud transformation
+    transposed_point_cloud_transform = np.array(data["object"]["children"][0]["matrix"]).reshape(4, 4).T
+    assert transposed_point_cloud_transform[3, 3] == 1.0
+
+    # bbox transformation
+    bbox_transform = np.array(data["object"]["children"][1]["matrix"]).reshape(4, 4).T
+    w, h, d = data["geometries"][1]["width"], data["geometries"][1]["height"], data["geometries"][1]["depth"]
+    temp = np.array([w, h, d]) / 2.0
+    bbox = np.array([-temp, temp])
+    bbox = np.concatenate([bbox, np.ones_like(bbox[:, 0:1])], axis=1)
+    bbox = (bbox_transform @ bbox.T).T[:, 0:3]
+
+    aabb = bbox  # rename to aabb because it's an axis-aligned bounding box
+    return torch.from_numpy(aabb).float(), torch.from_numpy(transposed_point_cloud_transform).float()
+
+
 @profiler.time_function
 def load_friends_data(basedir, downscale_factor=1, split="train", include_semantics=True, include_point_cloud=False):
     """_summary_
@@ -38,6 +60,10 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant
     # TODO: handle the splits differently
     image_filenames = [os.path.join(basedir, "images", image_path) for image_path in image_paths]
 
+    # -- set the bounding box ---
+    aabb, transposed_point_cloud_transform = get_aabb_and_transform(basedir)
+    scene_bounds = SceneBounds(aabb=aabb)
+
     # --- intrinsics ---
     cameras_data = read_cameras_binary(os.path.join(basedir, "colmap", "cameras.bin"))
     intrinsics = []
@@ -61,8 +87,9 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant
         camera_to_world.append(c2w)
     camera_to_world = torch.tensor(np.array(camera_to_world)).float()
     camera_to_world = torch.inverse(camera_to_world)
-    camera_to_world = camera_to_world[:, :3]
     camera_to_world[..., 1:3] *= -1
+    camera_to_world = transposed_point_cloud_transform @ camera_to_world
+    camera_to_world = camera_to_world[:, :3]
 
     # --- masks to mask out things (e.g., people) ---
 
@@ -87,19 +114,15 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant
             thing_filenames=thing_filenames,
         )
 
-    # --- possibly transform ---
-    # TODO:
-
-    # -- set the bounding box ---
-    # TODO:
-
     # Possibly include the sparse point cloud from COLMAP in the dataset inputs.
     # NOTE(ethan): this will be common across the different splits.
     point_cloud = PointCloud()
     if include_point_cloud:
         points_3d = read_pointsTD_binary(os.path.join(basedir, "colmap", "points3D.bin"))
-        xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()]))
-        rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()]))
+        xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])).float()
+        rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float()
+        xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1)
+        xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3]
         point_cloud.xyz = xyz
         point_cloud.rgb = rgb
 
@@ -110,5 +133,6 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant
         camera_to_world=camera_to_world,
         semantics=semantics,
         point_cloud=point_cloud,
+        scene_bounds=scene_bounds,
     )
     return dataset_inputs

mattport/nerf/dataset/image_dataset.py

@@ -27,6 +27,13 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False)
 
     batch = default_collate(batch_list)
     num_images, image_height, image_width, _ = batch["image"].shape
+
+    # # checking where mask is not zero
+    # # TODO(ethan): add this to the code with a config
+    # nonzero_indices = torch.nonzero(batch["mask"], as_tuple=False)
+    # chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch)
+    # indices = nonzero_indices[chosen_indices]
+
     indices = torch.floor(
         torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width])
     ).long()

mattport/nerf/fields/instant_ngp_field.py

@@ -10,6 +10,7 @@ import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
 from mattport.nerf.field_modules.encoding import Encoding, HashEncoding, SHEncoding
+from mattport.nerf.field_modules.field_heads import FieldHeadNames
 from mattport.nerf.fields.base import Field
 from mattport.nerf.fields.nerf_field import NeRFField
 from mattport.structures.rays import PointSamples
@@ -112,7 +113,7 @@ class TCNNInstantNGPField(Field):
         h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
         h = self.mlp_head(h).view(*point_samples.directions.shape[:-1], -1).to(dtype)
         rgb = torch.sigmoid(h)
-        return {"rgb": rgb}
+        return {FieldHeadNames.RGB: rgb}
 
 
 class TorchInstantNGPField(NeRFField):

mattport/nerf/graph/base.py

@@ -107,7 +107,7 @@ class Graph(AbstractGraph):
         # during training, keep only the rays that intersect the scene. discard the rest
         valid_mask = intersected_ray_bundle.valid_mask
         masked_intersected_ray_bundle = intersected_ray_bundle.get_masked_ray_bundle(valid_mask)
-        masked_batch = get_masked_dict(batch, valid_mask)
+        masked_batch = get_masked_dict(batch, valid_mask)  # NOTE(ethan): this is really slow if on CPU!
         outputs = self.get_outputs(masked_intersected_ray_bundle)
         loss_dict = self.get_loss_dict(outputs=outputs, batch=masked_batch)
         return outputs, loss_dict
@@ -128,6 +128,9 @@ class Graph(AbstractGraph):
     @torch.no_grad()
     def get_outputs_for_camera(self, intrinsics, camera_to_world, chunk_size=1024, training_camera_index=0):
         """Takes in camera parameters and computes the output of the graph."""
+        # NOTE(ethan): this function has a spike in CPU usage
+
+        device = self.get_device()
         assert len(intrinsics.shape) == 1
         num_intrinsics_params = len(intrinsics)
         camera_class = get_camera_model(num_intrinsics_params)
@@ -135,9 +138,9 @@ class Graph(AbstractGraph):
         camera_ray_bundle = camera.generate_camera_rays()
         # TODO(ethan): decide how to properly handle the image indices for validation images
         camera_ray_bundle.set_camera_indices(camera_index=training_camera_index)
+        camera_ray_bundle.move_to_device(device)
         image_height, image_width = camera_ray_bundle.origins.shape[:2]
 
-        device = self.get_device()
         num_rays = len(camera_ray_bundle)
 
         outputs = {}
@@ -146,7 +149,6 @@ class Graph(AbstractGraph):
             start_idx = i
             end_idx = i + chunk_size
             original_ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
-            original_ray_bundle.move_to_device(device)
             intersected_ray_bundle = self.collider(original_ray_bundle)
             outputs = self.get_outputs(intersected_ray_bundle)
             for output_name, output in outputs.items():
@@ -156,5 +158,5 @@ class Graph(AbstractGraph):
         return outputs
 
     @abstractmethod
-    def log_test_image_outputs(self, image_idx, step, image, outputs):
+    def log_test_image_outputs(self, image_idx, step, image, mask, outputs):
         """Writes the test image outputs."""

mattport/nerf/graph/instant_ngp.py

@@ -118,7 +118,7 @@ class NGPGraph(Graph):
         loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, image, outputs):
+    def log_test_image_outputs(self, image_idx, step, image, mask, outputs):
         rgb = outputs["rgb"]
         acc = visualization.apply_colormap(outputs["accumulation"])
         depth = visualization.apply_depth_colormap(

mattport/nerf/graph/nerfw.py

@@ -22,28 +22,27 @@ from mattport.structures.rays import RayBundle, RaySamples
 from mattport.utils import visualization, writer
 
 
-class NerfWField(NeRFField):
+class NerfWField(Field):
     """The NeRF-W field which has appearance and transient conditioning."""
 
     def __init__(
         self,
-        num_layers=8,
-        layer_width=256,
-        skip_connections: Tuple = (4,),
-        num_images: int = 0,
-        appearance_embedding_dim: int = 48,
-        transient_embedding_dim: int = 16,
+        position_encoding: Encoding = Identity(in_dim=3),
+        direction_encoding: Encoding = Identity(in_dim=3),
+        base_mlp_num_layers: int = 8,
+        base_mlp_layer_width: int = 256,
+        head_mlp_num_layers: int = 2,
+        head_mlp_layer_width: int = 128,
+        skip_connections: Tuple[int] = (4,),
+        field_heads: Tuple[FieldHead] = (RGBFieldHead(),),
     ) -> None:
-        assert num_images > 0
+        super().__init__()
         self.num_images = num_images
-        self.appearance_embedding_dim = appearance_embedding_dim
-        self.transient_embedding_dim = transient_embedding_dim
-        super().__init__(num_layers, layer_width, skip_connections)
+        self.appearance_embedding_dim = 48
+        self.transient_embedding_dim = 16
 
-    def build_mlp_base(self):
         self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
         self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim)
-        super().build_mlp_base()
         self.mlp_transient = MLP(
             in_dim=self.mlp_base.get_out_dim() + self.embedding_transient.get_out_dim(),
             out_dim=self.layer_width // 2,
@@ -52,8 +51,7 @@ class NerfWField(NeRFField):
             activation=nn.ReLU(),
         )
 
-    def build_mlp_rgb(self):
-        self.mlp_rgb = MLP(
+        self.mlp_head = MLP(
             in_dim=self.mlp_base.get_out_dim()
             + self.encoding_dir.get_out_dim()
             + self.embedding_appearance.get_out_dim(),
@@ -193,7 +191,7 @@ class NerfWGraph(NeRFGraph):
         loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, image, outputs):
+    def log_test_image_outputs(self, image_idx, step, image, mask, outputs):
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         rgb_fine_static = outputs["rgb_fine_static"]

mattport/nerf/graph/semantic_nerf.py

@@ -149,8 +149,8 @@ class SemanticNerfGraph(NeRFGraph):
         loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, image, outputs):
-        super().log_test_image_outputs(image_idx, step, image, outputs)
+    def log_test_image_outputs(self, image_idx, step, image, mask, outputs):
+        super().log_test_image_outputs(image_idx, step, image, mask, outputs)
         semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantic_logits"], dim=-1), dim=-1) / len(
             self.stuff_classes
         )

mattport/nerf/graph/vanilla_nerf.py

@@ -129,7 +129,7 @@ class NeRFGraph(Graph):
         loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict)
         return loss_dict
 
-    def log_test_image_outputs(self, image_idx, step, image, outputs):
+    def log_test_image_outputs(self, image_idx, step, image, mask, outputs):
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"])
@@ -150,10 +150,12 @@ class NeRFGraph(Graph):
         combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1)
         combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)
         combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)
+        mask = visualization.apply_depth_colormap(mask[..., None])
 
         writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img")
         writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation")
         writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth")
+        writer.put_image(name=f"image_idx_{image_idx}", image=mask, step=step, group="mask")
 
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]

mattport/nerf/optimizers.py

@@ -14,16 +14,26 @@ from mattport.utils import writer
 
 
 class ExponentialDecaySchedule(LambdaLR):
-    """Exponential learning rate decay function."""
+    """Exponential learning rate decay function.
+    See https://github.com/google-research/google-research/blob/
+    fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360
+    for details.
+    """
 
-    def __init__(self, optimizer, lr_init, lr_final, max_steps) -> None:
+    def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None:
         def func(step):
+            if lr_delay_steps > 0:
+                delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
+                    0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
+                )
+            else:
+                delay_rate = 1.0
             t = np.clip(step / max_steps, 0, 1)
             log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
             multiplier = (
                 log_lerp / lr_init
             )  # divided by lr_init because the multiplier is with the initial learning rate
-            return multiplier
+            return delay_rate * multiplier
 
         super().__init__(optimizer, lr_lambda=func)
 

mattport/nerf/ray_generator.py

@@ -30,6 +30,12 @@ class RayGenerator(nn.Module):
         self.camera_to_world = nn.Parameter(camera_to_world, requires_grad=False)
         # TODO(ethan): add learnable parameters that are deltas on the intrinsics and camera_to_world parameters
 
+        # NOTE(ethan): we currently assume all images have the same height and width
+        camera_index = 0
+        self.camera_class = get_camera_model(self.num_intrinsics_params)
+        camera = self.camera_class(*self.intrinsics[camera_index].tolist())
+        self.image_coords = nn.Parameter(camera.get_image_coords(), requires_grad=False)
+
     def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle:
         """Index into the cameras to generate the rays.
 
@@ -41,13 +47,10 @@ class RayGenerator(nn.Module):
         x = ray_indices[:, 2]  # col indices
         intrinsics = self.intrinsics[c]
         camera_to_world = self.camera_to_world[c]
-        # NOTE(ethan): we currently assume all images have the same height and width
-        camera_index = 0
-        camera_class = get_camera_model(self.num_intrinsics_params)
-        camera = camera_class(*self.intrinsics[camera_index].tolist())
-        image_coords = camera.get_image_coords().to(ray_indices.device)  # NOTE(ethan): I can actually avoid this step
-        coords = image_coords[y, x]
+        coords = self.image_coords[y, x]
 
-        ray_bundle = camera_class.generate_rays(intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords)
+        ray_bundle = self.camera_class.generate_rays(
+            intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords
+        )
         ray_bundle.camera_indices = c
         return ray_bundle

mattport/nerf/trainer.py

@@ -21,6 +21,7 @@ from mattport.nerf.optimizers import Optimizers
 from mattport.utils import profiler, writer
 from mattport.utils.callbacks import update_occupancy
 from mattport.utils.decorators import check_main_thread
+from mattport.utils.misc import get_dict_to_torch
 
 logging.getLogger("PIL").setLevel(logging.WARNING)
 
@@ -194,14 +195,13 @@ class Trainer:
 
     def train(self) -> None:
         """_summary_"""
-        # self.graph.train()
-        # assert self.graph.training, "Graph is not in train mode!"
         train_start = time()
         num_iterations = self.config.graph.max_num_iterations
         iter_dataset = iter(self.train_dataloader)
         for step in range(self.start_step, self.start_step + num_iterations):
             data_start = time()
             batch = next(iter_dataset)
+            batch = get_dict_to_torch(batch, device=self.device)
             writer.put_time(
                 name=writer.EventName.ITER_LOAD_TIME,
                 start_time=data_start,
@@ -257,9 +257,8 @@ class Trainer:
     @profiler.time_function
     def train_iteration(self, batch: dict, step: int, _callback: List[Callable] = None):
         """Run one iteration with a batch of inputs."""
-        # move batch to correct device
-        ray_indices = batch["indices"].to(self.device)
-        _, loss_dict = self.graph(ray_indices, batch=batch, step=step)
+        ray_indices = batch["indices"]
+        _, loss_dict = self.graph.forward(ray_indices, batch=batch, step=step)
         loss = loss_dict["aggregated_loss"]
         self.optimizers.zero_grad_all()
         loss.backward()
@@ -281,7 +280,9 @@ class Trainer:
         outputs = self.graph.get_outputs_for_camera(
             intrinsics, camera_to_world, chunk_size=chunk_size, training_camera_index=training_camera_index
         )
-        image = self.val_image_dataset[image_idx]["image"].to(self.device)
-        psnr = self.graph.log_test_image_outputs(image_idx, step, image, outputs)
+        val_image_data = self.val_image_dataset[image_idx]
+        image = val_image_data["image"].to(self.device)
+        mask = val_image_data["mask"].to(self.device)
+        psnr = self.graph.log_test_image_outputs(image_idx, step, image, mask, outputs)
         self.graph.train()
         return psnr

mattport/structures/cameras.py

@@ -106,7 +106,7 @@ class Camera:
         """
         height = self.get_image_height()
         width = self.get_image_width()
-        intrinsics = self.get_intrinsics().unsqueeze(0).repeat(height, height, 1)  # (num_rays, num_intrinsics_params)
+        intrinsics = self.get_intrinsics().unsqueeze(0).repeat(height, width, 1)  # (num_rays, num_intrinsics_params)
         camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1)  # (num_rays, 3, 4)
         coords = self.get_image_coords()
         ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords)

mattport/structures/rays.py

@@ -89,12 +89,19 @@ class RaySamples:
         delta_density = self.deltas * densities[..., 0]
         alphas = 1 - torch.exp(-delta_density)
 
+        # mip-nerf version of transmittance calculation:
         transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
         transmittance = torch.cat(
             [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1
         )
         transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
 
+        # most nerf codebases do the following:
+        # transmittance = torch.cat(
+        #     [torch.ones((*alphas.shape[:1], 1)).to(densities.device), 1.0 - alphas + 1e-10], dim=-1
+        # )
+        # transmittance = torch.cumprod(transmittance, dim=-1)[..., :-1]  # [..., "num_samples"]
+
         weights = alphas * transmittance  # [..., "num_samples"]
 
         return weights

tests/nerf/dataloader/configs/blender_lego.yaml

@@ -6,7 +6,7 @@ dataset:
 
 # dataloader options
 dataloader:
-  num_workers: 0
+  num_workers: 1
   shuffle: true
   num_images_to_sample_from: 8
   num_rays_per_batch: 1024