dmfenton/splatt3r

1import torch2from einops import rearrange, repeat3 4from .cuda_splatting import render_cuda5from utils.geometry import normalize_intrinsics6 7 8class DecoderSplattingCUDA(torch.nn.Module):9 10    def __init__(self, background_color):11        super().__init__()12        self.register_buffer(13            "background_color",14            torch.tensor(background_color, dtype=torch.float32),15            persistent=False,16        )17    18    def forward(self, batch, pred1, pred2, image_shape):19 20        base_pose = batch['context'][0]['camera_pose'] # [b, 4, 4]21        inv_base_pose = torch.inverse(base_pose)22 23        extrinsics = torch.stack([target_view['camera_pose'] for target_view in batch['target']], dim=1)24        intrinsics = torch.stack([target_view['camera_intrinsics'] for target_view in batch['target']], dim=1)25        intrinsics = normalize_intrinsics(intrinsics, image_shape)[..., :3, :3]26 27        # Rotate the ground truth extrinsics into the coordinate system used by MAST3R28        # --i.e. in the coordinate system of the first context view, normalized by the scene scale29        extrinsics = inv_base_pose[:, None, :, :] @ extrinsics30 31        means = torch.stack([pred1["means"], pred2["means_in_other_view"]], dim=1)32        covariances = torch.stack([pred1["covariances"], pred2["covariances"]], dim=1)33        harmonics = torch.stack([pred1["sh"], pred2["sh"]], dim=1)34        opacities = torch.stack([pred1["opacities"], pred2["opacities"]], dim=1)35 36        b, v, _, _ = extrinsics.shape37        near = torch.full((b, v), 0.1, device=means.device)38        far = torch.full((b, v), 1000.0, device=means.device)39 40        color = render_cuda(41            rearrange(extrinsics, "b v i j -> (b v) i j"),42            rearrange(intrinsics, "b v i j -> (b v) i j"),43            rearrange(near, "b v -> (b v)"),44            rearrange(far, "b v -> (b v)"),45            image_shape,46            repeat(self.background_color, "c -> (b v) c", b=b, v=v),47            repeat(rearrange(means, "b v h w xyz -> b (v h w) xyz"), "b g xyz -> (b v) g xyz", v=v),48            repeat(rearrange(covariances, "b v h w i j -> b (v h w) i j"), "b g i j -> (b v) g i j", v=v),49            repeat(rearrange(harmonics, "b v h w c d_sh -> b (v h w) c d_sh"), "b g c d_sh -> (b v) g c d_sh", v=v),50            repeat(rearrange(opacities, "b v h w 1 -> b (v h w)"), "b g -> (b v) g", v=v),51        )52        color = rearrange(color, "(b v) c h w -> b v c h w", b=b, v=v)53        return color, None