tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 03:14:00 +08:00
Code Activity

Merge branch 'master' of github.com:ethanweber/mattport into eng/logging_cleanup

Browse Source
This commit is contained in:
evonneng committed 2022-05-20 15:58:09 +00:00
2 parents 61238247f2 10d713f84b
commit 1817d1cf9d
5 files changed
+151 -160
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

mattport/nerf/fields/base.py
+5 -5
View File
@@ -25,23 +25,23 @@ class Field(nn.Module):
return density
@abstractmethod
def get_density(self, point_samples: PointSamples, valid_mask=None):
def get_density(self, point_samples: PointSamples):
"""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):
def forward(self, point_samples: PointSamples):
"""Evaluates the field at points along the ray."""
density, density_embedding = self.get_density(point_samples, valid_mask=valid_mask_density)
density, density_embedding = self.get_density(point_samples)
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):
def forward_with_weight_pruning(self, ray_samples: RaySamples):
"""Forward but with pruning based on weights, from the densities."""
density, density_embedding = self.get_density(ray_samples, valid_mask=valid_mask_density)
density, density_embedding = self.get_density(ray_samples)
weights = ray_samples.get_weights(density)
valid_mask_outputs = weights >= self.weights_threshold
field_outputs = self.get_outputs(
mattport/nerf/graph/vanilla_nerf.py
+3 -2
View File
@@ -81,9 +81,10 @@ class NeRFField(Field):
self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim())
self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())
def get_density(self, point_samples: PointSamples, valid_mask=None):
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
@@ -178,7 +179,7 @@ class NeRFGraph(Graph):
depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.ts)
# pdf sampling
ray_samples_pdf = self.sampler_pdf(ray_samples_uniform, weights_coarse)
ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)
# fine field:
field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples())
mattport/nerf/sampler.py
+64 -90
View File
@@ -2,99 +2,69 @@
Collection of sampling strategies
"""
from abc import abstractmethod
from typing import Optional
import torch
from torch import nn
from torchtyping import TensorType
from mattport.nerf.occupancy_grid import OccupancyGrid
from mattport.nerf.occupancy_grid import OccupancyGrid
from mattport.structures.rays import RayBundle, RaySamples
class OccupancyGridSampler(nn.Module):
"""Sample based on occupancy grid"""
class Sampler(nn.Module):
"""Generate Samples"""
def __init__(
self, num_samples: int, occupancy_field: Optional[OccupancyGrid] = None, weight_threshold: float = 1e-4
) -> None:
super().__init__()
self.num_samples = num_samples
self.occupancy_field = occupancy_field
self.weight_threshold = weight_threshold
@abstractmethod
def generate_ray_samples(self) -> RaySamples:
"""Generate Ray Samples"""
def forward(self, *args, **kwargs) -> RaySamples:
"""Generate ray samples with optional occupancy filtering"""
ray_samples = self.generate_ray_samples(*args, **kwargs)
if self.occupancy_field is not None:
densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True)
weights = ray_samples.get_weights(densities)
valid_mask = weights >= self.weight_threshold
ray_samples.set_valid_mask(valid_mask)
return ray_samples
class UniformSampler(Sampler):
"""Sample uniformly along a ray"""
def __init__(
self,
occupancy_grid: OccupancyGrid,
num_samples: int = 256,
weight_threshold: float = 1e-4,
num_samples: int,
train_stratified=True,
occupancy_field: Optional[OccupancyGrid] = None,
weight_threshold: float = 1e-4,
) -> None:
"""
Args:
"""
super().__init__()
self.num_samples = num_samples
self.train_stratified = train_stratified
self.weight_threshold = weight_threshold
self.occupancy_grid = occupancy_grid
@torch.no_grad()
def forward(
self,
ray_bundle: RayBundle,
num_samples: Optional[int] = None,
) -> RaySamples:
"""Generates position samples uniformly.
Args:
ray_bundle (RayBundle): Rays to generate samples for
num_samples (Optional[int]): Number of samples per ray
Returns:
RaySamples: Positions and deltas for samples along a ray
"""
assert ray_bundle.nears is not None
assert ray_bundle.fars is not None
num_samples = num_samples or self.num_samples
num_rays = ray_bundle.origins.shape[0]
bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device) # shape (num_samples+1,)
bins = ray_bundle.nears[:, None] + bins[None, :] * (
ray_bundle.fars[:, None] - ray_bundle.nears[:, None]
) # shape (num_rays, num_samples+1)
if self.train_stratified and self.training:
t_rand = torch.rand((num_rays, num_samples), dtype=bins.dtype, device=bins.device)
ts = bins[:, :-1] + t_rand * (bins[:, 1:] - bins[:, :-1]) # shape (num_rays, num_samples)
else:
ts = (bins[:, 1:] + bins[:, :-1]) / 2
ray_samples = RaySamples(
ts=ts,
ray_bundle=ray_bundle,
)
densities = self.occupancy_grid.get_densities(ray_samples.positions, update_iter_count=True)
weights = ray_samples.get_weights(densities)
valid_mask = weights >= self.weight_threshold
return ray_samples, weights, valid_mask
class UniformSampler(nn.Module):
"""Sample uniformly along a ray"""
def __init__(self, num_samples: int, train_stratified=True) -> None:
"""
Args:
near_plane (float): Minimum distance along ray to sample
far_plane (float): Maximum distance along ray to sample
num_samples (int): Number of samples per ray
train_stratified (bool): Use stratified sampling during training. Defults to True
occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides,
samples below weight_threshold as set as invalid.
weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided.
"""
super().__init__()
self.num_samples = num_samples
super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold)
self.train_stratified = train_stratified
@torch.no_grad()
def forward(
def generate_ray_samples(
self,
ray_bundle: RayBundle,
ray_bundle: RayBundle = None,
num_samples: Optional[int] = None,
) -> RaySamples:
"""Generates position samples uniformly.
@@ -106,6 +76,7 @@ class UniformSampler(nn.Module):
Returns:
RaySamples: Positions and deltas for samples along a ray
"""
assert ray_bundle is not None
assert ray_bundle.nears is not None
assert ray_bundle.fars is not None
@@ -123,35 +94,41 @@ class UniformSampler(nn.Module):
else:
ts = (bins[:, 1:] + bins[:, :-1]) / 2
ray_samples = RaySamples(
ts=ts,
ray_bundle=ray_bundle,
)
ray_samples = ray_bundle.get_ray_samples(ts)
return ray_samples
class PDFSampler(nn.Module):
class PDFSampler(Sampler):
"""Sample based on probability distribution"""
def __init__(self, num_samples: int, train_stratified: bool = True, include_original: bool = True) -> None:
def __init__(
self,
num_samples: int,
train_stratified: bool = True,
include_original: bool = True,
occupancy_field: OccupancyGrid = None,
weight_threshold: float = 1e-4,
) -> None:
"""
Args:
num_samples (int): Number of samples per ray
train_stratified: boolean: Randomize location within each bin during training. Defaults to True
include_original: Add original samples to ray. Defaults to True
occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides,
samples below weight_threshold as set as invalid.
weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided.
"""
super().__init__()
self.num_samples = num_samples
self.include_original = include_original
super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold)
self.train_stratified = train_stratified
self.include_original = include_original
@torch.no_grad()
def forward(
def generate_ray_samples(
self,
coarse_ray_samples: RaySamples,
weights: TensorType[..., "num_samples"],
ray_bundle: RayBundle = None,
ray_samples: RaySamples = None,
weights: TensorType[..., "num_samples"] = None,
num_samples: Optional[int] = None,
eps: float = 1e-5,
) -> RaySamples:
@@ -195,7 +172,7 @@ class PDFSampler(nn.Module):
indicies_g = torch.stack([below, above], -1)
matched_shape = (indicies_g.shape[0], indicies_g.shape[1], cdf.shape[-1])
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g)
bins_g = torch.gather(coarse_ray_samples.ts.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g)
bins_g = torch.gather(ray_samples.ts.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g)
denom = cdf_g[..., 1] - cdf_g[..., 0]
denom = torch.where(denom < eps, torch.ones_like(denom), denom)
@@ -203,16 +180,13 @@ class PDFSampler(nn.Module):
ts = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
if self.include_original:
ts, _ = torch.sort(torch.cat([coarse_ray_samples.ts, ts], -1), -1)
ts, _ = torch.sort(torch.cat([ray_samples.ts, ts], -1), -1)
else:
ts, _ = torch.sort(ts, -1)
# Stop gradients
ts = ts.detach()
ray_samples = RaySamples(
ts=ts,
ray_bundle=coarse_ray_samples.ray_bundle,
)
ray_samples = ray_bundle.get_ray_samples(ts)
return ray_samples
mattport/structures/rays.py
+78 -62
View File
@@ -4,8 +4,8 @@ Some ray datastructures.
import random
from dataclasses import dataclass
from typing import Optional
import torch
import torch
from torchtyping import TensorType
from mattport.utils.misc import is_not_none
@@ -18,6 +18,51 @@ class PointSamples:
positions: TensorType[..., 3] = None
directions: TensorType[..., 3] = None
camera_indices: TensorType[..., 1] = None
valid_mask: TensorType[..., 1] = None
@dataclass
class RaySamples:
"""Samples along a ray"""
positions: TensorType[..., 3] = None
directions: TensorType[..., 3] = None
camera_indices: TensorType[..., 1] = None
valid_mask: TensorType[..., 1] = None
ts: TensorType[..., 1] = None
deltas: TensorType[..., 1] = None
def to_point_samples(self) -> PointSamples:
"""Convert to PointSamples instance and return."""
# TODO: make this more interpretable
return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask)
def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]:
"""Return weights based on predicted densities
Args:
densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray
Returns:
TensorType[..., "num_samples"]: Weights for each sample
"""
delta_density = self.deltas * densities[..., 0]
alphas = 1 - torch.exp(-delta_density)
transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
transmittance = torch.cat(
[torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1
)
transmittance = torch.exp(-transmittance) # [..., "num_samples"]
weights = alphas * transmittance # [..., "num_samples"]
return weights
def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None:
"""Sets valid mask"""
self.valid_mask = valid_mask
@dataclass
@@ -85,6 +130,38 @@ class RayBundle:
valid_mask=self.valid_mask[valid_mask] if is_not_none(self.valid_mask) else None,
)
def get_ray_samples(self, ts: TensorType["num_rays", "num_samples"]) -> RaySamples:
"""
Args:
ts (TensorType["num_rays", "num_samples"]): _description_
Returns:
RaySamples: _description_
"""
positions = self.origins[:, None] + ts[:, :, None] * self.directions[:, None]
directions = self.directions.unsqueeze(1).repeat(1, positions.shape[1], 1)
valid_mask = torch.ones_like(ts, dtype=torch.bool)
dists = ts[..., 1:] - ts[..., :-1]
dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples]
deltas = dists * torch.norm(self.directions[..., None, :], dim=-1)
if self.camera_indices is not None:
camera_indices = self.camera_indices.unsqueeze(1).repeat(1, positions.shape[1])
else:
camera_indices = None
ray_samples = RaySamples(
positions=positions,
directions=directions,
camera_indices=camera_indices,
valid_mask=valid_mask,
ts=ts,
deltas=deltas,
)
return ray_samples
@dataclass
class CameraRayBundle:
@@ -130,64 +207,3 @@ class CameraRayBundle:
directions=self.directions.view(-1, 3)[start_idx:end_idx],
camera_indices=camera_indices,
)
class RaySamples:
"""_summary_"""
def __init__(
self,
ts: TensorType["num_rays", "num_samples"],
ray_bundle: RayBundle,
) -> None:
self.ray_bundle = ray_bundle
self.ts = ts
self.positions = self.get_positions(ray_bundle)
self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1)
self.deltas = self.get_deltas()
def to_point_samples(self):
"""Convert to PointSamples instance and return."""
# TODO: make this more interpretable
return PointSamples(positions=self.positions, directions=self.directions)
def get_camera_indices(self):
"""Returns camera indices."""
assert not isinstance(self.ray_bundle.camera_indices, type(None)), "Camera indices cannot be None here."
camera_indices = self.ray_bundle.camera_indices.unsqueeze(1).repeat(1, self.positions.shape[1])
return camera_indices
def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]:
"""Returns positions."""
return ray_bundle.origins[:, None] + self.ts[:, :, None] * ray_bundle.directions[:, None]
def get_deltas(self) -> TensorType[..., "num_samples"]:
"""Returns deltas."""
dists = self.ts[..., 1:] - self.ts[..., :-1]
dist_inf = 1e10 * torch.ones_like(dists[..., -1:])
dists = torch.cat([dists, dist_inf], -1) # [N_rays, N_samples]
deltas = dists * torch.norm(self.ray_bundle.directions[..., None, :], dim=-1)
return deltas
def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]:
"""Return weights based on predicted densities
Args:
densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray
Returns:
TensorType[..., "num_samples"]: Weights for each sample
"""
delta_density = self.deltas * densities[..., 0]
alphas = 1 - torch.exp(-delta_density)
transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
transmittance = torch.cat(
[torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1
)
transmittance = torch.exp(-transmittance) # [..., "num_samples"]
weights = alphas * transmittance # [..., "num_samples"]
return weights
tests/nerf/test_sampler.py
+1 -1
View File
@@ -47,7 +47,7 @@ def test_pdf_sampler():
# Just check that it doesn't crash. Need to add some actual tests.
pdf_sampler = PDFSampler(num_samples)
pdf_sampler(coarse_ray_samples, weights, num_samples)
pdf_sampler(ray_bundle, coarse_ray_samples, weights, num_samples)
# TODO Tancik: Add more precise tests