nerfstudio/docs/notebooks/visualize_samplers.ipynb

Ray Samplers

Spaced Samplers

A basic sampler that spaces samples based on a predefined function.

In [5]:

# COLLAPSED
import torch
from nerfactory.cameras.rays import RayBundle
from nerfactory.graphs.modules import ray_sampler
import plotly.graph_objects as go
from plotly.subplots import make_subplots


num_samples = 1000
near = 2
far = 5
train_stratified = False

samplers = [
    ray_sampler.UniformSampler,
    ray_sampler.LinearDisparitySampler,
    ray_sampler.SqrtSampler,
    ray_sampler.LogSampler,
]

fig = make_subplots(
    rows=2,
    cols=2,
    subplot_titles=("Uniform", "Linear in Disparity", "Square Root", "Log Sampler"),
    shared_xaxes=True,
    shared_yaxes=True,
    vertical_spacing=0.1,
)

for i, Sampler in enumerate(samplers):
    sampler = Sampler(num_samples=num_samples, train_stratified=train_stratified)

    ray_bundle = RayBundle(
        origins=torch.ones([1, 3]),
        directions=torch.ones([1, 3]),
        pixel_area=torch.ones([1, 1]),
        nears=torch.ones([1, 1]) * near,
        fars=torch.ones([1, 1]) * far,
    )

    samples = sampler.generate_ray_samples(ray_bundle)

    trace = go.Histogram(x=samples.frustums.starts[0, :, 0], nbinsx=50)
    fig.append_trace(trace, i // 2 + 1, i % 2 + 1)

fig.update_yaxes(title_text="# Samples", row=1, col=1)
fig.update_yaxes(title_text="# Samples", row=2, col=1)
fig.update_xaxes(title_text="Distance", row=2, col=1)
fig.update_xaxes(title_text="Distance", row=2, col=2)

# Overlay both histograms
fig.update_layout(height=700, hovermode=False, showlegend=False, margin=dict(l=20, r=20, t=50, b=20))
fig.update_yaxes(range=[0, 80])
fig.update_traces(opacity=0.7)
fig.show()

PDF Sampler

In [6]:

# COLLAPSED
import torch
from nerfactory.cameras.rays import RayBundle
from nerfactory.graphs.modules import ray_sampler
import plotly.graph_objects as go
from plotly.subplots import make_subplots

num_coarse_samples = 20
num_samples = 1000
near = 2
far = 5
train_stratified = False

fig = make_subplots(
    rows=1,
    cols=2,
    subplot_titles=("PDF", "Samples"),
)

uniform_sampler = ray_sampler.UniformSampler(num_samples=num_coarse_samples, train_stratified=train_stratified)
pdf_sampler = ray_sampler.PDFSampler(num_samples=num_samples, train_stratified=train_stratified, include_original=False)

ray_bundle = RayBundle(
    origins=torch.ones([1, 3]),
    directions=torch.ones([1, 3]),
    pixel_area=torch.ones([1, 1]),
    nears=torch.ones([1, 1]) * near,
    fars=torch.ones([1, 1]) * far,
)

coarse_ray_samples = uniform_sampler(ray_bundle)

# Generate arbitrary PDF
weights = torch.ones(num_coarse_samples)
weights += torch.sin(torch.linspace(0, 3 * torch.pi, num_coarse_samples))
weights += torch.sin(torch.linspace(0, 0.5 * torch.pi, num_coarse_samples))
weights -= torch.min(weights)
weights /= torch.sum(weights)

samples = pdf_sampler.generate_ray_samples(ray_bundle, coarse_ray_samples, weights[None, :, None], num_samples)

# Plotting stuff
x = torch.ones((num_coarse_samples * 2))
x[::2] = coarse_ray_samples.frustums.starts[0, :, 0]
x[1::2] = coarse_ray_samples.frustums.ends[0, :, 0]

y = torch.ones((num_coarse_samples * 2))
y[::2] = weights
y[1::2] = weights

pdf_trace = go.Scatter(x=x, y=y)
fig.append_trace(pdf_trace, 1, 1)

samples_trace = go.Histogram(x=samples.frustums.starts[0, :, 0], nbinsx=100)
fig.append_trace(samples_trace, 1, 2)

fig.update_yaxes(title_text="# Samples", row=1, col=2)
fig.update_xaxes(title_text="Distance", row=1, col=1)
fig.update_xaxes(title_text="Distance", row=1, col=2)

# Overlay both histograms
fig.update_layout(height=400, hovermode=False, showlegend=False, margin=dict(l=20, r=20, t=50, b=20))
fig.update_traces(opacity=0.7)
fig.show()