src/model/encoder/of3gs.py

import copy
import os
import sys
from dataclasses import dataclass
from typing import List, Literal, Optional
from einops import rearrange
import torch
import torch.nn.functional as F
from einops import rearrange
from jaxtyping import Float
from safetensors.torch import load_file
from src.dataset.shims.normalize_shim import apply_normalize_shim
from src.dataset.types import BatchedExample, DataShim

from src.model.encoder.heads.vggt_dpt_gs_head import VGGT_DPT_GS_Head
from src.model.encoder.vggt.utils.geometry import (
    batchify_unproject_depth_map_to_point_map,
    closed_form_inverse_se3,
)
from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri
from torch import nn, Tensor


def scatter_add(
    source: Tensor,
    index: Tensor,
    dim: int = 0,
    dim_size: int | None = None,
) -> Tensor:
    """Small torch_scatter-compatible helper for the dimensions used by OF3GS."""
    if dim != 0:
        raise ValueError("OF3GS scatter_add currently supports dim=0 only")
    size = dim_size
    if size is None:
        size = int(index.max().item()) + 1 if index.numel() else 0
    output = source.new_zeros((size, *source.shape[1:]))
    expanded_index = index.reshape(-1, *([1] * (source.ndim - 1))).expand_as(source)
    return output.scatter_add_(0, expanded_index, source)


def scatter_max(
    source: Tensor,
    index: Tensor,
    dim: int = 0,
    dim_size: int | None = None,
) -> tuple[Tensor, None]:
    """Return scatter maxima; OF3GS does not consume the argmax output."""
    if dim != 0:
        raise ValueError("OF3GS scatter_max currently supports dim=0 only")
    size = dim_size
    if size is None:
        size = int(index.max().item()) + 1 if index.numel() else 0
    output = source.new_full((size, *source.shape[1:]), -float("inf"))
    expanded_index = index.reshape(-1, *([1] * (source.ndim - 1))).expand_as(source)
    output.scatter_reduce_(0, expanded_index, source, reduce="amax", include_self=True)
    return output, None

from ..types import Gaussians
from .backbone import BackboneCfg

from .common.gaussian_adapter import (
    GaussianAdapter,
    GaussianAdapterCfg,
    UnifiedGaussianAdapter,
)
from .encoder import Encoder, EncoderOutput
from .visualization.encoder_visualizer_epipolar_cfg import EncoderVisualizerEpipolarCfg
from src.model.encoder.vggt.models.vggt import VGGT
from src.model.encoder.streamvggt.models.streamvggt import StreamVGGT

root_path = os.path.abspath(".")
sys.path.append(root_path)
inf = float("inf")


@dataclass
class OpacityMappingCfg:
    initial: float
    final: float
    warm_up: int


@dataclass
class GSHeadParams:
    dec_depth: int = 23
    patch_size: tuple[int, int] = (14, 14)
    enc_embed_dim: int = 2048
    dec_embed_dim: int = 2048
    feature_dim: int = 256
    depth_mode = ("exp", -inf, inf)
    conf_mode = True


@dataclass
class EncoderOF3GSCfg:
    name: Literal["of3gs"]
    anchor_feat_dim: int
    voxel_size: float
    n_offsets: int
    d_feature: int
    add_view: bool
    num_monocular_samples: int
    backbone: BackboneCfg
    visualizer: EncoderVisualizerEpipolarCfg
    gaussian_adapter: GaussianAdapterCfg
    apply_bounds_shim: bool
    opacity_mapping: OpacityMappingCfg
    gaussians_per_pixel: int
    num_surfaces: int
    gs_params_head_type: str
    input_mean: tuple[float, float, float] = (0.5, 0.5, 0.5)
    input_std: tuple[float, float, float] = (0.5, 0.5, 0.5)
    pretrained_weights: str = ""
    pose_free: bool = True
    pred_pose: bool = True
    gt_pose_to_pts: bool = False
    gs_prune: bool = False
    opacity_threshold: float = 0.001
    gs_keep_ratio: float = 1.0
    pred_head_type: Literal["depth", "point"] = "point"
    freeze_backbone: bool = False
    freeze_module: Literal[
        "all",
        "global",
        "frame",
        "patch_embed",
        "patch_embed+frame",
        "patch_embed+global",
        "global+frame",
        "None",
    ] = "None"
    distill: bool = False
    render_conf: bool = False
    opacity_conf: bool = False
    conf_threshold: float = 0.1
    intermediate_layer_idx: Optional[List[int]] = None
    voxelize: bool = False
    mode: str = ""
    pre_vggt_path: str | None = ""
    pre_svggt_path: str | None = ""
    pre_dav3_path: str | None = ""


