From 69fa901788523dcfab749561cacc70c1fcf3af0c Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 27 Apr 2022 12:08:39 -0700 Subject: [PATCH] Append input to encoding --- mattport/nerf/field_modules/encoding.py | 13 +++++++++++-- mattport/nerf/graph/vanilla_nerf.py | 8 ++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index 276f56d08..0ff7bcb26 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -61,7 +61,9 @@ class ScalingAndOffset(Encoding): class NeRFEncoding(Encoding): """Multi-scale sinousoidal encoding proposed in the original NeRF paper""" - def __init__(self, in_dim: int, num_frequencies: int, min_freq_exp: float, max_freq_exp: float) -> None: + def __init__( + self, in_dim: int, num_frequencies: int, min_freq_exp: float, max_freq_exp: float, include_input: bool = False + ) -> None: """Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp. Args: @@ -69,15 +71,20 @@ class NeRFEncoding(Encoding): num_frequencies (int): Number of encoded frequencies per axis min_freq_exp (float): Minimum frequency exponent max_freq_exp (float): Maximum frequency exponent + include_input (float): Append the input coordinate to the encoding """ super().__init__(in_dim) self.num_frequencies = num_frequencies self.min_freq = min_freq_exp self.max_freq = max_freq_exp + self.include_input = include_input def get_out_dim(self) -> int: - return self.in_dim * self.num_frequencies * 2 + out_dim = self.in_dim * self.num_frequencies * 2 + if self.include_input: + out_dim += self.in_dim + return out_dim def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: """ @@ -91,6 +98,8 @@ class NeRFEncoding(Encoding): scaled_inputs = 2 * torch.pi * in_tensor[..., None] * freqs # [..., "input_dim", "num_scales"] scaled_inputs = scaled_inputs.view(*scaled_inputs.shape[:-2], -1) # [..., "input_dim" * "num_scales"] encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1) + if self.include_input: + encoded_inputs = torch.cat([encoded_inputs, in_tensor], axis=-1) return encoded_inputs diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 074551120..c3af7c9dc 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -30,8 +30,12 @@ class NeRFField(nn.Module): self.layer_width = layer_width self.skip_connections = skip_connections - self.encoding_xyz = NeRFEncoding(in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0) - self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0) + self.encoding_xyz = NeRFEncoding( + in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ) + self.encoding_dir = NeRFEncoding( + in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True + ) self.mlp_base = MLP( in_dim=self.encoding_xyz.get_out_dim(), out_dim=self.layer_width,