| """MapGS: the full feed-forward model (§2.2-§2.5). |
| |
| Pipeline: ViT encoder fuses the context views into image tokens; the |
| TokenManager builds map-anchored + free (+ dynamic) query tokens; the DETR |
| decoder cross-attends them to the image tokens and self-attends among them |
| (dynamic->static causal mask); the GaussianHead decodes each token into ``N_G`` |
| Gaussians. Dynamic tokens are canonicalized and mapped to the target |
| timestamp's world frame (:mod:`mapgs.model.dynamic`). Appearance is decoded from |
| rendered feature maps by a light UNet (or direct RGB). |
| |
| ``forward`` returns the scene Gaussians (a :class:`DecodedGaussians`); |
| rasterization + appearance decoding are driven by the trainer/losses so that the |
| extra extrapolation renders (§2.6) reuse the same Gaussians. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from mapgs.config import MapGSConfig |
| from mapgs.model.encoder import ViTEncoder |
| from mapgs.model.decoder import DETRDecoder |
| from mapgs.model.tokens import TokenManager, TokenMeta |
| from mapgs.model.gaussian_head import GaussianHead, DecodedGaussians |
| from mapgs.model.dynamic import DynamicModule |
| from mapgs.model.unet import FeatureUNet |
|
|
|
|
| class MapGS(nn.Module): |
| def __init__(self, cfg: MapGSConfig): |
| super().__init__() |
| self.cfg = cfg |
| m = cfg.model |
| self.encoder = ViTEncoder( |
| dim=m.embed_dim, patch=m.patch_size, depth=m.enc_depth, n_heads=m.n_heads, |
| mlp_ratio=m.mlp_ratio, qk_norm=m.qk_norm, layerscale_init=m.layerscale_init, |
| ) |
| self.token_manager = TokenManager(m) |
| self.decoder = DETRDecoder( |
| dim=m.embed_dim, depth=m.dec_depth, n_heads=m.n_heads, mlp_ratio=m.mlp_ratio, |
| qk_norm=m.qk_norm, layerscale_init=m.layerscale_init, shared_kv=m.shared_decoder_kv, |
| ) |
| self.head = GaussianHead(m) |
| self.dynamic = DynamicModule(cfg) |
| if m.feature_color and m.use_unet: |
| self.unet = FeatureUNet(in_ch=m.feature_dim, out_ch=3) |
| else: |
| self.unet = None |
|
|
| |
| def forward( |
| self, |
| images: torch.Tensor, |
| plucker: torch.Tensor, |
| timestep_ids: torch.Tensor, |
| anchor_pos: torch.Tensor, |
| anchor_type: torch.Tensor, |
| anchor_normal: torch.Tensor, |
| s_t: float = 1.0, |
| dynamic: Optional[dict] = None, |
| ) -> DecodedGaussians: |
| """Decode scene Gaussians in canonical frame. Dynamic Gaussians are |
| placed at a target timestamp afterwards via :meth:`place_dynamics`.""" |
| image_tokens = self.encode(images, plucker, timestep_ids) |
| tokens, meta, self_mask = self.build_query_tokens( |
| anchor_pos, anchor_type, anchor_normal, image_tokens.shape[0], |
| image_tokens.device, dynamic) |
| return self.decode(tokens, image_tokens, self_mask, meta, s_t) |
|
|
| |
| def encode(self, images, plucker, timestep_ids) -> torch.Tensor: |
| return self.encoder(images, plucker, timestep_ids) |
|
|
| def build_query_tokens(self, anchor_pos, anchor_type, anchor_normal, B, device, |
| dynamic: Optional[dict] = None): |
| tokens, meta = self.token_manager.build_static(anchor_pos, anchor_type, anchor_normal) |
| self_mask = None |
| if dynamic is not None: |
| dyn_tokens, dyn_meta = self.dynamic.build_tokens(dynamic, B, device) |
| tokens = torch.cat([tokens, dyn_tokens], dim=1) |
| meta = TokenMeta.cat_tokens(meta, dyn_meta) |
| self_mask = self.dynamic.self_mask(meta) |
| return tokens, meta, self_mask |
|
|
| def decode(self, tokens, image_tokens, self_mask, meta, s_t: float) -> DecodedGaussians: |
| tokens = self.decoder(tokens, image_tokens, self_mask=self_mask) |
| return self.head(tokens, meta, s_t) |
|
|
| def place_dynamics(self, decoded: DecodedGaussians, dynamic: Optional[dict], |
| frame_idx: int) -> DecodedGaussians: |
| """Rigidly place dynamic Gaussians at ``frame_idx`` (identity if static-only).""" |
| if dynamic is None: |
| return decoded |
| return self.dynamic.place_at(decoded, dynamic, frame_idx) |
|
|
| |
| def set_grad_checkpoint(self, enabled: bool = True): |
| self.encoder.grad_checkpoint = enabled |
| self.decoder.grad_checkpoint = enabled |
|
|
| def feature_to_rgb(self, feature_maps: torch.Tensor) -> torch.Tensor: |
| """Decode rendered feature maps ``[V, N_c, H, W]`` to RGB ``[V, 3, H, W]``.""" |
| if self.unet is not None: |
| return self.unet(feature_maps) |
| return feature_maps[:, :3].clamp(0, 1) |
|
|
| @property |
| def uses_features(self) -> bool: |
| return self.cfg.model.feature_color |
|
|
