diff --git a/nerf/encodings.py b/nerf/encodings.py new file mode 100644 index 000000000..e89cd60cc --- /dev/null +++ b/nerf/encodings.py @@ -0,0 +1,57 @@ +""" +Encoding functions +""" +from torch import nn +from torchtyping import TensorType + + +class Encoding(nn.Module): + """Encode an input tensor. Intended to be subclassed""" + + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + """Encodes an input tensor. + + Args: + x (TensorType[..., "input_dim"]): Input tensor to be encoded + + Returns: + TensorType[..., "output_dim"]: A encoded input tensor + """ + raise NotImplementedError + + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + """Call forward""" + return self.encode(in_tensor) + + +class RFFEncoding(Encoding): + """Random Fourier Feature encoding""" + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + """Encodes an input tensor. + + Args: + x (TensorType[..., "input_dim"]): Input tensor to be encoded + + Returns: + TensorType[..., "output_dim"]: A encoded input tensor + """ + raise NotImplementedError + + +class HashEncoding(Encoding): + """Hash encoding""" + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + """Encodes an input tensor. + + Args: + x (TensorType[..., "input_dim"]): Input tensor to be encoded + + Returns: + TensorType[..., "output_dim"]: A encoded input tensor + """ + raise NotImplementedError + +encodings_dict = { + "RFF": RFFEncoding, + "hash": HashEncoding, +} \ No newline at end of file