tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Move masking abstraction

Browse Source
This commit is contained in:
Matt Tancik 2022-05-23 17:20:43 -07:00
parent ade3b632de
commit 800461e7ae
6 changed files with 82 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mattport/nerf/field_modules/encoding.py
View File

@ -215,7 +215,7 @@ class HashEncoding(Encoding):
assert in_tensor.shape[-1] == 3
in_tensor = in_tensor[..., None, :] # [..., 1, 3]
scaled = in_tensor * self.scalings.view(1, 1, -1, 1).to(in_tensor.device) # [..., L, 3]
scaled = in_tensor * self.scalings.view(-1, 1).to(in_tensor.device) # [..., L, 3]
scaled_c = torch.ceil(scaled).type(torch.int32)
scaled_f = torch.floor(scaled).type(torch.int32)

48
mattport/nerf/fields/base.py
View File

@ -3,13 +3,15 @@ Base class for the graphs.
"""
from abc import abstractmethod
from typing import Tuple
from typing import Optional, Tuple
from torch import nn
import torch
from torchtyping import TensorType
from mattport.nerf.field_modules.field_heads import FieldHeadNames
from mattport.structures.rays import PointSamples
from mattport.utils.misc import is_not_none
class Field(nn.Module):
@ -33,12 +35,48 @@ class Field(nn.Module):
"""
@abstractmethod
def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):
"""Computes and returns the colors."""
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."""
density, density_embedding = self.get_density(point_samples)
field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding)
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

2
mattport/nerf/fields/instant_ngp_field.py
View File

@ -102,7 +102,7 @@ class TCNNInstantNGPField(Field):
density = F.softplus(density_before_activation)
return density, base_mlp_out
def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):
def get_outputs(self, point_samples: PointSamples, density_embedding=None):
# TODO: add valid_mask masking!
assert is_not_none(density_embedding)
directions = point_samples.directions

43
mattport/nerf/fields/nerf_field.py
View File

@ -1,17 +1,17 @@
""" Classic NeRF field"""
from typing import Tuple
from typing import Optional, Tuple
import torch
from torch import nn
from torchtyping import TensorType
from mattport.nerf.field_modules.encoding import Encoding, Identity
from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, RGBFieldHead
from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead
from mattport.nerf.field_modules.mlp import MLP
from mattport.nerf.fields.base import Field
from mattport.structures.rays import PointSamples
from mattport.utils.misc import is_not_none
class NeRFField(Field):
@ -63,35 +63,18 @@ class NeRFField(Field):
def get_density(self, point_samples: PointSamples):
"""Computes and returns the densities."""
positions = point_samples.positions
valid_mask = point_samples.valid_mask
if not is_not_none(valid_mask):
valid_mask = torch.ones_like(positions[..., 0]).bool()
# placeholders for values to return
density = torch.zeros(*valid_mask.shape, 1, dtype=torch.float32, device=positions.device)
base_mlp_out = torch.zeros(
*valid_mask.shape, self.mlp_base.out_dim, dtype=torch.float32, device=positions.device
)
if not valid_mask.any(): # empty mask
return density, base_mlp_out
encoded_xyz = self.position_encoding(positions[valid_mask])
base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz)
density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask])
encoded_xyz = self.position_encoding(point_samples.positions)
base_mlp_out = self.mlp_base(encoded_xyz)
density = self.field_output_density(base_mlp_out)
return density, base_mlp_out
def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None):
directions = point_samples.directions
if not is_not_none(valid_mask):
valid_mask = torch.ones_like(directions[..., 0]).bool()
def get_outputs(
self, point_samples: PointSamples, density_embedding: Optional[TensorType] = None
) -> dict[FieldHeadNames, TensorType]:
"""Computes and returns the outputs."""
outputs = {}
for field_head in self.field_heads:
# placeholders for values to return
out = torch.zeros(*valid_mask.shape, field_head.out_dim, dtype=torch.float32, device=directions.device)
if not valid_mask.any(): # empty mask
return {field_head.field_head_name: out}
encoded_dir = self.direction_encoding(directions[valid_mask])
mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding[valid_mask]], dim=-1))
out[valid_mask] = field_head(mlp_out)
outputs[field_head.field_head_name] = out
encoded_dir = self.direction_encoding(point_samples.directions)
mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1))
outputs[field_head.field_head_name] = field_head(mlp_out)
return outputs

2
mattport/nerf/occupancy_grid.py
View File

@ -21,7 +21,7 @@ class OccupancyGrid(nn.Module):
self.num_cascades = num_cascades # the number of levels (i.e, cascades)
self.resolution = resolution
self.register_buffer("aabb", aabb) # axis-aligned bounding box
occupancy_grid = torch.zeros(
occupancy_grid = torch.ones(
[self.num_cascades] + [self.resolution] * 3
) # (num_cascades, x_resolution, y_resolution, z_resolution)
# not a module parameter, but it still should be saved and part of `state_dict`

23
mattport/structures/rays.py
View File

@ -27,6 +27,29 @@ class PointSamples:
camera_indices: TensorType[..., 1] = None
valid_mask: TensorType[...] = None
def apply_masks(self) -> "PointSamples":
"""Use valid_mask to mask samples.
Returns:
PointSamples: New set of masked samples.
"""
if is_not_none(self.valid_mask):
positions = self.positions[self.valid_mask]
directions = self.directions[self.valid_mask] if is_not_none(self.directions) else self.directions
camera_indices = (
self.camera_indices[self.valid_mask] if is_not_none(self.camera_indices) else self.camera_indices
)
return PointSamples(
positions=positions,
directions=directions,
camera_indices=camera_indices,
)
return PointSamples(
positions=self.positions,
directions=self.directions,
camera_indices=self.camera_indices,
)
@dataclass
class RaySamples: