Optimizing camera undistortion #2037

DreekFire · 2023-06-05T11:56:41Z

Only runs Newton's method until sub-pixel accuracy is reached instead of all 10 iterations. Then resamples images with bilinear sampling to get pixel color at points with error.

Also uses analytical formulas for pixel areas instead of distorting coordinates of neighboring pixels to get pixel area.

…pute pixel area

DreekFire · 2023-06-05T12:12:38Z

Need example of datasets with different camera types to finish testing. Also contains lots of junk from merging branches, will clean up.

tancik

Nice! Added some initial comments.

tancik · 2023-06-05T17:04:19Z

.vscode/settings.json

Remove the changes in this file.

tancik · 2023-06-05T17:10:07Z

nerfstudio/cameras/camera_utils.py

@@ -413,40 +417,226 @@ def radial_and_tangential_undistort(
    distortion_params: torch.Tensor,
    eps: float = 1e-3,
    max_iterations: int = 10,
-) -> torch.Tensor:
+    resolution: torch.Tensor = torch.tensor([1e-3, 1e-3]),


If you know the size, you should type it, ie
resolution: Float[torch.Tensor, "2"]

tancik · 2023-06-05T17:11:08Z

nerfstudio/cameras/camera_utils.py

        max_iterations: The maximum number of iterations to perform.
+        resolution: The resolution (w, h of each pixel, in units of multiples of focal length)



Add doc for tolerance

tancik · 2023-06-05T17:11:29Z

nerfstudio/cameras/camera_utils.py

@@ -413,40 +417,226 @@ def radial_and_tangential_undistort(
    distortion_params: torch.Tensor,
    eps: float = 1e-3,
    max_iterations: int = 10,
-) -> torch.Tensor:
+    resolution: torch.Tensor = torch.tensor([1e-3, 1e-3]),
+    tol: float = 0.5,


rename to tolerance

tancik · 2023-06-05T17:13:20Z

nerfstudio/cameras/camera_utils.py

+    # n_samples = coords.shape[0]
+    # n_iters = 0


Remove commented out code, here and elsewhere

tancik

Nice! Added some initial comments.

tancik · 2023-06-05T20:30:57Z

nerfstudio/data/datamanagers/base_datamanager.py

Many of the changes in this file don't seem related to this PR

Yeah would suggest to revert the things unrelated

liruilong940607 · 2023-06-05T21:28:01Z

nerfstudio/model_components/ray_generators.py

-        camera_opt_to_camera = self.pose_optimizer(c)
+        if self.pose_optimizer is not None:
+            camera_opt_to_camera = self.pose_optimizer(c)
+        else:
+            camera_opt_to_camera = None


Revert this change as unrelated to this PR?

liruilong940607 · 2023-06-05T21:28:31Z

nerfstudio/model_components/ray_generators.py

@@ -40,7 +41,8 @@ def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None:
        self.pose_optimizer = pose_optimizer
        self.register_buffer("image_coords", cameras.get_image_coords(), persistent=False)

-    def forward(self, ray_indices: Int[Tensor, "num_rays 3"]) -> RayBundle:
+    @profiler.time_function


Remove profiler here

liruilong940607 · 2023-06-05T21:28:39Z

nerfstudio/model_components/ray_generators.py

@@ -21,6 +21,7 @@
 from nerfstudio.cameras.camera_optimizers import CameraOptimizer
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.utils import profiler


And remove here

liruilong940607 · 2023-06-05T21:35:52Z

nerfstudio/cameras/camera_utils.py


    Returns:
        The residuals (fx, fy) and jacobians (fx_x, fx_y, fy_x, fy_y).
    """
+    assert distortion_params.shape[-1] == 8


The fisheye undistortion, which was tied to this function, is actually with correct formula. So instead of creating two separate newton functions for fisheye and perspective cameras, you can keep them in the same function like the old way. (The only minor issue with this function is the k4 for perspective camera but that's all zero for the current datasets anyway).

liruilong940607 · 2023-06-05T21:37:15Z

nerfstudio/cameras/camera_utils.py

    """Computes undistorted coords given opencv distortion parameters.
    Adapted from MultiNeRF
    https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509

    Args:
        coords: The distorted coordinates.
