diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py new file mode 100644 index 000000000..e7b3c4a46 --- /dev/null +++ b/mattport/nerf/renderers.py @@ -0,0 +1,94 @@ +""" +Collection of renderers +""" +from typing import Dict, Set + +import torch +from torch import nn +from torchtyping import TensorType + +from mattport.nerf import render_heads as rh + + +class Renderer(nn.Module): + """Base Renderer. Intended to be subclassed""" + + def required_render_heads(self) -> Set[rh.RenderHead]: + """ + Returns: + Set[RenderHead]: RenderHeads required for this renderer. + """ + raise NotImplementedError + + def forward( + self, + render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]], + deltas: TensorType[..., "num_samples"], + ) -> TensorType[..., "out_dim"]: + """Composite samples along ray and render image + + Args: + in_tensor (TensorType[..., "num_samples", "features"]): feature for each sample + deltas (TensorType[..., "num_samples"]): depth of each sample + + Returns: + TensorType[..., "out_dim"]: Rendered ray + """ + raise NotImplementedError + + +class RGB(Renderer): + """Standard volumetic rendering.""" + + def __init__(self, rgb_name, density_name) -> None: + super().__init__() + self.rgb_name = rgb_name + self.density_name = density_name + + def required_render_heads(self) -> Set[rh.RenderHead]: + return set(rh.RGBHead, rh.DensityHead) + + def forward( + self, + render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]], + deltas: TensorType[..., "num_samples"], + ) -> TensorType[..., "out_dim"]: + """Composite samples along ray and render color image + + Args: + in_tensor (TensorType[..., "num_samples", 3]): RGB for each sample + deltas (TensorType[..., "num_samples"]): depth influence for each sample + + Returns: + TensorType[..., "out_dim"]: Composited RGB ray + """ + + delta_density = deltas * render_head_outputs[self.density_name] + alphas = 1 - torch.exp(-delta_density) + + transmittance = torch.cat(0, torch.cumsum(delta_density[..., :-1], dim=-1), dim=-1) + transmittance = torch.exp(-transmittance) # [..., "num_samples"] + + weights = alphas * transmittance # [..., "num_samples"] + + rgb = torch.sum(weights[..., None] * render_head_outputs[self.rgb_name], dim=-2) + + return rgb + + +class Depth(Renderer): + """_summary_ + + Args: + Renderer (_type_): _description_ + """ + + def required_render_heads(self) -> Set[rh.RenderHead]: + return set(rh.DensityHead) + + def forward( + self, + render_head_outputs: Dict["str", TensorType[..., "num_samples", "features"]], + deltas: TensorType[..., "num_samples"], + ) -> TensorType[..., "out_dim"]: + raise NotImplementedError