"""Camera-aware unprojection into an explicit typed voxel state.""" from __future__ import annotations import math from collections.abc import Sequence import torch import torch.nn.functional as functional from torch import Tensor, nn from mcss.geometry import make_voxel_centers, project_world from mcss.model.encoder2d import ImageEncoder from mcss.types import Cameras, SceneState class Residual3DBlock(nn.Module): def __init__(self, channels: int) -> None: super().__init__() groups = _group_count(channels) self.block = nn.Sequential( nn.Conv3d(channels, channels, 3, padding=1, bias=False), nn.GroupNorm(groups, channels), nn.SiLU(inplace=True), nn.Conv3d(channels, channels, 3, padding=1, bias=False), nn.GroupNorm(groups, channels), ) self.activation = nn.SiLU(inplace=True) def forward(self, inputs: Tensor) -> Tensor: return self.activation(inputs + self.block(inputs)) class UnprojectiveStateEncoder(nn.Module): """Fuse calibrated image observations at fixed world-space voxel centers.""" def __init__( self, *, voxel_resolution: int | Sequence[int], image_feature_dim: int, state_feature_dim: int, refinement_blocks: int, ) -> None: super().__init__() if isinstance(voxel_resolution, int): resolution = (voxel_resolution,) * 3 else: resolution = tuple(voxel_resolution) if len(resolution) != 3 or min(resolution) < 4: raise ValueError("voxel_resolution must contain D, H, W values of at least four") if state_feature_dim < 4: raise ValueError("state_feature_dim must be at least four") if refinement_blocks < 0: raise ValueError("refinement_blocks cannot be negative") self.voxel_resolution = resolution self.image_feature_dim = image_feature_dim self.state_feature_dim = state_feature_dim self.image_encoder = ImageEncoder(image_feature_dim) observation_channels = image_feature_dim + 3 + 3 + 1 fused_channels = observation_channels * 2 + 1 hidden_channels = max(16, state_feature_dim * 2) self.input_projection = nn.Sequential( nn.Conv3d(fused_channels, hidden_channels, 3, padding=1, bias=False), nn.GroupNorm(_group_count(hidden_channels), hidden_channels), nn.SiLU(inplace=True), ) self.refinement = nn.Sequential( *(Residual3DBlock(hidden_channels) for _ in range(refinement_blocks)) ) self.density_head = nn.Conv3d(hidden_channels, 1, 1) self.color_head = nn.Conv3d(hidden_channels, 3, 1) self.variance_head = nn.Conv3d(hidden_channels, 1, 1) self.feature_head = nn.Conv3d(hidden_channels, state_feature_dim, 1) nn.init.constant_(self.density_head.bias, -2.0) nn.init.constant_(self.variance_head.bias, -3.0) def forward(self, context_rgb: Tensor, cameras: Cameras, bounds: Tensor) -> SceneState: if context_rgb.ndim != 5 or context_rgb.shape[2] != 3: raise ValueError("context_rgb must have shape [B, V, 3, H, W]") batch_size, view_count, _, image_height, image_width = context_rgb.shape if cameras.leading_shape != (batch_size, view_count): raise ValueError("context cameras must align with context_rgb") if cameras.image_size != (image_height, image_width): raise ValueError("context camera image size must match context_rgb") if bounds.shape != (batch_size, 2, 3): raise ValueError("bounds must have shape [B, 2, 3]") encoded = self.image_encoder( context_rgb.reshape(batch_size * view_count, 3, image_height, image_width) ) _, channels, feature_height, feature_width = encoded.shape encoded = encoded.reshape(batch_size, view_count, channels, feature_height, feature_width) centers = make_voxel_centers(bounds, self.voxel_resolution) point_count = math.prod(self.voxel_resolution) flat_centers = centers.reshape(batch_size, point_count, 3) points_per_view = flat_centers[:, None].expand(batch_size, view_count, point_count, 3) with torch.autocast(device_type=context_rgb.device.type, enabled=False): pixels, camera_depth, valid = project_world( points_per_view.float(), cameras.to(dtype=torch.float32) ) normalized_grid = _pixels_to_grid(pixels, image_height, image_width) sampled_features = _sample_images(encoded, normalized_grid) sampled_rgb = _sample_images(context_rgb, normalized_grid) camera_centers = cameras.c2w[..., :3, 3].unsqueeze(-2) view_direction = _normalize(points_per_view - camera_centers) scene_diagonal = torch.linalg.vector_norm(bounds[:, 1] - bounds[:, 0], dim=-1) normalized_depth = camera_depth / scene_diagonal[:, None, None].clamp_min(1e-6) observations = torch.cat( (sampled_features, sampled_rgb, view_direction, normalized_depth.unsqueeze(-1)), dim=-1 ) mask = valid.unsqueeze(-1).to(observations.dtype) count = mask.sum(dim=1).clamp_min(1.0) mean = (observations * mask).sum(dim=1) / count variance = (((observations - mean[:, None]) ** 2) * mask).sum(dim=1) / count coverage = mask.sum(dim=1) / view_count fused = torch.cat((mean, variance, coverage), dim=-1) fused = fused.reshape( batch_size, *self.voxel_resolution, fused.shape[-1], ).permute(0, 4, 1, 2, 3) hidden = self.refinement(self.input_projection(fused)) density_logits = self.density_head(hidden) color = torch.sigmoid(self.color_head(hidden)) log_variance = self.variance_head(hidden).clamp(-8.0, 4.0) features = self.feature_head(hidden) return SceneState(density_logits, color, log_variance, bounds, features) def _sample_images(images: Tensor, grid: Tensor) -> Tensor: batch_size, views, channels, height, width = images.shape point_count = grid.shape[-2] sampled = functional.grid_sample( images.reshape(batch_size * views, channels, height, width), grid.reshape(batch_size * views, point_count, 1, 2), mode="bilinear", padding_mode="zeros", align_corners=True, ) return sampled.reshape(batch_size, views, channels, point_count).permute(0, 1, 3, 2) def _pixels_to_grid(pixels: Tensor, height: int, width: int) -> Tensor: x = pixels[..., 0] / max(width - 1, 1) * 2.0 - 1.0 y = pixels[..., 1] / max(height - 1, 1) * 2.0 - 1.0 return torch.stack((x, y), dim=-1) def _normalize(values: Tensor) -> Tensor: return values / torch.linalg.vector_norm(values, dim=-1, keepdim=True).clamp_min(1e-8) def _group_count(channels: int) -> int: for groups in (8, 4, 2): if channels % groups == 0: return groups return 1