File size: 3,941 Bytes
9e15541
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
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])
    )