| """Contains decoder head for direct prediction of delta values. |
| |
| For licensing see accompanying LICENSE file. |
| Copyright (C) 2025 Apple Inc. All Rights Reserved. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
| from .gaussian_decoder import ImageFeatures |
|
|
|
|
| class DirectPredictionHead(nn.Module): |
| """Decodes features into delta values using convolutions.""" |
|
|
| def __init__(self, feature_dim: int, num_layers: int) -> None: |
| """Initialize DirectGaussianPredictor. |
| |
| Args: |
| feature_dim: Number of input features. |
| num_layers: The number of layers of Gaussians to predict. |
| """ |
| super().__init__() |
| self.num_layers = num_layers |
|
|
| |
| self.geometry_prediction_head = nn.Conv2d(feature_dim, 3 * num_layers, 1) |
| self.geometry_prediction_head.weight.data.zero_() |
| assert self.geometry_prediction_head.bias is not None |
| self.geometry_prediction_head.bias.data.zero_() |
|
|
| self.texture_prediction_head = nn.Conv2d(feature_dim, (14 - 3) * num_layers, 1) |
| self.texture_prediction_head.weight.data.zero_() |
| assert self.texture_prediction_head.bias is not None |
| self.texture_prediction_head.bias.data.zero_() |
|
|
| def forward(self, image_features: ImageFeatures) -> torch.Tensor: |
| """Predict deltas for 3D Gaussians. |
| |
| Args: |
| image_features: Image features from decoder. |
| |
| Returns: |
| The predicted deltas for Gaussian attributes. |
| """ |
| delta_values_geometry = self.geometry_prediction_head(image_features.geometry_features) |
| delta_values_texture = self.texture_prediction_head(image_features.texture_features) |
| delta_values_geometry = delta_values_geometry.unflatten(1, (3, self.num_layers)) |
| delta_values_texture = delta_values_texture.unflatten(1, (14 - 3, self.num_layers)) |
| delta_values = torch.cat([delta_values_geometry, delta_values_texture], dim=1) |
| return delta_values |
|
|
