Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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 | |
| 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)), | |
| } | |
| 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), | |
| } | |
| 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 | |
| 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 | |
| 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 | |