import io import time import zipfile import zlib from dataclasses import dataclass, field from typing import Literal, Optional, List import numpy as np import torch from einops import rearrange, repeat from torch import nn from ...dataset.shims.patch_shim import apply_patch_shim from ...dataset.types import BatchedExample, DataShim from ...geometry.projection import sample_image_grid from ..types import Gaussians from .common.depth_guided_mv_refiner import DepthGuidedMVRefiner from .common.gaussian_adapter import GaussianAdapter, GaussianAdapterCfg from .common.gs_decoder import GaussianLatentDecoder, GaussianLatentDepthDecoder from .encoder import Encoder from .vae import coding from .vae import VariableRateLossyVAE from .vae.config import Config as VaeConfig import torch.nn.functional as F from numpy.lib.format import write_array from .unimatch.mv_unimatch import MultiViewUniMatch from .unimatch.dpt_head import DPTHead @dataclass class EncoderCodecSplatCfg: name: Literal["codecsplat"] num_depth_candidates: int num_surfaces: int gaussian_adapter: GaussianAdapterCfg downscale_factor: int shim_patch_size: int # mv_unimatch num_scales: int upsample_factor: int lowest_feature_resolution: int depth_unet_channels: int grid_sample_disable_cudnn: bool # color branch init_sh_input_img: bool gaussian_regressor_channels: int # loss config supervise_intermediate_depth: bool return_depth: bool # monodepth config monodepth_vit_type: str # multi-view matching local_mv_match: int # feature codec path enable_feature_codec: bool = False enable_depth_guided_mv_refiner: bool = False enable_post_vae_mv_refiner: bool = False codec_loss_weight: float = 1.0 codec_depth_loss_weight: float = 0.1 codec_lmb_range: List[float] = field(default_factory=lambda: [16.0, 1024.0]) freeze_legacy_modules_for_codec: bool = False codec_estimated_feature_num_bytes: bool = False # Test-only ablation: entropy-code the stage-1 latent feature/depth tensors # with plain deflate instead of using the learned VAE codec. deflate_feature_ablation: bool = False deflate_feature_dtype: Literal["float16", "float32"] = "float16" deflate_feature_level: int = 9 class EncoderCodecSplat(Encoder[EncoderCodecSplatCfg]): def __init__(self, cfg: EncoderCodecSplatCfg) -> None: super().__init__(cfg) if cfg.enable_feature_codec: assert not cfg.init_sh_input_img, ( "Feature codec mode requires init_sh_input_img=False because " "decoder-side reconstruction cannot access original input images." ) if cfg.deflate_feature_ablation: assert not cfg.enable_feature_codec, ( "deflate_feature_ablation is a non-learned baseline and should be " "run with enable_feature_codec=False." ) self.depth_predictor = MultiViewUniMatch( num_scales=cfg.num_scales, upsample_factor=cfg.upsample_factor, lowest_feature_resolution=cfg.lowest_feature_resolution, num_depth_candidates=cfg.num_depth_candidates, vit_type=cfg.monodepth_vit_type, unet_channels=cfg.depth_unet_channels, grid_sample_disable_cudnn=cfg.grid_sample_disable_cudnn, ) self.mv_refiner = None self.post_mv_refiner = None self.feature_codec = None self.gs_decoder = None self.codec_depth_head = None # upsample features to the original resolution model_configs = { 'vits': {'in_channels': 384, 'features': 64, 'out_channels': [48, 96, 192, 384]}, 'vitb': {'in_channels': 768, 'features': 96, 'out_channels': [96, 192, 384, 768]}, 'vitl': {'in_channels': 1024, 'features': 128, 'out_channels': [128, 256, 512, 1024]}, } self.feature_upsampler = DPTHead(**model_configs[cfg.monodepth_vit_type], downsample_factor=cfg.upsample_factor, return_feature=True, num_scales=cfg.num_scales, ) feature_upsampler_channels = model_configs[cfg.monodepth_vit_type]["features"] # gaussians adapter self.gaussian_adapter = GaussianAdapter(cfg.gaussian_adapter) # If a pre-refiner is enabled, let it handle explicit depth conditioning and keep # the gaussian_regressor input focused on image / match / feature cues. self.gaussian_regressor_uses_depth = not cfg.enable_depth_guided_mv_refiner gaussian_regressor_in_ch = 3 + 1 + feature_upsampler_channels if self.gaussian_regressor_uses_depth: gaussian_regressor_in_ch += 1 channels = self.cfg.gaussian_regressor_channels self.gaussian_regressor = nn.Sequential( nn.Conv2d(gaussian_regressor_in_ch, channels, 3, 1, 1), nn.GELU(), nn.Conv2d(channels, channels, 3, 1, 1), ) # predict gaussian parameters: scale, q, sh, offset, opacity num_gaussian_parameters = self.gaussian_adapter.d_in + 2 + 1 self.num_gaussian_parameters = num_gaussian_parameters self.gaussian_regressor_out_channels = channels if self.cfg.enable_depth_guided_mv_refiner: self.mv_refiner = DepthGuidedMVRefiner( feature_channels=channels, hidden_channels=channels, concat_ref_depth=True, ) if self.cfg.enable_feature_codec: codec_lmb_range = tuple(float(x) for x in self.cfg.codec_lmb_range) self.feature_codec = VariableRateLossyVAE( channels, VaeConfig( lmb_range=codec_lmb_range, lmb=float(codec_lmb_range[-1]), ), ) self.codec_depth_head = GaussianLatentDepthDecoder( latent_channels=channels, hidden_channels=channels, ) if self.cfg.enable_post_vae_mv_refiner: self.post_mv_refiner = DepthGuidedMVRefiner( feature_channels=channels, hidden_channels=channels, ) self.gs_decoder = GaussianLatentDecoder( latent_channels=channels, hidden_channels=channels, num_gaussian_parameters=num_gaussian_parameters, init_sh_input_img=self.cfg.init_sh_input_img, ) if self.cfg.enable_feature_codec: if self.cfg.freeze_legacy_modules_for_codec: self.freeze_legacy_modules_for_codec() def freeze_legacy_modules_for_codec(self) -> None: for param in self.parameters(): param.requires_grad_(False) modules_to_train = [ self.gaussian_regressor, self.mv_refiner, self.feature_codec, self.codec_depth_head, self.post_mv_refiner, self.gs_decoder, ] for module in modules_to_train: if module is None: continue for param in module.parameters(): param.requires_grad_(True) def forward( self, context: dict, global_step: int, visualization_dump: Optional[dict] = None, scene_names: Optional[list] = None, return_bitstream: bool = False, ): if self.training: return self.normal_forward( context, global_step, visualization_dump=visualization_dump, scene_names=scene_names, ) if self.cfg.deflate_feature_ablation: return self.deflate_feature_forward( context, global_step, visualization_dump=visualization_dump, scene_names=scene_names, ) if self.feature_codec is not None: return self.codec_forward( context, global_step, visualization_dump=visualization_dump, scene_names=scene_names, return_bitstream=return_bitstream, ) return self.normal_forward( context, global_step, visualization_dump=visualization_dump, scene_names=scene_names, ) def compute_view_neighbor_indices(self, extrinsics: torch.Tensor): v = extrinsics.shape[1] if v <= 3: return None with torch.no_grad(): xyzs = extrinsics[:, :, :3, -1].detach().float() cameras_dist_matrix = torch.cdist(xyzs, xyzs, p=2) cameras_dist_index = torch.argsort(cameras_dist_matrix) return cameras_dist_index[:, :, :(self.cfg.local_mv_match + 1)] def encode_camera_stream(self, context: dict) -> bytes: _, v, _, h, w = context["image"].shape camera_arrays = { "image_shape": np.array([h, w], dtype=np.int32), "num_context_views": np.array([v], dtype=np.int32), "intrinsics": context["intrinsics"][0].detach().cpu().float().numpy(), "extrinsics": context["extrinsics"][0].detach().cpu().float().numpy(), "near": context["near"][0].detach().cpu().float().numpy(), "far": context["far"][0].detach().cpu().float().numpy(), } camera_buffer = io.BytesIO() with zipfile.ZipFile( camera_buffer, mode="w", compression=zipfile.ZIP_DEFLATED, compresslevel=9, ) as camera_zip: for name, array in camera_arrays.items(): array_buffer = io.BytesIO() write_array(array_buffer, np.asanyarray(array), allow_pickle=False) camera_zip.writestr(f"{name}.npy", array_buffer.getvalue()) camera_stream = camera_buffer.getvalue() camera_buffer.close() return camera_stream @torch.no_grad() def compress(self, context: dict, features: torch.Tensor) -> tuple[bytes, dict]: assert self.feature_codec is not None b, v, _, h, w = context["image"].shape assert b == 1, "Scene bitstream export currently supports batch size 1." assert features.shape[0] == v, ( f"Expected one feature tensor per view, got {features.shape[0]=}, {v=}." ) assert v > 1, "Scene feature stream packing expects at least two context views." if features.is_cuda: torch.cuda.synchronize(features.device) start = time.perf_counter() camera_stream = self.encode_camera_stream(context) feature_streams = [ self.feature_codec.compress(features[i : i + 1]) for i in range(features.shape[0]) ] feature_stream = coding.concat_bytes_list(feature_streams) camera_num_bytes = len(camera_stream) assert camera_num_bytes < (1 << 24), ( f"Camera stream is too large for the 3-byte length header: {camera_num_bytes}" ) bytestream = ( camera_num_bytes.to_bytes(3, byteorder="big", signed=False) + camera_stream + feature_stream ) if features.is_cuda: torch.cuda.synchronize(features.device) compress_seconds = time.perf_counter() - start return bytestream, { "codec_compress_seconds": compress_seconds, "codec_camera_num_bytes": float(camera_num_bytes), "codec_feature_num_bytes": float(len(feature_stream)), "codec_num_bytes": float(len(bytestream)), } @torch.no_grad() def run_deflate_feature_ablation( self, context: dict, features: torch.Tensor, depth: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, dict]: assert not self.training, "deflate_feature_ablation is intended for evaluation." device = features.device if features.is_cuda: torch.cuda.synchronize(device) compress_start = time.perf_counter() dtype = self.cfg.deflate_feature_dtype np_dtype = np.float16 if dtype == "float16" else np.float32 level = int(self.cfg.deflate_feature_level) def tensor_to_npy_bytes(tensor: torch.Tensor) -> bytes: tensor_np = tensor.detach().cpu().to(torch.float32).numpy().astype( np_dtype, copy=False, ) tensor_buffer = io.BytesIO() write_array(tensor_buffer, np.asanyarray(tensor_np), allow_pickle=False) stream = tensor_buffer.getvalue() tensor_buffer.close() return stream camera_stream = self.encode_camera_stream(context) feature_raw_stream = tensor_to_npy_bytes(features) depth_raw_stream = tensor_to_npy_bytes(depth) feature_stream = zlib.compress(feature_raw_stream, level=level) depth_stream = zlib.compress(depth_raw_stream, level=level) if features.is_cuda: torch.cuda.synchronize(device) compress_seconds = time.perf_counter() - compress_start if features.is_cuda: torch.cuda.synchronize(device) decompress_start = time.perf_counter() features_hat = torch.from_numpy( np.load(io.BytesIO(zlib.decompress(feature_stream)), allow_pickle=False) ).to(device=device, dtype=features.dtype) depth_hat = torch.from_numpy( np.load(io.BytesIO(zlib.decompress(depth_stream)), allow_pickle=False) ).to(device=device, dtype=depth.dtype) if features.is_cuda: torch.cuda.synchronize(device) decompress_seconds = time.perf_counter() - decompress_start header_num_bytes = 24 # three uint64 lengths for camera/feature/depth streams codec_num_bytes = ( header_num_bytes + len(camera_stream) + len(feature_stream) + len(depth_stream) ) return features_hat, depth_hat, { "codec_compress_seconds": compress_seconds, "codec_decompress_seconds": decompress_seconds, "codec_camera_num_bytes": float(len(camera_stream)), "codec_feature_num_bytes": float(len(feature_stream) + len(depth_stream)), "codec_latent_feature_num_bytes": float(len(feature_stream)), "codec_depth_num_bytes": float(len(depth_stream)), "codec_feature_raw_num_bytes": float(len(feature_raw_stream)), "codec_depth_raw_num_bytes": float(len(depth_raw_stream)), "codec_num_bytes": float(codec_num_bytes), } @torch.no_grad() def decompress( self, bytestream: bytes, visualization_dump: Optional[dict] = None, ) -> tuple[Gaussians, Optional[torch.Tensor], dict, dict]: assert self.feature_codec is not None param = next(self.parameters()) device = param.device if param.is_cuda: torch.cuda.synchronize(device) start = time.perf_counter() camera_num_bytes = int.from_bytes( bytestream[:3], byteorder="big", signed=False ) camera_stream = bytestream[3 : 3 + camera_num_bytes] feature_stream = bytestream[3 + camera_num_bytes :] with np.load(io.BytesIO(camera_stream)) as camera_data: image_shape = tuple(int(x) for x in camera_data["image_shape"]) num_context_views = int(camera_data["num_context_views"][0]) intrinsics_np = camera_data["intrinsics"] extrinsics_np = camera_data["extrinsics"] near_np = camera_data["near"] far_np = camera_data["far"] feature_streams = coding.split_bytes_list(feature_stream, num_context_views) features_hat = torch.cat( [self.feature_codec.decompress(stream) for stream in feature_streams], dim=0, ) decoded_context = {} decoded_context["image_shape"] = image_shape decoded_context["intrinsics"] = torch.from_numpy(intrinsics_np).to( device=device, dtype=torch.float32, ).unsqueeze(0) decoded_context["extrinsics"] = torch.from_numpy(extrinsics_np).to( device=device, dtype=torch.float32, ).unsqueeze(0) decoded_context["near"] = torch.from_numpy(near_np).to( device=device, dtype=torch.float32, ).unsqueeze(0) decoded_context["far"] = torch.from_numpy(far_np).to( device=device, dtype=torch.float32, ).unsqueeze(0) neighbor_indices = self.compute_view_neighbor_indices( decoded_context["extrinsics"] ) decode_state = { "device": device, "b": 1, "v": num_context_views, "h": image_shape[0], "w": image_shape[1], "refiner_source_indices": ( neighbor_indices[..., 1:] if neighbor_indices is not None else None ), } gaussians_raw, depth, _, _ = self.decode_codec_features( features_hat, decoded_context, decode_state, codec_depth_targets=None, ) decode_state["depth_preds"] = [depth] gaussians, depths = self.adapt_gaussians( decoded_context, decoded_context, gaussians_raw, depth, decode_state, visualization_dump, ) if param.is_cuda: torch.cuda.synchronize(device) decompress_seconds = time.perf_counter() - start return gaussians, depths, decoded_context, { "codec_decompress_seconds": decompress_seconds, } def encode_context_features(self, context: dict): device = context["image"].device b, v, _, h, w = context["image"].shape cameras_dist_index = self.compute_view_neighbor_indices(context["extrinsics"]) refiner_source_indices = ( cameras_dist_index[..., 1:] if cameras_dist_index is not None else None ) results_dict = self.depth_predictor( context["image"], attn_splits_list=[2], min_depth=1. / context["far"], max_depth=1. / context["near"], intrinsics=context["intrinsics"], extrinsics=context["extrinsics"], nn_matrix=cameras_dist_index, supervise_intermediate_depth=self.cfg.supervise_intermediate_depth, ) depth_preds = results_dict['depth_preds'] depth = depth_preds[-1] features = self.feature_upsampler( results_dict["features_mono_intermediate"], cnn_features=results_dict["features_cnn_all_scales"][::-1], mv_features=results_dict["features_mv"][0] if self.cfg.num_scales == 1 else results_dict["features_mv"][::-1], ) match_prob = results_dict['match_probs'][-1] match_prob = torch.max(match_prob, dim=1, keepdim=True)[0] match_prob = F.interpolate(match_prob, size=depth.shape[-2:], mode='nearest') regressor_inputs = [ rearrange(context["image"], "b v c h w -> (b v) c h w"), match_prob, features, ] if self.mv_refiner is None: regressor_inputs.insert(1, rearrange(depth, "b v h w -> (b v) () h w")) out = self.gaussian_regressor(torch.cat(regressor_inputs, dim=1)) if self.mv_refiner is not None: out = rearrange(out, "(b v) c h w -> b v c h w", b=b, v=v) out = self.mv_refiner( out, depth, context["intrinsics"], context["extrinsics"], source_view_indices=refiner_source_indices, ) out = rearrange(out, "b v c h w -> (b v) c h w") return { "device": device, "b": b, "v": v, "h": h, "w": w, "depth_preds": depth_preds, "depth": depth, "teacher_depth": depth, "features": features, "match_prob": match_prob, "out": out, "refiner_source_indices": refiner_source_indices, } def decode_codec_features( self, out_hat: torch.Tensor, decode_context: dict, state: dict, codec_depth_targets: Optional[torch.Tensor] = None, ): device = state["device"] b = state["b"] v = state["v"] near_flat = rearrange(decode_context["near"], "b v -> (b v) () () ()") far_flat = rearrange(decode_context["far"], "b v -> (b v) () () ()") ( codec_depth_flat, _codec_inv_depth_flat, codec_norm_inv_depth_flat, ) = self.codec_depth_head( out_hat, near_flat, far_flat, return_inverse=True, return_normalized_inverse=True, ) codec_depth = rearrange(codec_depth_flat, "(b v) () h w -> b v h w", b=b, v=v) codec_depth_loss_pred = None codec_depth_loss_target = None if codec_depth_targets is not None: codec_norm_inv_depth = rearrange( codec_norm_inv_depth_flat, "(b v) () h w -> b v h w", b=b, v=v ) with torch.autocast(device_type=device.type, enabled=False): near = rearrange(decode_context["near"].float(), "b v -> b v () ()") far = rearrange(decode_context["far"].float(), "b v -> b v () ()") inv_min = far.clamp_min(torch.finfo(torch.float32).eps).reciprocal() inv_max = near.clamp_min(torch.finfo(torch.float32).eps).reciprocal() inv_range = (inv_max - inv_min).clamp_min(torch.finfo(torch.float32).eps) teacher_inv_depth = codec_depth_targets.float().clamp_min( torch.finfo(torch.float32).eps ).reciprocal() teacher_inv_depth = teacher_inv_depth.clamp(min=inv_min, max=inv_max) codec_depth_loss_target = ((teacher_inv_depth - inv_min) / inv_range).clamp(0.0, 1.0) codec_depth_loss_pred = codec_norm_inv_depth if self.post_mv_refiner is not None: out_hat = rearrange(out_hat, "(b v) c h w -> b v c h w", b=b, v=v) out_hat = self.post_mv_refiner( out_hat, codec_depth, decode_context["intrinsics"], decode_context["extrinsics"], source_view_indices=state["refiner_source_indices"], ) out_hat = rearrange(out_hat, "b v c h w -> (b v) c h w") depth_flat = rearrange(codec_depth, "b v h w -> (b v) () h w") assert self.gs_decoder is not None gaussians = self.gs_decoder(out_hat, depth_flat) return gaussians, codec_depth, codec_depth_loss_pred, codec_depth_loss_target def adapt_gaussians( self, context: dict, adapter_context: dict, gaussians: torch.Tensor, depth: torch.Tensor, state: dict, visualization_dump: Optional[dict], ): device = state["device"] b = state["b"] v = state["v"] h = state["h"] w = state["w"] depth_preds = state["depth_preds"] gaussians = rearrange(gaussians, "(b v) c h w -> b v c h w", b=b, v=v) depths = rearrange(depth, "b v h w -> b v (h w) () ()") raw_gaussians = rearrange(gaussians, "b v c h w -> b v (h w) c") if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1: num_depths = len(depth_preds) intermediate_depths = torch.cat(depth_preds[:(num_depths - 1)], dim=0) intermediate_depths = rearrange( intermediate_depths, "b v h w -> b v (h w) () ()" ) depths = torch.cat((intermediate_depths, depths), dim=0) raw_gaussians = torch.cat([raw_gaussians] * num_depths, dim=0) b *= num_depths with torch.autocast(device_type=device.type, enabled=False): raw_gaussians = raw_gaussians.float() depths = depths.float() opacities = raw_gaussians[..., :1].sigmoid().unsqueeze(-1) raw_gaussians = raw_gaussians[..., 1:] xy_ray, _ = sample_image_grid((h, w), device) xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy") gaussians = rearrange( raw_gaussians, "... (srf c) -> ... srf c", srf=self.cfg.num_surfaces, ) offset_xy = gaussians[..., :2].sigmoid() pixel_size = 1 / torch.tensor((w, h), dtype=torch.float32, device=device) xy_ray = xy_ray + (offset_xy - 0.5) * pixel_size sh_input_images = context.get("image") if self.cfg.init_sh_input_img: assert sh_input_images is not None, ( "init_sh_input_img=True requires input images, but decoded " "bitstreams intentionally do not contain original RGB images." ) if self.cfg.supervise_intermediate_depth and len(depth_preds) > 1: context_extrinsics = torch.cat( [adapter_context["extrinsics"]] * len(depth_preds), dim=0 ) context_intrinsics = torch.cat( [adapter_context["intrinsics"]] * len(depth_preds), dim=0 ) gaussians = self.gaussian_adapter.forward( rearrange(context_extrinsics, "b v i j -> b v () () () i j"), rearrange(context_intrinsics, "b v i j -> b v () () () i j"), rearrange(xy_ray, "b v r srf xy -> b v r srf () xy"), depths, opacities, rearrange(gaussians[..., 2:], "b v r srf c -> b v r srf () c"), (h, w), input_images=sh_input_images.repeat(len(depth_preds), 1, 1, 1, 1) if self.cfg.init_sh_input_img else None, ) else: gaussians = self.gaussian_adapter.forward( rearrange(adapter_context["extrinsics"], "b v i j -> b v () () () i j"), rearrange(adapter_context["intrinsics"], "b v i j -> b v () () () i j"), rearrange(xy_ray, "b v r srf xy -> b v r srf () xy"), depths, opacities, rearrange(gaussians[..., 2:], "b v r srf c -> b v r srf () c"), (h, w), input_images=sh_input_images if self.cfg.init_sh_input_img else None, ) if visualization_dump is not None: visualization_dump["depth"] = rearrange( depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w ) visualization_dump["scales"] = rearrange( gaussians.scales, "b v r srf spp xyz -> b (v r srf spp) xyz" ) visualization_dump["rotations"] = rearrange( gaussians.rotations, "b v r srf spp xyzw -> b (v r srf spp) xyzw" ) gaussians = Gaussians( means=rearrange(gaussians.means, "b v r srf spp xyz -> b (v r srf spp) xyz"), covariances=rearrange( gaussians.covariances, "b v r srf spp i j -> b (v r srf spp) i j" ), harmonics=rearrange( gaussians.harmonics, "b v r srf spp c d_sh -> b (v r srf spp) c d_sh" ), opacities=rearrange(gaussians.opacities, "b v r srf spp -> b (v r srf spp)"), scales=rearrange(gaussians.scales, "b v r srf spp xyz -> b (v r srf spp) xyz"), rotations=rearrange( gaussians.rotations, "b v r srf spp xyzw -> b (v r srf spp) xyzw" ), rotations_unnorm=rearrange( gaussians.rotations_unnorm, "b v r srf spp xyzw -> b (v r srf spp) xyzw" ), ) if self.cfg.return_depth: depths = rearrange( depths, "b v (h w) srf s -> b v h w srf s", h=h, w=w ).squeeze(-1).squeeze(-1) return gaussians, depths @torch.no_grad() def codec_forward( self, context: dict, global_step: int, visualization_dump: Optional[dict] = None, scene_names: Optional[list] = None, return_bitstream: bool = False, ): b = context["image"].shape[0] assert b == 1, "Codec bitstream evaluation currently supports batch size 1." assert self.feature_codec is not None assert not self.training input_device = context["image"].device if context["image"].is_cuda: torch.cuda.synchronize(input_device) encode_context_start = time.perf_counter() state = self.encode_context_features(context) if context["image"].is_cuda: torch.cuda.synchronize(input_device) encode_context_seconds = time.perf_counter() - encode_context_start assert "depth_only_result" not in state, "Codec bitstream evaluation needs feature latents." estimated_feature_num_bytes = None if self.cfg.codec_estimated_feature_num_bytes: estimated_feature_bpp, _, _ = self.feature_codec.train_forward( state["out"], lmb=self.feature_codec.cfg.lmb, ) feature_height, feature_width = state["out"].shape[-2:] estimated_feature_num_bytes = ( estimated_feature_bpp.float().sum() * float(feature_height * feature_width) / 8.0 ) bytestream, compress_stats = self.compress(context, state["out"]) gaussians, depths, codec_decode_context, decompress_stats = self.decompress( bytestream, visualization_dump=visualization_dump, ) assert codec_decode_context["image_shape"] == (state["h"], state["w"]), ( f"Decoded image_shape={codec_decode_context['image_shape']} " f"does not match context image_shape={(state['h'], state['w'])}." ) assert codec_decode_context["intrinsics"].shape[1] == state["v"], ( "Decoded camera stream context view count does not match " f"the current context: {codec_decode_context['intrinsics'].shape[1]} vs {state['v']}." ) device = state["device"] codec_num_bytes = torch.tensor( compress_stats["codec_num_bytes"], device=device, dtype=torch.float32 ) total_pixels_per_scene = state["out"].shape[-2] * state["out"].shape[-1] result = { "gaussians": gaussians, "codec_num_bytes": codec_num_bytes, "codec_bpp": codec_num_bytes * 8.0 / float(total_pixels_per_scene), "codec_encode_context_seconds": torch.tensor( encode_context_seconds, device=device, dtype=torch.float32 ), "codec_compress_seconds": torch.tensor( compress_stats["codec_compress_seconds"], device=device, dtype=torch.float32 ), "codec_decompress_seconds": torch.tensor( decompress_stats["codec_decompress_seconds"], device=device, dtype=torch.float32 ), "codec_camera_num_bytes": torch.tensor( compress_stats["codec_camera_num_bytes"], device=device, dtype=torch.float32 ), "codec_feature_num_bytes": torch.tensor( compress_stats["codec_feature_num_bytes"], device=device, dtype=torch.float32 ), } if self.cfg.return_depth: result["depths"] = depths if estimated_feature_num_bytes is not None: result["codec_estimated_feature_num_bytes"] = ( estimated_feature_num_bytes ) if return_bitstream: result["bitstream"] = bytestream return result @torch.no_grad() def deflate_feature_forward( self, context: dict, global_step: int, visualization_dump: Optional[dict] = None, scene_names: Optional[list] = None, ): assert not self.training assert self.feature_codec is None assert context["image"].shape[0] == 1, ( "Deflate feature evaluation currently supports batch size 1." ) state = self.encode_context_features(context) assert "depth_only_result" not in state, ( "Deflate feature evaluation needs Gaussian-generation features." ) device = state["device"] out, depth, deflate_stats = self.run_deflate_feature_ablation( context, state["out"], state["depth"], ) state["depth"] = depth state["depth_preds"] = [*state["depth_preds"][:-1], depth] depth_flat = rearrange(depth, "b v h w -> (b v) () h w") assert self.gs_decoder is not None gaussians_raw = self.gs_decoder(out, depth_flat) gaussians, depths = self.adapt_gaussians( context, context, gaussians_raw, depth, state, visualization_dump, ) codec_num_bytes = torch.tensor( deflate_stats["codec_num_bytes"], device=device, dtype=torch.float32, ) total_pixels_per_scene = out.shape[-2] * out.shape[-1] result = { "gaussians": gaussians, "codec_num_bytes": codec_num_bytes, "codec_bpp": codec_num_bytes * 8.0 / float(total_pixels_per_scene), } result.update( { name: torch.tensor(value, device=device, dtype=torch.float32) for name, value in deflate_stats.items() if name != "codec_num_bytes" } ) if self.cfg.return_depth: result["depths"] = depths return result def normal_forward( self, context: dict, global_step: int, visualization_dump: Optional[dict] = None, scene_names: Optional[list] = None, ): state = self.encode_context_features(context) b = state["b"] v = state["v"] out = state["out"] depth = state["depth"] codec_loss = None codec_lmb = None codec_depth_targets = None codec_depth_loss_pred = None codec_depth_loss_target = None if self.feature_codec is not None: assert self.training, ( "Learned codec evaluation must use the complete bitstream path." ) codec_depth_targets = state["teacher_depth"].detach() scene_lmb = self.feature_codec.sample_lmb(b, out.device) view_lmb = repeat(scene_lmb, "b -> (b v)", v=v) codec_loss, out_hat, _ = self.feature_codec(out, lmb=view_lmb) codec_loss = rearrange(codec_loss, "(b v) -> b v", b=b, v=v) codec_loss = codec_loss.sum(dim=1) codec_lmb = scene_lmb gaussians_raw, depth, codec_depth_loss_pred, codec_depth_loss_target = ( self.decode_codec_features( out_hat, context, state, codec_depth_targets=codec_depth_targets, ) ) adapter_context = context else: depth_flat = rearrange(depth, "b v h w -> (b v) () h w") assert self.gs_decoder is not None gaussians_raw = self.gs_decoder(out, depth_flat) adapter_context = context gaussians, depths = self.adapt_gaussians( context, adapter_context, gaussians_raw, depth, state, visualization_dump, ) if self.cfg.return_depth: result = {"gaussians": gaussians, "depths": depths} if codec_depth_targets is not None: result["codec_depth_targets"] = codec_depth_targets if codec_depth_loss_pred is not None: result["codec_depth_loss_pred"] = codec_depth_loss_pred if codec_depth_loss_target is not None: result["codec_depth_loss_target"] = codec_depth_loss_target if codec_loss is not None: return result, codec_loss, codec_lmb return result if codec_loss is not None: return gaussians, codec_loss, codec_lmb return gaussians def get_data_shim(self) -> DataShim: def data_shim(batch: BatchedExample) -> BatchedExample: batch = apply_patch_shim( batch, patch_size=self.cfg.shim_patch_size * self.cfg.downscale_factor, ) return batch return data_shim