kaolin-wisp icon indicating copy to clipboard operation
kaolin-wisp copied to clipboard

Tracer / NeF: background values for additional channels

Open orperel opened this issue 2 years ago • 0 comments

The current PackedRFTracer implementation allows some control for bg_color to the RGB channel. We're missing this flexibility for additional channels (i.e. whose "clear color" may be random or some other non-zero value).

There are actually 3 entangled issues here that need to be addressed:

  1. Allow to set the background value for any channel type, not just RGB.
  2. The premultiplied alpha flag should be orthogonal to the bg color to avoid confusion, as i.e. #143
  3. Properly introduce neural fields with an env MLP that learns the bg color

The suggested fix is to extend BaseNeuralField with the concept of background functions one could register, similar to forward ones:

class BaseNeuralField:

    def _register_bg_function(self, fn, channels):
        """Registers a background function.

        Args:
            fn (function): Function to register.
            channels (list of str): Channel output names.
        """
        if isinstance(channels, str):
            channels = [channels]
        self._bg_functions[fn] = set(channels)

    @abstractmethod
    def register_background_functions(self):
         pass

class NeuralRadianceField:
    def register_background_functions(self):
        register_background_functions(self.bg_rgb, ["rgb"])
        register_background_functions(self.bg_depth, ["depth"])

    def bg_rgb(self, rays: Ray):
           ... # Return color according to bg param, or from some env decoder

    def bg_depth(self, rays: Ray):
           ... # Return color according to far plane value == rays.dist_max

The tracer should now be agnostic to background color - and instead fetch it from the neural field per channel, similar to how foreground color is obtained.

The premultiplied alpha flag will remain within the tracer, as it determines how blending should take place.

orperel avatar May 02 '23 08:05 orperel