-        distortion_params: The distortion parameters [k1, k2, k3, k4, p1, p2].
-        eps: The epsilon for the convergence.
+        distortion_params: The distortion parameters. Supports 0, 1, 2, 4, 8 parameters, in


Revert here if the _compute_residual_and_jacobian is reverted.

liruilong940607 · 2023-06-05T21:37:29Z

nerfstudio/cameras/camera_utils.py

+    assert distortion_params.shape[-1] in [0, 1, 2, 4, 8]
+
+    if distortion_params.shape[-1] == 0:
+        return coords, torch.eye(2, device=coords.device), coords
+
+    if distortion_params.shape[-1] < 8:
+        distortion_params = F.pad(distortion_params, (0, 8 - distortion_params.shape[-1]), "constant", 0.0)
+    assert distortion_params.shape[-1] == 8


liruilong940607 · 2023-06-05T21:39:11Z

nerfstudio/data/datamanagers/base_datamanager.py

Yeah would suggest to revert the things unrelated

liruilong940607 · 2023-06-06T01:13:11Z

Thanks @DreekFire for the updates. This is not an easy one.

My main question is on the fisheye formula. Have you run any tests to verify the fisheye formula is correct?

The reason I'm asking this is that the old version actually have a correct fisheye undistortion formula by combining

nerfstudio/nerfstudio/cameras/camera_utils.py

Line 411 in 43c399e

def radial_and_tangential_undistort(

and

nerfstudio/nerfstudio/cameras/cameras.py

Lines 724 to 735 in 43c399e

    
           theta = torch.sqrt(torch.sum(coord_stack**2, dim=-1)) 
        
           theta = torch.clip(theta, 0.0, math.pi) 
        
           sin_theta = torch.sin(theta) 
        
           directions_stack[..., 0][mask] = torch.masked_select( 
        
               coord_stack[..., 0] * sin_theta / theta, mask 
        
           ).float() 
        
           directions_stack[..., 1][mask] = torch.masked_select( 
        
               coord_stack[..., 1] * sin_theta / theta, mask 
        
           ).float() 
        
           directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask).float()

And this PR seems to change the radial_and_tangential_undistort function but not the second part for fisheye. I'm not sure if the formula is correct anymore.

kerrj · 2023-06-13T18:20:53Z

What's the speed/quality difference on this PR?

DreekFire · 2023-06-13T20:56:16Z

Some tests for correctness:

