mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 10:55:57 +08:00
83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
"""
|
|
Base class for the graphs.
|
|
"""
|
|
|
|
from abc import abstractmethod
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
from torch import nn
|
|
import torch
|
|
from torchtyping import TensorType
|
|
from pyrad.nerf.field_modules.field_heads import FieldHeadNames
|
|
|
|
from pyrad.structures.rays import PointSamples
|
|
from pyrad.utils.misc import is_not_none
|
|
|
|
|
|
class Field(nn.Module):
|
|
"""Base class for fields."""
|
|
|
|
def density_fn(self, positions):
|
|
"""Returns only the density. Used primarily with the occupancy grid."""
|
|
point_samples = PointSamples(positions=positions)
|
|
density, _ = self.get_density(point_samples)
|
|
return density
|
|
|
|
@abstractmethod
|
|
def get_density(self, point_samples: PointSamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]:
|
|
"""Computes and returns the densities.
|
|
|
|
Args:
|
|
point_samples (PointSamples): Samples locations to compute density.
|
|
|
|
Returns:
|
|
Tuple[TensorType[...,1], TensorType[...,"num_features"]]: A tensor of densities and a tensor of features.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_outputs(
|
|
self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None
|
|
) -> Dict[FieldHeadNames, TensorType]:
|
|
"""Computes and returns the colors.
|
|
|
|
Args:
|
|
point_samples (PointSamples): Samples locations to compute outputs.
|
|
density_embedding (TensorType, optional): Density embeddings to condition on.
|
|
|
|
Returns:
|
|
Dict[FieldHeadNames, TensorType]: Output field values.
|
|
"""
|
|
|
|
def forward(self, point_samples: PointSamples):
|
|
"""Evaluates the field at points along the ray."""
|
|
valid_mask = point_samples.valid_mask
|
|
|
|
if is_not_none(valid_mask):
|
|
# Hacky handling of empty masks. Tests on a single ray but doesn't use results
|
|
if not valid_mask.any():
|
|
point_samples = PointSamples(
|
|
positions=point_samples.positions[0, :], directions=point_samples.directions[0, :]
|
|
)
|
|
else:
|
|
point_samples = point_samples.apply_masks()
|
|
density_masked, density_embedding_masked = self.get_density(point_samples)
|
|
field_outputs_masked = self.get_outputs(point_samples, density_embedding=density_embedding_masked)
|
|
|
|
field_outputs = {}
|
|
for k, value in field_outputs_masked.items():
|
|
zeros = torch.zeros(*valid_mask.shape, value.shape[-1], dtype=torch.float32, device=valid_mask.device)
|
|
if valid_mask.any():
|
|
zeros[valid_mask] = value
|
|
else:
|
|
zeros[0, :] = value
|
|
field_outputs[k] = zeros
|
|
density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=valid_mask.device)
|
|
if valid_mask.any():
|
|
density[valid_mask] = density_masked
|
|
else:
|
|
density, density_embedding = self.get_density(point_samples)
|
|
field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding)
|
|
|
|
field_outputs[FieldHeadNames.DENSITY] = density
|
|
return field_outputs
|