From 7c9b837ba91b820d1df80b86d882621a39849aaf Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 18 Jul 2022 15:37:27 -0700 Subject: [PATCH] Instant-ngp fuse encoding with network (#171) Instant-ngp fuse encoding with network --- pyrad/data/image_dataset.py | 2 +- pyrad/fields/instant_ngp_field.py | 27 +++++++++++---------------- pyrad/fields/modules/field_heads.py | 11 +++++------ pyrad/graphs/instant_ngp.py | 4 ++-- 4 files changed, 19 insertions(+), 25 deletions(-) diff --git a/pyrad/data/image_dataset.py b/pyrad/data/image_dataset.py index a34a116ea..f88c87aae 100644 --- a/pyrad/data/image_dataset.py +++ b/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: diff --git a/pyrad/fields/instant_ngp_field.py b/pyrad/fields/instant_ngp_field.py index dd9f56497..6a35858e2 100644 --- a/pyrad/fields/instant_ngp_field.py +++ b/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 diff --git a/pyrad/fields/modules/field_heads.py b/pyrad/fields/modules/field_heads.py index 24c6b0cd3..8fcb660ee 100644 --- a/pyrad/fields/modules/field_heads.py +++ b/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 diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index 78a280971..c074bde16 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/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, ...]