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