NVIDIAGameWorks / kaolin-wisp

NVIDIA Kaolin Wisp is a PyTorch library powered by NVIDIA Kaolin Core to work with neural fields (including NeRFs, NGLOD, instant-ngp and VQAD).
Other
1.46k stars 131 forks source link

Tracer / NeF: background values for additional channels #152

Open orperel opened 1 year ago

orperel commented 1 year ago

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.