Instant-ngp fuse encoding with network (#171)

Instant-ngp fuse encoding with network

pyrad/data/image_dataset.py

@@ -91,7 +91,7 @@ class ImageDataset(torch.utils.data.Dataset):
     def get_image(self, image_idx: int):
         """Returns a 3 channel image."""
         image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0)
-        if self.alpha_color is not None:
+        if self.alpha_color is not None and image.shape[-1] == 4:
             assert image.shape[-1] == 4
             image = image[:, :, :3] * image[:, :, -1:] + self.alpha_color * (1.0 - image[:, :, -1:])
         else:

pyrad/fields/instant_ngp_field.py

@@ -63,18 +63,6 @@ class TCNNInstantNGPField(Field):
         # TODO: set this properly based on the aabb
         per_level_scale = 1.4472692012786865
 
-        self.position_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={
-                "otype": "HashGrid",
-                "n_levels": 16,
-                "n_features_per_level": 2,
-                "log2_hashmap_size": 19,
-                "base_resolution": 16,
-                "per_level_scale": per_level_scale,
-            },
-        )
-
         self.direction_encoding = tcnn.Encoding(
             n_input_dims=3,
             encoding_config={
@@ -83,9 +71,17 @@ class TCNNInstantNGPField(Field):
             },
         )
 
-        self.mlp_base = tcnn.Network(
-            n_input_dims=32,
+        self.mlp_base = tcnn.NetworkWithInputEncoding(
+            n_input_dims=3,
             n_output_dims=1 + self.geo_feat_dim,
+            encoding_config={
+                "otype": "HashGrid",
+                "n_levels": 16,
+                "n_features_per_level": 2,
+                "log2_hashmap_size": 19,
+                "base_resolution": 16,
+                "per_level_scale": per_level_scale,
+            },
             network_config={
                 "otype": "FullyFusedMLP",
                 "activation": "ReLU",
@@ -112,8 +108,7 @@ class TCNNInstantNGPField(Field):
         positions = get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
         positions_flat = positions.view(-1, 3)
         dtype = positions_flat.dtype
-        x = self.position_encoding(positions_flat)
-        h = self.mlp_base(x).view(*ray_samples.frustums.get_positions().shape[:-1], -1).to(dtype)
+        h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1).to(dtype)
         density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
 
         # Rectifying the density with an exponential is much more stable than a ReLU or

pyrad/fields/modules/field_heads.py

@@ -16,7 +16,7 @@
 Collection of render heads
 """
 from enum import Enum
-from typing import Optional
+from typing import Callable, Optional, Union
 
 from torch import nn
 from torchtyping import TensorType
@@ -51,7 +51,7 @@ class FieldHead(FieldModule):
         out_dim: int,
         field_head_name: FieldHeadNames,
         in_dim: Optional[int] = None,
-        activation: Optional[nn.Module] = None,
+        activation: Optional[Union[nn.Module, Callable]] = None,
     ) -> None:
 
         super().__init__()
@@ -69,10 +69,7 @@ class FieldHead(FieldModule):
         self._construct_net()
 
     def _construct_net(self):
-        layers = [nn.Linear(self.in_dim, self.out_dim)]
-        if self.activation:
-            layers.append(self.activation)
-        self.net = nn.Sequential(*layers)
+        self.net = nn.Linear(self.in_dim, self.out_dim)
 
     def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType:
         """Process network output for renderer
@@ -86,6 +83,8 @@ class FieldHead(FieldModule):
         if not self.net:
             raise SystemError("in_dim not set. Must be provided to construtor, or set_in_dim() should be called.")
         out_tensor = self.net(in_tensor)
+        if self.activation:
+            out_tensor = self.activation(out_tensor)
         return out_tensor
 
 

pyrad/graphs/instant_ngp.py

@@ -157,8 +157,8 @@ class NGPGraph(Graph):
         writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
         writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
 
-        depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"])
-        writer.put_image(name=f"depth_occupancy_grid/image_idx_{image_idx}", image=combined_depth, step=step)
+        occupancy_depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"])
+        writer.put_image(name=f"depth_occupancy_grid/image_idx_{image_idx}", image=occupancy_depth, step=step)
 
         # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
         image = torch.moveaxis(image, -1, 0)[None, ...]