From ee3a4480cc40ce1f4c9a35df9ba15a8d049e2e1f Mon Sep 17 00:00:00 2001 From: Ruilong Li Date: Fri, 16 Sep 2022 02:36:43 +0000 Subject: [PATCH 1/3] support stratified sampling --- README.md | 2 +- examples/trainval.py | 3 ++- nerfacc/utils.py | 4 ++++ nerfacc/volumetric_rendering.py | 2 ++ 4 files changed, 9 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 2dbe7f03..e40f985c 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ Ours on TITAN RTX : | trainval | Lego | Mic | Materials | | - | - | - | - | | Time | 300s | 272s | 258s | -| PSNR | 36.28 | 36.16 | 29.76 | +| PSNR | 36.61 | 37.45 | 30.15 | | FPS | 11.49 | 21.48 | 8.86 | Instant-NGP paper (5 min) on 3090: diff --git a/examples/trainval.py b/examples/trainval.py index 1dd4eb86..588c68ab 100644 --- a/examples/trainval.py +++ b/examples/trainval.py @@ -46,6 +46,7 @@ def render_image(radiance_field, rays, render_bkgd, render_step_size): scene_resolution=occ_field.resolution, render_bkgd=render_bkgd, render_step_size=render_step_size, + stratified=radiance_field.training, ) results.append(chunk_results) rgb, depth, acc, counter, compact_counter = [ @@ -65,7 +66,7 @@ def render_image(radiance_field, rays, render_bkgd, render_step_size): torch.manual_seed(42) device = "cuda:0" - scene = "lego" + scene = "materials" # setup dataset train_dataset = SubjectLoader( diff --git a/nerfacc/utils.py b/nerfacc/utils.py index cd6ffe12..6e923a0e 100644 --- a/nerfacc/utils.py +++ b/nerfacc/utils.py @@ -44,6 +44,7 @@ def volumetric_marching( t_min: torch.Tensor = None, t_max: torch.Tensor = None, render_step_size: float = 1e-3, + stratified: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Volumetric marching with occupancy test. @@ -62,6 +63,7 @@ def volumetric_marching( t_max: Optional. Ray far planes. Tensor with shape (n_ray,) If not given it will be calculated using aabb test. Default is None. render_step_size: Marching step size. Default is 1e-3. + stratified: Whether to use stratified sampling. Default is False. Returns: packed_info: Stores infomation on which samples belong to the same ray. @@ -83,6 +85,8 @@ def volumetric_marching( == scene_resolution[0] * scene_resolution[1] * scene_resolution[2] ), f"Shape {scene_occ_binary.shape} is not right!" + if stratified: + t_min = t_min + torch.rand_like(t_min) * render_step_size ( packed_info, frustum_origins, diff --git a/nerfacc/volumetric_rendering.py b/nerfacc/volumetric_rendering.py index e491ecee..767acd95 100644 --- a/nerfacc/volumetric_rendering.py +++ b/nerfacc/volumetric_rendering.py @@ -19,6 +19,7 @@ def volumetric_rendering( scene_resolution: Tuple[int, int, int], render_bkgd: torch.Tensor, render_step_size: int, + stratified: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """A *fast* version of differentiable volumetric rendering.""" n_rays = rays_o.shape[0] @@ -47,6 +48,7 @@ def volumetric_rendering( scene_occ_binary=scene_occ_binary, # sampling render_step_size=render_step_size, + stratified=stratified, ) frustum_positions = ( frustum_origins + frustum_dirs * (frustum_starts + frustum_ends) / 2.0 From 30ba03224335cf5a490f823ca4df355796112cc3 Mon Sep 17 00:00:00 2001 From: Ruilong Li Date: Fri, 16 Sep 2022 02:44:35 +0000 Subject: [PATCH 2/3] bump version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d7342411..313a8865 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "nerfacc" -version = "0.0.3" +version = "0.0.4" authors = [{name = "Ruilong", email = "ruilongli94@gmail.com"}] license = { text="MIT" } requires-python = ">=3.8" From dabf37a097d64bda510c4fae381ee6ca66159040 Mon Sep 17 00:00:00 2001 From: Ruilong Li Date: Fri, 16 Sep 2022 03:04:17 +0000 Subject: [PATCH 3/3] more benchmark --- README.md | 18 +++++++++--------- acc_binary_test.png | Bin 0 -> 2557 bytes acc_binary_train.png | Bin 0 -> 2359 bytes examples/trainval.py | 2 +- rgb_train.png | Bin 0 -> 266910 bytes 5 files changed, 10 insertions(+), 10 deletions(-) create mode 100644 acc_binary_test.png create mode 100644 acc_binary_train.png create mode 100644 rgb_train.png diff --git a/README.md b/README.md index e40f985c..760db555 100644 --- a/README.md +++ b/README.md @@ -12,17 +12,17 @@ python examples/trainval.py Ours on TITAN RTX : -| trainval | Lego | Mic | Materials | -| - | - | - | - | -| Time | 300s | 272s | 258s | -| PSNR | 36.61 | 37.45 | 30.15 | -| FPS | 11.49 | 21.48 | 8.86 | +| trainval | Lego | Mic | Materials | Chair | Hotdog | +| - | - | - | - | - | - | +| Time | 300s | 272s | 258s | 331s | 287s | +| PSNR | 36.61 | 37.45 | 30.15 | 36.06 | 38.17 | +| FPS | 11.49 | 21.48 | 8.86 | 15.61 | 7.38 | -Instant-NGP paper (5 min) on 3090: +Instant-NGP paper (5 min) on 3090 (w/ mask): -| trainval | Lego | Mic | Materials | -| - | - | - | - | -| PSNR | 36.39 | 36.22 | 29.78 | +| trainval | Lego | Mic | Materials | Chair | Hotdog | +| - | - | - | - | - | - | +| PSNR | 36.39 | 36.22 | 29.78 | 35.00 | 37.40 |