diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index 7ea8b9b98..9c361186a 100644 --- a/.github/workflows/code_checks.yml +++ b/.github/workflows/code_checks.yml @@ -33,7 +33,7 @@ jobs: pylint mattport tests \ --max-line-length=120 \ --generated-members=numpy.*,torch.*,cv2.*,cv.* \ - --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts \ + --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u \ --disable=too-many-arguments \ --disable=fixme - name: Test with pytest diff --git a/.vscode/settings.json b/.vscode/settings.json index cea078e23..fa7771bbd 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -29,7 +29,7 @@ "python.linting.pylintArgs": [ "--generated-members=numpy.*,torch.*,cv2.*,cv.*", "--max-line-length=120", - "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts", + "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u", "--disable=too-many-arguments", "--disable=fixme" ], diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 439e09a7f..83ec23989 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -9,54 +9,13 @@ from torchtyping import TensorType from mattport.structures.rays import RayBundle, RaySamples -class Sampler(nn.Module): - """Base Sampler. Intended to be subclassed""" +class UniformSampler(nn.Module): + """Sample uniformly along a ray""" - def __init__(self, near_plane: float, far_plane: float, num_samples: int) -> None: - """Inialize sampler - - 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 - """ - super().__init__() - - self.near_plane = near_plane - self.far_plane = far_plane - self.num_samples = num_samples - - def generate_samples( - self, ray_bundle: RayBundle, near_plane: float, far_plane: float, num_samples: int - ) -> RaySamples: - """Encodes an input tensor. - - Args: - camera_rays (cameras.Rays): Rays to generate samples for - TensorType[..., "output_dim"]: A encoded input tensor - 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 - - Returns: - RaySamples: Positions and deltas for samples along a ray - """ - raise NotImplementedError - - def forward(self, ray_bundle: RayBundle) -> TensorType[..., "num_samples", 3]: - """Call forward""" - return self.generate_samples( - ray_bundle, near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_samples - ) - - -class UniformSampler(Sampler): - """Base Sampler. Intended to be subclassed""" - - def generate_samples( - self, ray_bundle: RayBundle, near_plane: float, far_plane: float, num_samples: int - ) -> RaySamples: - """Encodes an input tensor. + def forward( + self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int + ) -> TensorType[..., "num_samples", 3]: + """Generates position samples uniformly. Args: camera_rays (cameras.Rays): Rays to generate samples for @@ -72,7 +31,76 @@ class UniformSampler(Sampler): far_plane = far_plane or self.far_plane num_samples = num_samples or self.num_samples - bins = torch.linspace(near_plane, far_plane, num_samples + 1).to(ray_bundle.origins.device) # (num_samples,) + bins = torch.linspace(near_plane, far_plane, num_samples + 1) + + t = (bins[..., 1:] + bins[..., :-1]) / 2.0 + deltas = bins[..., 1:] - bins[..., :-1] + + positions = camera_rays.origin + camera_rays.direction * t + + return RaySamples(locations=positions, deltas=deltas) + + +class PDFSampler(nn.Module): + """Sample based on probability distribution""" + + def forward( + self, + camera_rays: cameras.Rays, + num_samples: int, + bins: TensorType[..., "num_samples"], + weights: TensorType[..., "num_samples"], + randomized: bool = True, + eps: float = 1e-5, + ) -> TensorType[..., "num_samples", 3]: + """Generates position samples given a distribution. + + Args: + camera_rays (cameras.Rays): Rays to generate samples for + TensorType[..., "output_dim"]: A encoded input tensor + num_samples (int): Number of samples per ray + bins (TensorType[..., "num_samples"]): Ray bins + weights: (TensorType[..., "num_samples"]): Weights for each bin + randomized: boolean: Randomize location within each bin. Defaults to True + eps: float: Small value to prevent numerical issues. Defaults to 1e-5 + + Returns: + RaySamples: Positions and deltas for samples along a ray + """ + + # Add small offset to rays with zero weight to prevent NaNs + weights_sum = torch.sum(weights, dim=-1, keepdim=True) + padding = torch.relu(eps - weights_sum) + weights += padding / weights.shape[-1] + weights_sum += padding + + pdf = weights / weights_sum + cdf = torch.min(1, torch.cumsum(pdf, dim=-1)) + cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1) + + if randomized: + u = torch.rand(size=(*cdf.shape[:-1], num_samples)) + else: + u = torch.linspace(0.0, 1.0, steps=num_samples) + u = torch.expand(size=(*cdf.shape[:-1], num_samples)) + + # u = u.contigous() + indicies = torch.searchsorted(cdf, u, right=True) + below = torch.max(torch.zeros_like(indicies), indicies - 1) + above = torch.min(cdf.shape[-1] - 1, indicies) + + indicies_g = torch.stack([below, above], -1) + cdf_g = torch.gather(cdf, indicies_g, axis=-1, batch_dims=len(indicies_g.shape) - 2) + bins_g = torch.gather(bins, indicies_g, axis=-1, batch_dims=len(indicies_g.shape) - 2) + + denom = cdf_g[..., 1] - cdf_g[..., 0] + denom = torch.where(denom < eps, torch.ones_like(denom), denom) + t = (u - cdf_g[..., 0]) / denom + samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0]) + + deltas = samples[..., 1:] - samples[..., :-1] + + positions = camera_rays.origin + camera_rays.direction * t num_rays = ray_bundle.origins.shape[0] ray_samples = RaySamples( diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py new file mode 100644 index 000000000..a7cc951f7 --- /dev/null +++ b/tests/nerf/test_sampler.py @@ -0,0 +1,27 @@ +""" +Test samplers +""" +from random import uniform +import pytest +import torch + +from mattport.nerf import sampler +from mattport.structures.cameras import Rays + + +def test_uniform_sampler(): + """Test uniform sampler""" + uniform_sampler = sampler.UniformSampler() + + origins = torch.zeros((10, 3)) + directions = torch.ones_like(origins) + camera_rays = Rays(origin=origins, direction=directions) + + +def test_pdf_sampler(): + """Test uniform sampler""" + # TODO Tancik: Add more precise tests + + +if __name__ == "__main__": + test_uniform_sampler()