mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
Ethan Weber
fb440fedf3
* occupancy grid in progress * visualizing occupancy grid. currently restructing field and renderer head outputs * black and pylint. broke nerfw and semantic nerf but it's okay * get_project_root
52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
"""
|
|
Base class for the graphs.
|
|
"""
|
|
|
|
from abc import abstractmethod
|
|
|
|
from torch import nn
|
|
|
|
from mattport.structures.rays import PointSamples, RaySamples
|
|
|
|
|
|
class Field(nn.Module):
|
|
"""Base class for fields."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
# minimum weight for it to be valid. used to skip querying the field
|
|
# TODO(ethan): avoid the hardcoding
|
|
self.weights_threshold = 1e-4
|
|
|
|
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, valid_mask=None):
|
|
"""Computes and returns the densities."""
|
|
|
|
@abstractmethod
|
|
def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):
|
|
"""Computes and returns the colors."""
|
|
|
|
def forward(self, point_samples: PointSamples, valid_mask_density=None):
|
|
"""Evaluates the field at points along the ray."""
|
|
density, density_embedding = self.get_density(point_samples, valid_mask=valid_mask_density)
|
|
field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding)
|
|
field_outputs["density"] = density
|
|
return field_outputs
|
|
|
|
def forward_with_weight_pruning(self, ray_samples: RaySamples, valid_mask_density=None):
|
|
"""Forward but with pruning based on weights, from the densities."""
|
|
density, density_embedding = self.get_density(ray_samples, valid_mask=valid_mask_density)
|
|
weights = ray_samples.get_weights(density)
|
|
valid_mask_outputs = weights >= self.weights_threshold
|
|
field_outputs = self.get_outputs(
|
|
ray_samples, density_embedding=density_embedding, valid_mask=valid_mask_outputs
|
|
)
|
|
field_outputs["density"] = density
|
|
return field_outputs
|