Jax-Volumetric-Renderer

A Jax based renderer for plenoxels, work in progress...

This is an ongoing project on Plenoxels, for didactic purposes, from this paper: https://alexyu.net/plenoxels/

I am attempting to make a fast renderer in Jax, but using only a 8gb GPU. The full models as are 512x512x512, but we render only 128x128x128, more is possible but requires more batching and potentially some other tricks like voxel partitionning.

Usage:

The main work loop is in renderer_script.py, you will have to replace the relative paths so that the script can find the weights to the model.
The best is to run in the interpreter, there are a few calls to Open3d just to visualize what is happening, but actually not necessary.
The end result is in the end in an image buffer.

Other small things I've done to make things fit in memory:

Work in float16
Render one channel at a time
Find the silhouette of the model to be rendered

Things I've tried but did not work out:

Used Octrees, they just don't work with Jax due to memory coalescence.

Need to download the model weights from here: https://drive.google.com/drive/folders/128yBriW1IG_3NJ5Rp7APSTZsJqdJdfc1

To have all the requirements installed, you need:

Jax
Opencv
Numpy
Open3d, just to visualize some stuff.

Tips to install Jax:

Follow the steps on the website: https://jax.readthedocs.io/en/latest/installation.html
If you work in Pycharm, first launch the virtual env from the terminal, then launch Pycharm from the same terminal
If you call both TF2 or Torch as you call Jax (or even import them) you will likely run out of memory. They don't like each other.

TODO:

Figure out where the random black dots are coming from on the render
Provide install script
Make python install wheel
Work on the optimization/training steps
Get a GPU with more VRAM

szat/Jax-Volumetric-Renderer