[WIP] Unified depth rasterizer

[maturk]

Testing out unified rgb and depth rasterization in a single forward pass. Collaborating with @nikmo33 !

Currently these are the speeds on my machine (4090 GPU):

Forward pass:

• Unified RGB+Depth Rasterization: 0.105 s
• Separate RGB followed by Depth Rasterization: 0.124 s

Backward pass:

• Unified RGB+Depth Backward: 0.098 s
• Separate RGB followed by Depth backward: 0.100 s

Made an examples/depth_trainer.py script to test:

Now that there is the option to return_alpha in rasterize.py, I am not sure what would be the cleanest way to merge an additional return_depth option. Any suggestions?

[kerrj]

Cool! have you tried this in 3D scenes with supervision? For API I think including render_depth is reasonable, although it might be annoying if there's an eventual blowup of the tuple output.... Another discussion we had was outputting the image as either 3 or 4 channel based on return_alpha, and then depth could be included as another element in the tuple

[vye16]

Awesome thank you for doing this! With the depth rasterization, I think it would make sense to pack the RGBD into a float4 array, rather than a separate float* depth ad float3* colors. It might improve cache hits and performance -- thoughts @kerrj?

[nikmo33]

Awesome thank you for doing this! With the depth rasterization, I think it would make sense to pack the RGBD into a float4 array, rather than a separate float* depth ad float3* colors. It might improve cache hits and performance -- thoughts @kerrj?

@vye16 I think that is reasonable. One of the main points of contention is whether we want to add this to the default rasteriser or if we want to have a separate one that supports depth backward passes. If we are happy to just add this directly to the default rasteriser then I think things could be simpler and we could add some of the optimisations you mentioned such as using float4 as well.

[ethanweber]

I think it might make sense to have a separate function for rasterize 4-dim. We have a rasterize 3-dim (for rgb) and N-dim, but we could also have one for 4-dim. Maybe it just checks the dimension of the input like the current code does to switch between 3, 4, and N? Then you can pass in the depths to compute the rendered depth as rgbd.

[nikmo33]

I think it might make sense to have a separate function for rasterize 4-dim. We have a rasterize 3-dim (for rgb) and N-dim, but we could also have one for 4-dim. Maybe it just checks the dimension of the input like the current code does to switch between 3, 4, and N? Then you can pass in the depths to compute the rendered depth as rgbd.

Hmm so would we also have an N+1 for n-dim with depth backward support? I dont think we can treat depth as just another feature channel as the gradients are different to the 3d means than if it was just treated as a separate feature?

[ethanweber]

@nikmo33 Interesting, I'll have to check out your new code. Is the current code in Splatfacto incorrect with its gradients after dividing out alpha?