class CameraDec(nn.Module):
    def __init__(self, dim_in=2048):
        super().__init__()
        output_dim = dim_in
        self.backbone = nn.Sequential(
            nn.Linear(output_dim, output_dim),
            nn.ReLU(),
            nn.Linear(output_dim, output_dim),
            nn.ReLU(),
        )
        self.fc_fov = nn.Sequential(nn.Linear(output_dim, 1), nn.ReLU())

    def forward(self, feat):
        B, N, C = feat.shape
        feat_single = feat[:, 0, :]
        combined_feat = torch.cat([feat_single], dim=0)
        x = self.backbone(combined_feat)
        out = self.fc_fov(x.float())
        out_fov_single = out[:B].reshape(B, 1, 1)
        return out_fov_single


class EncoderOF3GS(Encoder[EncoderOF3GSCfg]):
    backbone: nn.Module
    gaussian_adapter: GaussianAdapter

    def __init__(self, cfg: EncoderOF3GSCfg) -> None:
        super().__init__(cfg)
        self.freeze_backbone = cfg.freeze_backbone
        self.distill = cfg.distill
        self.pred_pose = cfg.pred_pose

        if cfg.mode == "train":
            model_full = VGGT()
            ckpt = load_file(cfg.pre_vggt_path)
            model_full.load_state_dict(ckpt, strict=True)

            self.aggregator = model_full.aggregator.to(torch.float16)
            self.camera_head = model_full.camera_head

            if self.distill:
                self.distill_aggregator = copy.deepcopy(self.aggregator).to("cuda")
                self.distill_camera_head = copy.deepcopy(self.camera_head).to("cuda")
                for module in [
                    self.distill_aggregator,
                    self.distill_camera_head,
                ]:
                    for param in module.parameters():
                        param.requires_grad = False
            del model_full

        print("Initializing StreamVGGT model...")
        model_full = StreamVGGT()
        if cfg.mode == "train":
            ckpt = load_file(cfg.pre_svggt_path)
            model_full.load_state_dict(ckpt, strict=False)
        self.aggregator = model_full.aggregator

        self.cam_dec = CameraDec()
        if cfg.mode == "train":
            cam_ckpt = load_file(cfg.pre_dav3_path)
            cam_ckpt = {
                key.replace("model.", ""): value for key, value in cam_ckpt.items()
            }

            cam_ckpt = {
                key: value for key, value in cam_ckpt.items() if "cam_dec" in key
            }
            cam_ckpt = {
                key.replace("cam_dec.", ""): value for key, value in cam_ckpt.items()
            }
            if "fc_fov.0.weight" in cam_ckpt:
                cam_ckpt["fc_fov.0.weight"] = cam_ckpt["fc_fov.0.weight"][:1, :]

            if "fc_fov.0.bias" in cam_ckpt:
                cam_ckpt["fc_fov.0.bias"] = cam_ckpt["fc_fov.0.bias"][:1]
            self.cam_dec.load_state_dict(cam_ckpt, strict=False)
        for param in self.cam_dec.parameters():
            param.requires_grad = True

        self.camera_head = model_full.camera_head
        
        for module in [
            self.aggregator,
            self.camera_head,
        ]:
            for param in module.parameters():
                param.requires_grad = False
                        
        del model_full

        self.pose_free = cfg.pose_free
        if self.pose_free:
            self.gaussian_adapter = UnifiedGaussianAdapter(cfg.gaussian_adapter)
        else:
            self.gaussian_adapter = GaussianAdapter(cfg.gaussian_adapter)

        self.raw_gs_dim = 1 + self.gaussian_adapter.d_in  # 1 for opacity
        self.voxel_size = cfg.voxel_size
        self.gs_params_head_type = cfg.gs_params_head_type
        # fake backbone for head parameters
        head_params = GSHeadParams()
        self.gaussian_param_head = VGGT_DPT_GS_Head(
            dim_in=2048,
            patch_size=head_params.patch_size,
            output_dim=2,
            activation="norm_exp",
            conf_activation="expp1",
            features=head_params.feature_dim,
        )
        feature_dim = 256
        self.sh_degree = 0
        self.nums_sh = (self.sh_degree + 1) ** 2
        gaussian_raw_channels = 12 + self.nums_sh * 3
        self.gs_head = nn.Sequential(
            nn.Conv2d(
                feature_dim // 2, feature_dim, kernel_size=3, padding=1, bias=False
            ),
            nn.ReLU(True),
            nn.Conv2d(feature_dim, gaussian_raw_channels, kernel_size=1),
        )

    def map_pdf_to_opacity(
        self,
        pdf: Float[Tensor, " *batch"],
        global_step: int,
    ) -> Float[Tensor, " *batch"]:
        # https://www.desmos.com/calculator/opvwti3ba9

        # Figure out the exponent.
        cfg = self.cfg.opacity_mapping
        x = cfg.initial + min(global_step / cfg.warm_up, 1) * (cfg.final - cfg.initial)
        exponent = 2**x

        # Map the probability density to an opacity.
        return 0.5 * (1 - (1 - pdf) ** exponent + pdf ** (1 / exponent))

    def forward(
        self,
        image: torch.Tensor,
        global_step: int = 0,
    ) -> Gaussians:
        device = image.device
        b, v, _, h, w = image.shape
        if b != 1:
            raise ValueError(
                "EncoderOF3GS currently expects batch size 1 because voxel pruning "
                "can produce a different number of Gaussians per sample. Set "
                "data_loader.train.batch_size=1."
            )
        if v == 1:
            ctx_img_num = 1
        else:
            ctx_img_num = int(v * 0.5)
        ctx_img = image[:, :ctx_img_num, ...]
        distill_infos = {}
        pred_all_extrinsic = None
        pred_all_intrinsic = None

        if self.distill:
            distill_image = image.clone().detach()
            with torch.no_grad():
                with torch.amp.autocast("cuda", enabled=True, dtype=torch.float16):
                    distill_aggregated_tokens_list, _ = self.distill_aggregator(
                        distill_image.to(torch.float16)
                    )
                distill_aggregated_tokens_list = [
                    token.float() for token in distill_aggregated_tokens_list
                ]
                with torch.amp.autocast("cuda", enabled=False):
                    distill_pred_pose_enc_list = self.distill_camera_head(
                        distill_aggregated_tokens_list
                    )
                    last_distill_pred_pose_enc = distill_pred_pose_enc_list[-1]

                    pred_all_extrinsic, pred_all_intrinsic = (
                        pose_encoding_to_extri_intri(
                            last_distill_pred_pose_enc, image.shape[-2:]
                        )
                    )
                    extrinsic_padding = (
                        torch.tensor(
                            [0, 0, 0, 1],
                            device=pred_all_extrinsic.device,
                            dtype=pred_all_extrinsic.dtype,
                        )
                        .view(1, 1, 1, 4)
                        .repeat(b, distill_image.shape[1], 1, 1)
                    )
                    pred_all_extrinsic = torch.cat(
                        [pred_all_extrinsic, extrinsic_padding], dim=2
                    ).inverse()
                distill_infos["pred_pose_enc_list"] = last_distill_pred_pose_enc
                torch.cuda.empty_cache()

        image = ctx_img
        b, v, _, h, w = image.shape
        with torch.amp.autocast("cuda", enabled=True, dtype=torch.float16):
            aggregated_tokens_list, patch_start_idx = self.aggregator(
                image.to(torch.float16)
            )

        with torch.amp.autocast("cuda", enabled=False):
            pred_pose_enc_list = self.camera_head(aggregated_tokens_list)
            last_pred_pose_enc = pred_pose_enc_list[-1]

            pred_pose_enc_list = self.cam_dec(aggregated_tokens_list[-1][:, :, 0])

            extrinsic, intrinsic = pose_encoding_to_extri_intri(
                torch.cat(
                    (last_pred_pose_enc[..., :-2], pred_pose_enc_list.repeat(1, v, 2)),
                    dim=-1,
                ),
                image.shape[-2:],
            )
            gt_ex = closed_form_inverse_se3(extrinsic.flatten(0, 1))[..., :3, :]
            K = torch.zeros(
                (b, v, 3, 3), device=extrinsic.device, dtype=extrinsic.dtype
            )
            f_val = pred_all_intrinsic[0, 0, 0, 0]
            K[:, :, 0, 0] = f_val
            K[:, :, 1, 1] = f_val
            K[:, :, 0, 2] = intrinsic[..., 0, 2]
            K[:, :, 1, 2] = intrinsic[..., 1, 2]
            K[:, :, 2, 2] = 1.0
            gt_ix = K

            out, gs_depth, gs_depth_conf = self.gaussian_param_head(
                aggregated_tokens_list,
                image,
                image,
                patch_start_idx=patch_start_idx,
                image_size=(h, w),
            )
        del aggregated_tokens_list, patch_start_idx
        torch.cuda.empty_cache()

        depth_map = gs_depth
        depth_conf = gs_depth_conf
        conf_valid = torch.quantile(depth_conf.flatten(0, 1), self.cfg.conf_threshold)
        conf_valid_mask = depth_conf > conf_valid
        distill_infos["conf_mask"] = conf_valid_mask

        from . import act_gs, sh_utils

        gs_feats_reshape = rearrange(out, "b s c h w -> (b s) c h w")
        with torch.amp.autocast("cuda", enabled=False):
            gs_params = self.gs_head(gs_feats_reshape)
        splats = {}
        gs_params = rearrange(gs_params, "(b s) c h w -> b s h w c", b=b, s=v)
        quats, scales, opacities, residual_sh, weights, offsets = torch.split(
            gs_params, [4, 3, 1, self.nums_sh * 3, 1, 3], dim=-1
        )

        pts_all = batchify_unproject_depth_map_to_point_map(
            depth_map, gt_ex.detach(), gt_ix.detach()
        )

        offsets = offsets.reshape(b, v * h * w, 3)
        splats["offsets"] = offsets
        splats["quats"] = act_gs.reg_dense_rotation(quats.reshape(b, v * h * w, 4))
        splats["scales"] = act_gs.reg_dense_scales(
            scales.reshape(b, v * h * w, 3)
        ).clamp_max(0.3)
        densities = act_gs.reg_dense_opacities(opacities.reshape(b, v * h * w))
        splats["opacities"] = self.map_pdf_to_opacity(densities, global_step)
        residual_sh = act_gs.reg_dense_sh(
            residual_sh.reshape(b, v * h * w, self.nums_sh * 3)
        )
        new_sh = torch.zeros_like(residual_sh)
        new_sh[..., 0, :] = sh_utils.RGB2SH(
            image.permute(0, 1, 3, 4, 2).reshape(b, v * h * w, 3)
        )
        splats["sh"] = new_sh + residual_sh
        splats["residual_sh"] = residual_sh
        splats["weights"] = weights.reshape(b, v * h * w)
        means = pts_all.flatten(1, 3)
        splats["means"] = means + offsets

        def prune_gs(splats, voxel_size=0.002):
            B = splats["means"].shape[0]
            merged_splats_list = []
            device = splats["means"].device
            output = {}

            for i in range(B):
                splats_i = {
                    k: splats[k][i]
                    for k in ["means", "quats", "scales", "opacities", "sh", "weights"]
                }
                coords = splats_i["means"]
                voxel_indices = (coords / voxel_size).round().int()
                unique_voxels, inverse_indices, counts = torch.unique(
                    voxel_indices, dim=0, return_inverse=True, return_counts=True
                )
                K = len(unique_voxels)
                conf_flat = splats_i["weights"].flatten()  # [N]
                conf_voxel_max, _ = scatter_max(conf_flat, inverse_indices, dim=0)

                conf_exp = torch.exp(conf_flat - conf_voxel_max[inverse_indices])
                voxel_weights_sum = scatter_add(conf_exp, inverse_indices, dim=0)
                softmax_alpha = conf_exp / (voxel_weights_sum[inverse_indices] + 1e-6)
                alpha_mask = softmax_alpha.unsqueeze(-1)
                merged = {
                    "means": torch.zeros((K, 3), device=device),
                    "quats": torch.zeros((K, 4), device=device),
                    "scales": torch.zeros((K, 3), device=device),
                    "opacities": torch.zeros(K, device=device),
                    "sh": torch.zeros((K, self.nums_sh, 3), device=device),
                }

                # Means
                weighted_means = splats_i["means"] * alpha_mask
                merged["means"] = scatter_add(weighted_means, inverse_indices, dim=0)

                # SH
                weighted_sh = splats_i["sh"] * alpha_mask.unsqueeze(-1)
                merged["sh"] = scatter_add(weighted_sh, inverse_indices, dim=0)

                # Opacity
                weighted_opacities = splats_i["opacities"].flatten() * softmax_alpha
                merged["opacities"] = scatter_add(
                    weighted_opacities, inverse_indices, dim=0
                )

                # Scales
                weighted_scales = splats_i["scales"] * alpha_mask
                merged["scales"] = scatter_add(weighted_scales, inverse_indices, dim=0)

                # Quaternions
                weighted_quats = splats_i["quats"] * alpha_mask
                merged["quats"] = scatter_add(weighted_quats, inverse_indices, dim=0)
                quat_norms = torch.norm(merged["quats"], dim=1, keepdim=True)
                merged["quats"] = merged["quats"] / torch.clamp(quat_norms, min=1e-8)

                merged_splats_list.append(merged)

            for key in ["means", "sh", "opacities", "scales", "quats"]:
                output[key] = [merged[key] for merged in merged_splats_list]

            return output

        for k in splats:
            splats[k] = splats[k].float()

        splats = prune_gs(splats, voxel_size=self.voxel_size)

        key_mapping = {"quats": "rotations", "sh": "harmonics"}
        gaussians = {
            key_mapping.get(k, k): v[0].unsqueeze(0) for k, v in splats.items()
        }
        gaussians = Gaussians(**gaussians)

        extrinsic_padding = (
            torch.tensor([0, 0, 0, 1], device=device, dtype=extrinsic.dtype)
            .view(1, 1, 1, 4)
            .repeat(b, v, 1, 1)
        )
        intrinsic = gt_ix.clone()
        intrinsic = torch.stack(
            [intrinsic[:, :, 0] / w, intrinsic[:, :, 1] / h, intrinsic[:, :, 2]], dim=2
        )
        return (
            EncoderOutput(
                gaussians=gaussians,
                pred_pose_enc_list=pred_pose_enc_list,
                pred_context_pose=dict(
                    extrinsic=torch.cat(
                        [extrinsic, extrinsic_padding], dim=2
                    ).inverse(),
                    intrinsic=intrinsic,
                ),
                depth_dict=dict(depth=depth_map, conf_valid_mask=conf_valid_mask),
                distill_infos=distill_infos,
            ),
            pred_all_extrinsic,
            ctx_img_num,
        )

    def inference(
        self,
        image: torch.Tensor,
        tgt_image: torch.Tensor,
        global_step: int = 0,
    ) -> Gaussians:
        device = image.device
        b, v, _, h, w = image.shape
        tv = v

        distill_image = image.clone().detach()
        tgt_image = tgt_image.clone().detach()

        distill_image = torch.cat((distill_image, tgt_image), dim=1).to(torch.float16)

        all_pose = []
        past_key_values = [None] * self.aggregator.depth
        past_key_values_camera = [None] * self.camera_head.trunk_depth
        for i in range(distill_image.shape[1]):
            img = distill_image[:, i : i + 1]
            with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
                aggregated_tokens_list, patch_start_idx, past_key_values = (
                    self.aggregator(
                        img.to(torch.bfloat16),
                        past_key_values=past_key_values,
                        use_cache=True,
                        past_frame_idx=i,
                    )
                )

            with torch.amp.autocast("cuda", enabled=False):
                pred_pose_enc_list, past_key_values_camera = self.camera_head(
                    aggregated_tokens_list,
                    past_key_values_camera=past_key_values_camera,
                    use_cache=True,
                )
                last_pred_pose_enc = pred_pose_enc_list[-1]
                pred_pose_enc_list = self.cam_dec(aggregated_tokens_list[-1][:, :, 0])
                extrinsic, intrinsic = pose_encoding_to_extri_intri(
                    torch.cat(
                        (
                            last_pred_pose_enc[..., :-2],
                            pred_pose_enc_list.repeat(1, 1, 2),
                        ),
                        dim=-1,
                    ),
                    image.shape[-2:],
                )
            all_pose.append({"pose": extrinsic})

        extrinsics = []
        for pose in all_pose:
            extrinsics.append(pose["pose"])

        extrinsic_seq = torch.cat(extrinsics, dim=1)
        b_size = extrinsic_seq.shape[0]
        v_size = extrinsic_seq.shape[1]
        extrinsic_padding = (
            torch.tensor([0, 0, 0, 1], device=device, dtype=extrinsic_seq.dtype)
            .view(1, 1, 1, 4)
            .repeat(b_size, v_size, 1, 1)
        )
        pred_all_extrinsic = torch.cat(
            [extrinsic_seq, extrinsic_padding], dim=2
        ).inverse()
        del extrinsics
        torch.cuda.empty_cache()

        past_merged = {
            "means": None,
            "sum_w": None,
            "sum_w2": None,
            "sum_mw": None,
            "sum_shw": None,
            "sum_sw": None,
            "sum_qw": None,
        }
        all_ress = []

        past_key_values = [None] * self.aggregator.depth
        past_key_values_camera = [None] * self.camera_head.trunk_depth
        for i in range(image.shape[1]):
            img = image[:, i : i + 1]
            with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
                aggregated_tokens_list, patch_start_idx, past_key_values = (
                    self.aggregator(
                        img.to(torch.bfloat16),
                        past_key_values=past_key_values,
                        use_cache=True,
                        past_frame_idx=i,
                    )
                )

            with torch.amp.autocast("cuda", enabled=False):
                pred_pose_enc_list, past_key_values_camera = self.camera_head(
                    aggregated_tokens_list,
                    past_key_values_camera=past_key_values_camera,
                    use_cache=True,
                )
                last_pred_pose_enc = pred_pose_enc_list[-1]

                pred_pose_enc_list = self.cam_dec(aggregated_tokens_list[-1][:, :, 0])
                extrinsic, intrinsic = pose_encoding_to_extri_intri(
                    torch.cat(
                        (
                            last_pred_pose_enc[..., :-2],
                            pred_pose_enc_list.repeat(1, 1, 2),
                        ),
                        dim=-1,
                    ),
                    image.shape[-2:],
                )
                gt_ex = closed_form_inverse_se3(extrinsic.flatten(0, 1))
                if i == 0:
                    gt_ix = intrinsic
                    gt_ix[0, 0, 1, 1] = gt_ix[0, 0, 0, 0]

            with torch.amp.autocast("cuda", enabled=False):
                out, gs_depth, gs_depth_conf = self.gaussian_param_head(
                    aggregated_tokens_list,
                    img,
                    img,
                    patch_start_idx=patch_start_idx,
                    image_size=(h, w),
                )
            depth_map = gs_depth
            pts_all = batchify_unproject_depth_map_to_point_map(depth_map, gt_ex, gt_ix)
            b = 1
            v = 1
            from . import act_gs, sh_utils

            gs_feats_reshape = rearrange(out, "b s c h w -> (b s) c h w")
            with torch.amp.autocast("cuda", enabled=False):
                gs_params = self.gs_head(gs_feats_reshape)
            splats = {}
            gs_params = rearrange(gs_params, "(b s) c h w -> b s h w c", b=b, s=v)
            splats["gs_feats"] = gs_params.reshape(b, v * h * w, -1)
            quats, scales, opacities, residual_sh, weights, offsets = torch.split(
                gs_params, [4, 3, 1, self.nums_sh * 3, 1, 3], dim=-1
            )

            torch.cuda.empty_cache()

            offsets = offsets.reshape(b, v * h * w, 3)
            splats["offsets"] = offsets
            splats["quats"] = act_gs.reg_dense_rotation(quats.reshape(b, v * h * w, 4))
            splats["scales"] = act_gs.reg_dense_scales(
                scales.reshape(b, v * h * w, 3)
            ).clamp_max(0.3)
            densities = act_gs.reg_dense_opacities(opacities.reshape(b, v * h * w))
            splats["opacities"] = self.map_pdf_to_opacity(densities, global_step)
            residual_sh = act_gs.reg_dense_sh(
                residual_sh.reshape(b, v * h * w, self.nums_sh * 3)
            )
            new_sh = torch.zeros_like(residual_sh)
            new_sh[..., 0, :] = sh_utils.RGB2SH(
                img.permute(0, 1, 3, 4, 2).reshape(b, v * h * w, 3)
            )
            splats["sh"] = new_sh + residual_sh
            splats["residual_sh"] = residual_sh
            splats["weights"] = weights.reshape(b, v * h * w)
            means = pts_all.flatten(1, 3)
            splats["means"] = means + offsets

            def prune_gs(splats, voxel_size=0.002, past=None):
                B = splats["means"].shape[0]
                merged_splats_list = []
                device = splats["means"].device
                output = {}

                for i in range(B):
                    curr_conf = splats["weights"][i].flatten()
                    curr_means = splats["means"][i]
                    curr_sh = splats["sh"][i]
                    curr_opacities = splats["opacities"][i].flatten()
                    curr_scales = splats["scales"][i]
                    curr_quats = splats["quats"][i]
                    if past is not None and past.get("means") is not None:
                        all_means_to_index = torch.cat(
                            [curr_means, past["means"]], dim=0
                        )
                        voxel_indices = (all_means_to_index / voxel_size).round().int()
                        unique_voxels, inverse_indices = torch.unique(
                            voxel_indices, dim=0, return_inverse=True
                        )
                        K_combined = len(unique_voxels)

                        curr_inv_idx = inverse_indices[: len(curr_means)]
                        past_inv_idx = inverse_indices[len(curr_means) :]

                        m_old = torch.full((K_combined,), -float("inf"), device=device)
                        e_old = torch.zeros((K_combined,), device=device)

                        m_old.scatter_(0, past_inv_idx, past["max_logit"])
                        e_old.scatter_(0, past_inv_idx, past["sum_ew"])

                        def remap_acc(past_acc, p_idx, K):
                            out = torch.zeros((K,) + past_acc.shape[1:], device=device)
                            view_shape = (p_idx.shape[0],) + (1,) * (past_acc.ndim - 1)
                            out.scatter_add_(
                                0, p_idx.view(view_shape).expand_as(past_acc), past_acc
                            )
                            return out

                        s_mw_old = remap_acc(past["sum_mw"], past_inv_idx, K_combined)
                        s_shw_old = remap_acc(past["sum_shw"], past_inv_idx, K_combined)
                        s_ow_old = remap_acc(past["sum_ow"], past_inv_idx, K_combined)
                        s_sw_old = remap_acc(past["sum_sw"], past_inv_idx, K_combined)
                        s_qw_old = remap_acc(past["sum_qw"], past_inv_idx, K_combined)

                        curr_max = torch.full(
                            (K_combined,), -float("inf"), device=device
                        )
                        curr_max = curr_max.scatter_reduce(
                            0,
                            curr_inv_idx,
                            curr_conf,
                            reduce="amax",
                            include_self=False,
                        )

                        m_new = torch.max(m_old, curr_max)
                        m_new[m_new == -float("inf")] = 0

                        scale_past = torch.exp(m_old - m_new)
                        scale_past[torch.isnan(scale_past)] = 0

                        exp_curr = torch.exp(curr_conf - m_new[curr_inv_idx])

                        curr_sum_ew = scatter_add(
                            exp_curr, curr_inv_idx, dim=0, dim_size=K_combined
                        )
                        new_sum_ew = (e_old * scale_past) + curr_sum_ew

                        def update_acc(old_acc, curr_val, exp_w, inv_idx, s_past, K):
                            v_dim = curr_val.ndim
                            s_mask = s_past.view(-1, *([1] * (v_dim - 1)))
                            e_mask = exp_w.view(-1, *([1] * (v_dim - 1)))
                            curr_term = scatter_add(
                                curr_val * e_mask, inv_idx, dim=0, dim_size=K
                            )
                            return (old_acc * s_mask) + curr_term

                        new_sum_mw = update_acc(
                            s_mw_old,
                            curr_means,
                            exp_curr,
                            curr_inv_idx,
                            scale_past,
                            K_combined,
                        )
                        new_sum_shw = update_acc(
                            s_shw_old,
                            curr_sh,
                            exp_curr,
                            curr_inv_idx,
                            scale_past,
                            K_combined,
                        )
                        new_sum_ow = update_acc(
                            s_ow_old,
                            curr_opacities.unsqueeze(-1),
                            exp_curr,
                            curr_inv_idx,
                            scale_past,
                            K_combined,
                        )
                        new_sum_sw = update_acc(
                            s_sw_old,
                            curr_scales,
                            exp_curr,
                            curr_inv_idx,
                            scale_past,
                            K_combined,
                        )
                        new_sum_qw = update_acc(
                            s_qw_old,
                            curr_quats,
                            exp_curr,
                            curr_inv_idx,
                            scale_past,
                            K_combined,
                        )

                    else:
                        voxel_indices = (curr_means / voxel_size).round().int()
                        unique_voxels, curr_inv_idx = torch.unique(
                            voxel_indices, dim=0, return_inverse=True
                        )
                        K_combined = len(unique_voxels)

                        m_new, _ = scatter_max(curr_conf, curr_inv_idx, dim=0)
                        exp_curr = torch.exp(curr_conf - m_new[curr_inv_idx])
                        new_sum_ew = scatter_add(exp_curr, curr_inv_idx, dim=0)

                        e_mask = exp_curr.view(-1, 1)
                        new_sum_mw = scatter_add(
                            curr_means * e_mask, curr_inv_idx, dim=0
                        )
                        new_sum_shw = scatter_add(
                            curr_sh * e_mask.unsqueeze(-1), curr_inv_idx, dim=0
                        )
                        new_sum_ow = scatter_add(
                            curr_opacities.unsqueeze(-1) * e_mask, curr_inv_idx, dim=0
                        )
                        new_sum_sw = scatter_add(
                            curr_scales * e_mask, curr_inv_idx, dim=0
                        )
                        new_sum_qw = scatter_add(
                            curr_quats * e_mask, curr_inv_idx, dim=0
                        )

                    denom_safe = (new_sum_ew + 1e-6).view(-1, 1)

                    merged = {
                        "means": new_sum_mw / denom_safe,
                        "scales": new_sum_sw / denom_safe,
                        "sh": new_sum_shw / denom_safe.unsqueeze(-1),
                        "opacities": (new_sum_ow / denom_safe).squeeze(-1),
                        "quats": F.normalize(new_sum_qw / denom_safe, p=2, dim=-1),
                        "_acc": {
                            "max_logit": m_new,
                            "sum_ew": new_sum_ew,
                            "sum_mw": new_sum_mw,
                            "sum_shw": new_sum_shw,
                            "sum_ow": new_sum_ow,
                            "sum_sw": new_sum_sw,
                            "sum_qw": new_sum_qw,
                        },
                    }
                    merged_splats_list.append(merged)

                output = {
                    key: [m[key] for m in merged_splats_list]
                    for key in ["means", "sh", "opacities", "scales", "quats"]
                }
                return output, merged_splats_list[-1]["_acc"]

            splats, acc = prune_gs(splats, voxel_size=self.voxel_size, past=past_merged)
            past_merged["means"] = splats["means"][0]
            past_merged["quats"] = splats["quats"][0]
            past_merged["scales"] = splats["scales"][0]
            past_merged["opacities"] = splats["opacities"][0]
            past_merged["sh"] = splats["sh"][0]

            past_merged["max_logit"] = acc["max_logit"]  # M_new
            past_merged["sum_ew"] = acc["sum_ew"]  # E_new
            past_merged["sum_mw"] = acc["sum_mw"]  # S_means_new
            past_merged["sum_shw"] = acc["sum_shw"]  # S_sh_new
            past_merged["sum_ow"] = acc["sum_ow"]  # S_opacity_new
            past_merged["sum_sw"] = acc["sum_sw"]  # S_scales_new
            past_merged["sum_qw"] = acc["sum_qw"]  # S_quats_new

            past_merged["weights"] = acc["max_logit"]

            key_mapping = {"quats": "rotations", "sh": "harmonics"}
            gaussians = {
                key_mapping.get(k, k): v[0].unsqueeze(0) for k, v in splats.items()
            }
            gaussians = Gaussians(**gaussians)

            all_ress.append({"camera_pose": extrinsic, "in": gt_ix})

        v = tv
        extrinsic = []
        intrinsic = []
        for ress in all_ress:
            extrinsic.append(ress["camera_pose"])
            intrinsic.append(ress["in"])
        extrinsic = torch.cat(extrinsic, dim=1)
        intrinsic = torch.cat(intrinsic, dim=1)
        extrinsic_padding = (
            torch.tensor([0, 0, 0, 1], device=device, dtype=extrinsic.dtype)
            .view(1, 1, 1, 4)
            .repeat(b, v, 1, 1)
        )
        intrinsic = intrinsic.clone()
        intrinsic = torch.stack(
            [intrinsic[:, :, 0] / w, intrinsic[:, :, 1] / h, intrinsic[:, :, 2]], dim=2
        )
        pred_context_pose = dict(
            extrinsic=torch.cat([extrinsic, extrinsic_padding], dim=2).inverse(),
            intrinsic=intrinsic,
        )

        return (gaussians, pred_all_extrinsic, pred_context_pose)

    def get_data_shim(self) -> DataShim:
        def data_shim(batch: BatchedExample) -> BatchedExample:
            batch = apply_normalize_shim(
                batch,
                self.cfg.input_mean,
                self.cfg.input_std,
            )

            return batch

        return data_shim