Iteratively undistort the normalized tracks using the given distortion parameters. Args: params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN. tracks_normalized (torch.Tensor or numpy.ndarray): Normalized tracks tensor of shape [batch_size, num_tracks,
(params, tracks_normalized, max_iterations=100, max_step_norm=1e-10, rel_step_size=1e-6)
| 49 | |
| 50 | |
| 51 | def iterative_undistortion(params, tracks_normalized, max_iterations=100, max_step_norm=1e-10, rel_step_size=1e-6): |
| 52 | """ |
| 53 | Iteratively undistort the normalized tracks using the given distortion parameters. |
| 54 | |
| 55 | Args: |
| 56 | params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN. |
| 57 | tracks_normalized (torch.Tensor or numpy.ndarray): Normalized tracks tensor of shape [batch_size, num_tracks, 2]. |
| 58 | max_iterations (int): Maximum number of iterations for the undistortion process. |
| 59 | max_step_norm (float): Maximum step norm for convergence. |
| 60 | rel_step_size (float): Relative step size for numerical differentiation. |
| 61 | |
| 62 | Returns: |
| 63 | torch.Tensor: Undistorted normalized tracks tensor. |
| 64 | """ |
| 65 | params = _ensure_torch(params) |
| 66 | tracks_normalized = _ensure_torch(tracks_normalized) |
| 67 | |
| 68 | B, N, _ = tracks_normalized.shape |
| 69 | u, v = tracks_normalized[..., 0].clone(), tracks_normalized[..., 1].clone() |
| 70 | original_u, original_v = u.clone(), v.clone() |
| 71 | |
| 72 | eps = torch.finfo(u.dtype).eps |
| 73 | for idx in range(max_iterations): |
| 74 | u_undist, v_undist = apply_distortion(params, u, v) |
| 75 | dx = original_u - u_undist |
| 76 | dy = original_v - v_undist |
| 77 | |
| 78 | step_u = torch.clamp(torch.abs(u) * rel_step_size, min=eps) |
| 79 | step_v = torch.clamp(torch.abs(v) * rel_step_size, min=eps) |
| 80 | |
| 81 | J_00 = (apply_distortion(params, u + step_u, v)[0] - apply_distortion(params, u - step_u, v)[0]) / (2 * step_u) |
| 82 | J_01 = (apply_distortion(params, u, v + step_v)[0] - apply_distortion(params, u, v - step_v)[0]) / (2 * step_v) |
| 83 | J_10 = (apply_distortion(params, u + step_u, v)[1] - apply_distortion(params, u - step_u, v)[1]) / (2 * step_u) |
| 84 | J_11 = (apply_distortion(params, u, v + step_v)[1] - apply_distortion(params, u, v - step_v)[1]) / (2 * step_v) |
| 85 | |
| 86 | J = torch.stack([torch.stack([J_00 + 1, J_01], dim=-1), torch.stack([J_10, J_11 + 1], dim=-1)], dim=-2) |
| 87 | |
| 88 | delta = torch.linalg.solve(J, torch.stack([dx, dy], dim=-1)) |
| 89 | |
| 90 | u += delta[..., 0] |
| 91 | v += delta[..., 1] |
| 92 | |
| 93 | if torch.max((delta**2).sum(dim=-1)) < max_step_norm: |
| 94 | break |
| 95 | |
| 96 | return torch.stack([u, v], dim=-1) |
| 97 | |
| 98 | |
| 99 | def apply_distortion(extra_params, u, v): |
no test coverage detected