Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| # def project_to_image( | |
| # pts: torch.Tensor, calib: dict[str, Any] | |
| # ) -> tuple[torch.Tensor, torch.Tensor]: | |
| # """Project pts in camera coordinates into image coordinates. | |
| # Args: | |
| # pts (torch.Tensor): B, n_views, n_pts, 3 | |
| # Ks (torch.Tensor): B, n_views, 3, 3 | |
| # Returns: | |
| # tuple[torch.Tensor, torch.Tensor]: (B, n_views, n_pts, 2), (B, n_views, n_pts, 1) | |
| # """ | |
| # pts = pts / torch.norm(pts, dim=-1, keepdim=True) | |
| # x = pts[:, 0] | |
| # y = pts[:, 1] | |
| # z = pts[:, 2] | |
| # xi_src = calib["mirror_parameters"]["xi"] | |
| # x = x / (z + xi_src) | |
| # y = y / (z + xi_src) | |
| # k1 = calib["distortion_parameters"]["k1"] | |
| # k2 = calib["distortion_parameters"]["k2"] | |
| # r = x * x + y * y | |
| # factor = 1 + k1 * r + k2 * r * r | |
| # x = x * factor | |
| # y = y * factor | |
| # gamma0 = calib["projection_parameters"]["gamma1"] | |
| # gamma1 = calib["projection_parameters"]["gamma2"] | |
| # u0 = calib["projection_parameters"]["u0"] | |
| # v0 = calib["projection_parameters"]["v0"] | |
| # x = x * gamma0 + u0 | |
| # y = y * gamma1 + v0 | |
| # return torch.stack([x, y], dim=-1), z | |
| # TODO: lookup | |
| EPS = 1.0e-6 | |
| def normalize_calib(calib: torch.Tensor, img_sizes: torch.Tensor) -> torch.Tensor: | |
| """Normalize the calibration matrices for fisheye cameras based on the image size | |
| Args: | |
| calib (torch.Tensor): B, n_views, 7, [xi, k1, k2, gamma1, gamma2, u0, v0] | |
| img_sizes (torch.Tensor): B, n_views, 2 | |
| Returns: | |
| torch.Tensor: B, n_views 7 | |
| """ | |
| calib[..., 3:5] = calib[..., 3:5] / img_sizes * 2.0 | |
| calib[..., 5:7] = calib[..., 5:7] / img_sizes * 2.0 - 1.0 | |
| return calib | |
| def unnormalize_calib(calib: torch.Tensor, img_sizes: torch.Tensor) -> torch.Tensor: | |
| """Unnormalize the calibration matrices for fisheye cameras based on the image size | |
| Args: | |
| calib (torch.Tensor): B, n_views, 7, [xi, k1, k2, gamma1, gamma2, u0, v0] | |
| img_sizes (torch.Tensor): B, n_views, 2 | |
| Returns: | |
| torch.Tensor: B, n_views 7 | |
| """ | |
| calib[..., 3:5] = calib[..., 3:5] * img_sizes / 2.0 | |
| calib[..., 5:7] = (calib[..., 5:7] + 1.0) * img_sizes / 2.0 | |
| return calib | |
| def project_to_image( | |
| pts: torch.Tensor, calib: torch.Tensor | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Project pts in camera coordinates into image coordinates. | |
| Args: | |
| pts (torch.Tensor): B, n_views, n_pts, 3 | |
| Ks (torch.Tensor): B, n_views, 7 | |
| Returns: | |
| tuple[torch.Tensor, torch.Tensor]: (B, n_views, n_pts, 2), (B, n_views, n_pts, 1) | |
| """ | |
| pts = pts / torch.norm(pts, dim=-1, keepdim=True) | |
| xy = pts[..., 0:2] | |
| z = pts[..., 2:3] | |
| xi_src = calib[..., 0:1].unsqueeze(-2) # B, n_views, 1, 1 | |
| xy = xy / (z + xi_src) | |
| r = torch.sum(torch.square(xy), dim=-1) | |
| factor = 1 + calib[..., 1:2] * r + calib[..., 2:3] * torch.square(r) | |
| xy = xy * factor.unsqueeze(-1) | |
| xy = xy * calib[..., 3:5].unsqueeze(-2) + calib[..., 5:7].unsqueeze(-2) | |
| return xy, z | |
| # TODO: define for fisheye | |
| def outside_frustum( | |
| xy: torch.Tensor, | |
| z: torch.Tensor, | |
| limits_x: tuple[float, float] | tuple[int, int] = (-1.0, 1.0), | |
| limits_y: tuple[float, float] | tuple[int, int] = (-1.0, 1.0), | |
| limit_z: float = EPS, | |
| ) -> torch.Tensor: | |
| """_summary_ | |
| Args: | |
| xy (torch.Tensor): _description_ | |
| z (torch.Tensor): _description_ | |
| limits_x (tuple[float, float] | tuple[int, int], optional): _description_. Defaults to (-1.0, 1.0). | |
| limits_y (tuple[float, float] | tuple[int, int], optional): _description_. Defaults to (-1.0, 1.0). | |
| limit_z (float, optional): _description_. Defaults to EPS. | |
| Returns: | |
| torch.Tensor: _description_ | |
| """ | |
| return ( | |
| (z <= limit_z) | |
| | (xy[..., :1] < limits_x[0]) | |
| | (xy[..., :1] > limits_x[1]) | |
| | (xy[..., 1:2] < limits_y[0]) | |
| | (xy[..., 1:2] > limits_y[1]) | |
| ) | |