54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
"""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
|
|
|
|
# 14 is 3 means, 3 scales, 4 quaternions, 3 colors and 1 opacity
|
|
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
|