Warp kpts0 from I0 to I1 with depth, K and Rt Also check covisibility and depth consistency. Depth is consistent if relative error < 0.2 (hard-coded). # https://github.com/zju3dv/LoFTR/blob/94e98b695be18acb43d5d3250f52226a8e36f839/src/loftr/utils/geometry.py adapted from here Args:
(kpts0, depth0, depth1, T_0to1, K0, K1, smooth_mask = False, return_relative_depth_error = False, depth_interpolation_mode = "bilinear", relative_depth_error_threshold = 0.05)
| 145 | |
| 146 | @torch.no_grad() |
| 147 | def warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, smooth_mask = False, return_relative_depth_error = False, depth_interpolation_mode = "bilinear", relative_depth_error_threshold = 0.05): |
| 148 | """Warp kpts0 from I0 to I1 with depth, K and Rt |
| 149 | Also check covisibility and depth consistency. |
| 150 | Depth is consistent if relative error < 0.2 (hard-coded). |
| 151 | # https://github.com/zju3dv/LoFTR/blob/94e98b695be18acb43d5d3250f52226a8e36f839/src/loftr/utils/geometry.py adapted from here |
| 152 | Args: |
| 153 | kpts0 (torch.Tensor): [N, L, 2] - <x, y>, should be normalized in (-1,1) |
| 154 | depth0 (torch.Tensor): [N, H, W], |
| 155 | depth1 (torch.Tensor): [N, H, W], |
| 156 | T_0to1 (torch.Tensor): [N, 3, 4], |
| 157 | K0 (torch.Tensor): [N, 3, 3], |
| 158 | K1 (torch.Tensor): [N, 3, 3], |
| 159 | Returns: |
| 160 | calculable_mask (torch.Tensor): [N, L] |
| 161 | warped_keypoints0 (torch.Tensor): [N, L, 2] <x0_hat, y1_hat> |
| 162 | """ |
| 163 | ( |
| 164 | n, |
| 165 | h, |
| 166 | w, |
| 167 | ) = depth0.shape |
| 168 | if depth_interpolation_mode == "combined": |
| 169 | # Inspired by approach in inloc, try to fill holes from bilinear interpolation by nearest neighbour interpolation |
| 170 | if smooth_mask: |
| 171 | raise NotImplementedError("Combined bilinear and NN warp not implemented") |
| 172 | valid_bilinear, warp_bilinear = warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, |
| 173 | smooth_mask = smooth_mask, |
| 174 | return_relative_depth_error = return_relative_depth_error, |
| 175 | depth_interpolation_mode = "bilinear", |
| 176 | relative_depth_error_threshold = relative_depth_error_threshold) |
| 177 | valid_nearest, warp_nearest = warp_kpts(kpts0, depth0, depth1, T_0to1, K0, K1, |
| 178 | smooth_mask = smooth_mask, |
| 179 | return_relative_depth_error = return_relative_depth_error, |
| 180 | depth_interpolation_mode = "nearest-exact", |
| 181 | relative_depth_error_threshold = relative_depth_error_threshold) |
| 182 | nearest_valid_bilinear_invalid = (~valid_bilinear).logical_and(valid_nearest) |
| 183 | warp = warp_bilinear.clone() |
| 184 | warp[nearest_valid_bilinear_invalid] = warp_nearest[nearest_valid_bilinear_invalid] |
| 185 | valid = valid_bilinear | valid_nearest |
| 186 | return valid, warp |
| 187 | |
| 188 | |
| 189 | kpts0_depth = F.grid_sample(depth0[:, None], kpts0[:, :, None], mode = depth_interpolation_mode, align_corners=False)[ |
| 190 | :, 0, :, 0 |
| 191 | ] |
| 192 | kpts0 = torch.stack( |
| 193 | (w * (kpts0[..., 0] + 1) / 2, h * (kpts0[..., 1] + 1) / 2), dim=-1 |
| 194 | ) # [-1+1/h, 1-1/h] -> [0.5, h-0.5] |
| 195 | # Sample depth, get calculable_mask on depth != 0 |
| 196 | # nonzero_mask = kpts0_depth != 0 |
| 197 | # Sample depth, get calculable_mask on depth > 0 |
| 198 | nonzero_mask = kpts0_depth > 0 |
| 199 | |
| 200 | # Unproject |
| 201 | kpts0_h = ( |
| 202 | torch.cat([kpts0, torch.ones_like(kpts0[:, :, [0]])], dim=-1) |
| 203 | * kpts0_depth[..., None] |
| 204 | ) # (N, L, 3) |