| 20 | |
| 21 | |
| 22 | class DecoderSplattingCUDA(Decoder[DecoderSplattingCUDACfg]): |
| 23 | background_color: Float[Tensor, "3"] |
| 24 | |
| 25 | def __init__( |
| 26 | self, |
| 27 | cfg: DecoderSplattingCUDACfg, |
| 28 | ) -> None: |
| 29 | super().__init__(cfg) |
| 30 | self.make_scale_invariant = cfg.make_scale_invariant |
| 31 | self.register_buffer( |
| 32 | "background_color", |
| 33 | torch.tensor(cfg.background_color, dtype=torch.float32), |
| 34 | persistent=False, |
| 35 | ) |
| 36 | |
| 37 | def forward( |
| 38 | self, |
| 39 | gaussians: Gaussians, |
| 40 | extrinsics: Float[Tensor, "batch view 4 4"], |
| 41 | intrinsics: Float[Tensor, "batch view 3 3"], |
| 42 | near: Float[Tensor, "batch view"], |
| 43 | far: Float[Tensor, "batch view"], |
| 44 | image_shape: tuple[int, int], |
| 45 | depth_mode: DepthRenderingMode | None = None, |
| 46 | cam_rot_delta: Float[Tensor, "batch view 3"] | None = None, |
| 47 | cam_trans_delta: Float[Tensor, "batch view 3"] | None = None, |
| 48 | ) -> DecoderOutput: |
| 49 | b, v, _, _ = extrinsics.shape |
| 50 | color, depth = render_cuda( |
| 51 | rearrange(extrinsics, "b v i j -> (b v) i j"), |
| 52 | rearrange(intrinsics, "b v i j -> (b v) i j"), |
| 53 | rearrange(near, "b v -> (b v)"), |
| 54 | rearrange(far, "b v -> (b v)"), |
| 55 | image_shape, |
| 56 | repeat(self.background_color, "c -> (b v) c", b=b, v=v), |
| 57 | repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v), |
| 58 | repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v), |
| 59 | repeat(gaussians.harmonics, "b g c d_sh -> (b v) g c d_sh", v=v), |
| 60 | repeat(gaussians.opacities, "b g -> (b v) g", v=v), |
| 61 | scale_invariant=self.make_scale_invariant, |
| 62 | cam_rot_delta=rearrange(cam_rot_delta, "b v i -> (b v) i") if cam_rot_delta is not None else None, |
| 63 | cam_trans_delta=rearrange(cam_trans_delta, "b v i -> (b v) i") if cam_trans_delta is not None else None, |
| 64 | ) |
| 65 | color = rearrange(color, "(b v) c h w -> b v c h w", b=b, v=v) |
| 66 | |
| 67 | depth = rearrange(depth, "(b v) h w -> b v h w", b=b, v=v) |
| 68 | return DecoderOutput(color, depth) |
nothing calls this directly
no outgoing calls
no test coverage detected