Visually inspected results on floating-tree and poster datasets - looks good (todo: equirectangular and stereo, even though these have no undistortion, just to make sure nothing broke)
Examined max, mean, and min pixel_areas under new method and old method: approximately matches (expect some difference because new method uses more accurate analytical formula which accounts for parallelogram-shaped pixels)
Bilinear sampling: manual inspection on dummy data
Measured max difference between rays returned by new and old methods: within specified tolerance for perspective camera, different for fisheye camera BUT new method is more accurate for fisheye camera (from ~1e-3 mean abs error to ~1e-8 (less than floating point precision) mean abs error.

And speed:
Speed fluctuates, measured manually at two different iteration ranges:
Units in thousand rays per second, on puppy.bair
Iteration 500-1800: 133-137 new fisheye, 131-132 new perspective, 126-127 old fisheye, 125-127 old perspective
Iteration 2000-3000: 150 new fisheye, 143-145 new perspective, 135-140 old fisheye, 133-135 old perspective
Speedup not as great as before. During development, I merged a new version of nerfstudio. Before and after the merge, the new method ran at about the same speed, but the older version sped up a bit, so the difference is not as great as expected.

kerrj · 2023-06-15T06:25:24Z

what machine are you testing on? is there any particular sort of hardware this implementation is optimized for? eg would you see stronger benefits on weaker CPUs or GPUs, etc

tancik · 2023-06-24T00:39:06Z

@kerrj Will this conflict with any of the datamanager stuff that you are working on?

DreekFire · 2023-08-01T09:56:30Z

Ended up removing the resampling because it didn't appear to be worth it - once a ray was already within sub-pixel accuracy, Newton's method usually achieved very good accuracy in just one more iteration. This makes the code simpler as well. Now, the key optimizations are early stopping for undistortion and analytical formulas for pixel areas instead of numerical ones (should also be more accurate because it correctly accounts for skewed pixels).

The speed fluctuates throughout training, but it is generally about 10% faster at any given iteration than the main branch on puppy.

Not sure why the tests are failing - seems to be something going wrong with torch.compile. It passes on both my local machine as well as my env on the puppy machine.

anc2001 · 2023-09-12T18:50:48Z

Thanks @DreekFire for the updates. This is not an easy one.

My main question is on the fisheye formula. Have you run any tests to verify the fisheye formula is correct?

The reason I'm asking this is that the old version actually have a correct fisheye undistortion formula by combining

nerfstudio/nerfstudio/cameras/camera_utils.py

Line 411 in 43c399e

def radial_and_tangential_undistort(

and

nerfstudio/nerfstudio/cameras/cameras.py

Lines 724 to 735 in 43c399e

theta = torch.sqrt(torch.sum(coord_stack**2, dim=-1))

theta = torch.clip(theta, 0.0, math.pi)

sin_theta = torch.sin(theta)

directions_stack[..., 0][mask] = torch.masked_select(

coord_stack[..., 0] * sin_theta / theta, mask

).float()

directions_stack[..., 1][mask] = torch.masked_select(

coord_stack[..., 1] * sin_theta / theta, mask

).float()

directions_stack[..., 2][mask] = -torch.masked_select(torch.cos(theta), mask).float()

And this PR seems to change the radial_and_tangential_undistort function but not the second part for fisheye. I'm not sure if the formula is correct anymore.

Hi @liruilong940607, would you mind elaborating on why the old version has the correct fisheye undistortion. I'm a little confused as to how the math is correct in this case. The radial_and_tangential_undistort function implements undistortion different from the OpenCV fisheye distortion specification and it's unclear to me how the second part corrects this.

DreekFire · 2023-09-13T08:23:01Z

Hi @anc2001, nerfstudio does not use all of the distortion parameters simultaneously. For perspective cameras, it uses two radial and two tangential distortion parameters, while for fisheye cameras, it uses 4 radial distortion parameters. The old formula produces results identical to the formulas implemented in OpenCV (OpenCV uses different formulas for perspective and fisheye cameras) as long as you do not try to use the third or fourth radial distortion parameters for a perspective camera, or any tangential distortion parameters for a fisheye camera.

anc2001 · 2023-09-25T18:45:16Z

Hi @DreekFire thanks for the reply. At the time of my original comment I was unsure that the current implementation of fisheye undistortion matched up with the math described in the OpenCV documentation. I worked it out and understand why it's correct now. Are you referring to the implementation on this branch or the current implementation on the main branch? To my understanding the current implementation uses 4 radial distortion parameters and 2 tangential where the 2 tangential are just 0 when the camera is fisheye. For perspective undistortion under the OpenCV distortion specification it will be incorrect if the 4th radial distortion parameter is included, but shouldn't it be correct up to the 3rd radial distortion parameter? Also shouldn't nerfstudio support the full OpenCV perspective camera model (maybe minus the thin prism coefficients)?

DreekFire and others added 27 commits March 6, 2023 20:24

replace default colormap option with turbo for single-channel output

db4ab80

Merge branch 'nerfstudio-project:main' into main

0a458d2

Merge branch 'main' into main

94bf96d

profiling

07698df

set pose_optimizer to None if it's mode is "off"

39e6c97

radial_and_tangential_undistort in nerfacc

0cc1683

fix camera distortions

2fff930

update to nerfacc 0.5.1. Upgrade instant-ngp model.

04ccc32

revert nerfacc to 0.5.1

0054a66

format

abd7558

allows to fall back to the pytorch undistortion

cd78b78

format

b02f376

fix linting

97c411e

fix test

74bb54e

fix test

e48e932

add termination condition to camera undistort and use Jacobian to com…

cdc0c86

…pute pixel area

add profile annotations

10fba45

use explicit representation for convergence mask

9ce0ffc

fix syntax error

3204396

calculate pixel area, backwards compatibility with custom datamanager

90b00b9

fix shapes and devices

9e8dc10

Merge branch 'main' of github.com:DreekFire/nerfstudio

2b96186

testing

aa7d57c

fix resampling axes and camera ids

11945de

send coords to CPU

0579f1e

fix merge conflicts

a279ca7

fix merge conflicts again

fb4b7b8

DreekFire marked this pull request as draft June 5, 2023 12:07

add tolerance to undistort for backwards compatibility

1134da0

clean up merge and fix squeeze

fcf9aa1

tancik reviewed Jun 5, 2023

View reviewed changes

liruilong940607 reviewed Jun 6, 2023

View reviewed changes

DreekFire and others added 4 commits June 7, 2023 11:50

clean up comments and annotations, expand undistort params if necessary

1ed7258

fix rename tol to tolerance

81e2d89

select distortion parameters for each model, clean up edge cases

b8504e1

Merge branch 'main' of https://github.com/nerfstudio-project/nerfstudio

66eb324

DreekFire added 2 commits June 20, 2023 18:17

Merge branch 'main' of github.com:DreekFire/nerfstudio

8f3c9b3

Merge branch 'main' of github.com:nerfstudio-project/nerfstudio

4125c1f

DreekFire and others added 6 commits July 10, 2023 12:14

remove resample

2bd23fa

Merge branch 'main' of github.com:nerfstudio-project/nerfstudio

67d855e

fix typing

36adee7

fix merge conflicts

ca45ec2

change how tolerance is calculated

9b44940

fix typing again, remove stack

770f05c

DreekFire force-pushed the optim branch from b98a553 to 770f05c Compare August 1, 2023 09:03

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Optimizing camera undistortion #2037

Optimizing camera undistortion #2037

DreekFire commented Jun 5, 2023

DreekFire commented Jun 5, 2023 •

edited

tancik left a comment

tancik Jun 5, 2023

tancik Jun 5, 2023

tancik Jun 5, 2023

tancik Jun 5, 2023

tancik Jun 5, 2023

tancik left a comment

tancik Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 Jun 5, 2023

liruilong940607 commented Jun 6, 2023

kerrj commented Jun 13, 2023

DreekFire commented Jun 13, 2023

kerrj commented Jun 15, 2023

tancik commented Jun 24, 2023

DreekFire commented Aug 1, 2023

anc2001 commented Sep 12, 2023 •

edited

DreekFire commented Sep 13, 2023

anc2001 commented Sep 25, 2023

		max_iterations: The maximum number of iterations to perform.
		resolution: The resolution (w, h of each pixel, in units of multiples of focal length)

Optimizing camera undistortion #2037

Are you sure you want to change the base?

Optimizing camera undistortion #2037

Conversation

DreekFire commented Jun 5, 2023

DreekFire commented Jun 5, 2023 • edited

tancik left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tancik left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liruilong940607 commented Jun 6, 2023

kerrj commented Jun 13, 2023

DreekFire commented Jun 13, 2023

kerrj commented Jun 15, 2023

tancik commented Jun 24, 2023

DreekFire commented Aug 1, 2023

anc2001 commented Sep 12, 2023 • edited

DreekFire commented Sep 13, 2023

anc2001 commented Sep 25, 2023

DreekFire commented Jun 5, 2023 •

edited

anc2001 commented Sep 12, 2023 •

edited