Added median depth rasterization for 3dgs [fwd only]#930
Added median depth rasterization for 3dgs [fwd only]#930sangeethagramas wants to merge 1 commit intonerfstudio-project:mainfrom
Conversation
| render_median, \ | ||
| median_ids \ |
There was a problem hiding this comment.
I think we should make them optional tensors -- to make sure we dont introduce any perf regression when median depth is not needed.
| scalar_t | ||
| *__restrict__ render_median, // [I, image_height, image_width, 1] | ||
| // depth of the Gaussian that causes | ||
| // transmittance to cross 0.5, or the | ||
| // last-hit depth as a fallback. | ||
| int32_t *__restrict__ median_ids // [I, image_height, image_width] |
There was a problem hiding this comment.
related to optional: we want to allow them to be null ptr when not needed. and in the function we skip compute median depth and write out so we dont waste compute if they are not needed
|
|
||
| // Track last hit depth (depth is stored as the last channel of | ||
| // colors, consistent with the existing D/ED plumbing and 2DGS). | ||
| const float dep = c_ptr[CDIM - 1]; |
There was a problem hiding this comment.
This silently assume render mode is in ED or D. Should assert this to protect in the rendering.py
There was a problem hiding this comment.
Also for the best performance, I think you can push this link and next into the if (!median_found && T > 0.5f && next_T <= 0.5f) { condition?
| // Median depth: if transmittance never crossed 0.5 fall back to the | ||
| // last contributing Gaussian's depth (0.0 if there were no hits). | ||
| if (!median_found) { | ||
| median_depth = last_hit_depth; | ||
| } | ||
| render_median[pix_id] = median_depth; | ||
| median_ids[pix_id] = static_cast<int32_t>(median_idx); |
There was a problem hiding this comment.
Skip write if median is not required (render_median is a null ptr)
| at::Tensor renders, // [..., image_height, image_width, channels] | ||
| at::Tensor alphas, // [..., image_height, image_width] | ||
| at::Tensor last_ids, // [..., image_height, image_width] | ||
| at::Tensor render_median, // [..., image_height, image_width, 1] |
There was a problem hiding this comment.
to be consistent on the naming convension in gsplat, suggested naming is median_depths
| last_ids.data_ptr<int32_t>() | ||
| last_ids.data_ptr<int32_t>(), | ||
| render_median.data_ptr<float>(), | ||
| median_ids.data_ptr<int32_t>() |
There was a problem hiding this comment.
update for optional
| packed: If True, the input tensors are expected to be packed with shape [nnz, ...]. Default: False. | ||
| absgrad: If True, the backward pass will compute a `.absgrad` attribute for `means2d`. Default: False. | ||
| return_median: If True, additionally returns per-pixel median depth, defined | ||
| as the value of the last channel of ``colors`` for the Gaussian whose |
There was a problem hiding this comment.
"the value of the last channel of colors " is weird to read. Maybe just say something like "expect the last channel to be depth"
| ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" | ||
|
|
||
| render_colors, render_alphas = _RasterizeToPixels.apply( | ||
| render_colors, render_alphas, render_median, _median_ids = _RasterizeToPixels.apply( |
There was a problem hiding this comment.
if we are not returning median ids we might just remove it entirely?
| extra_signals_sh_degree: Optional[ | ||
| int | ||
| ] = None, # Currently only None or 3 is accepted. | ||
| depth_mode: Literal["expected", "median"] = "expected", |
There was a problem hiding this comment.
"depth_mode" is a bit weird here because we already say render mode has ED or D. this "depth mode" between expected and median could be very confusing on the relationship between the ED and D in render mode.
My suggested interface change:
use argument return_median_depth: bool = False here. Then additionally return median depth (and id if you need it) if this is True. If not only return the original returns, and do not waste any compute in the kernel for median depth. This way we are backward compatible and cause minimal confusion.
2 participants
Change : Add median depth rasterization for 3dgs for the forward pass only
Description : This commit currently supports median depth calculation for 3D gaussians only on the forward pass. We currently do not anticipate the need for median depth support in the backward pass. We can add it in the future if the need arises.
The changes in this PR follows the style similar to 2D gaussian median depth rasterization.
Test cases :