3d-geometry-on-stereoids/project_point.py at master · hiteshhedwig/3d-geometry-on-stereoids · GitHub

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
import numpy as np

# projection considering distortion
# idea is
# Pw -> Pc
# normalize it with last vector of Pc
# normalized vector [x`, y`, 1]
# apply distortion
# now pass through K matrix
# now we have - pixel coordinate


def project_point(R : np.ndarray, T : np.ndarray, P_w : np.ndarray, K : np.ndarray, k_1 , k_2, p1, p2) :
    # Pw -> Pc
    P_c = R @ P_w + T

    # normalize
    P_c_norm = [P_c[0]/P_c[2], P_c[1]/P_c[2], 1]

    # apply distortion - x` = x * S_r + Delx_t <- tangential
    # 1. radial distortion only for now
    x = P_c_norm[0]
    y = P_c_norm[1]
    r_2 = x**2 + y**2

    # tangential dist -
    del_xt, del_yt = get_tangential_dist (x, y, p1, p2)

    x_d = x * (1 + k_1 * r_2 + k_2 * r_2**2) + del_xt
    y_d = y * (1 + k_1 * r_2 + k_2 * r_2**2) + del_yt


    P_d = [x_d, y_d, 1].T

    pixel_point = K @ P_d
    return pixel_point


def get_tangential_dist(x, y, p1, p2):
    # tangential distortion
    # del_xt = 2p1 xy + p2(r^2 + 2x^2)
    # del_yt = p1(r^2 + 2y^2) + 2p2 xy
    r2 = x**2 + y**2

    del_xt = 2*p1*x*y + p2*(r2 + 2*x**2)
    del_yt = p1*(r2 + 2*y**2) + 2*p2*x*y

    return del_xt, del_yt


def get_radial_dist(x, y, k1, k2) :
    r2 = x**2 + y**2
    return (1 + k1*r2 + k2*r2**2), (1 + k1*r2 + k2*r2**2)


import numpy as np

def distort_forward(x, y, k1, k2, p1, p2):
    # radius of the assunmed sphere
    r2 = x*x + y*y
    # scale of the radial distortion
    s  = 1.0 + k1*r2 + k2*(r2*r2)
    # scale of tangential distortion
    dx_t = 2*p1*x*y + p2*(r2 + 2*x*x)
    dy_t = p1*(r2 + 2*y*y) + 2*p2*x*y
    # overall distortion
    xd = x*s + dx_t
    yd = y*s + dy_t
    return xd, yd

def undistort_normalized(xd, yd, k1, k2, p1, p2, iters=5):
    # start from distorted as initial guess
    x, y = xd, yd
    for _ in range(iters):
        r2 = x*x + y*y
        s  = 1.0 + k1*r2 + k2*(r2*r2)
        dx_t = 2*p1*x*y + p2*(r2 + 2*x*x)
        dy_t = p1*(r2 + 2*y*y) + 2*p2*x*y
        x = (xd - dx_t) / s
        y = (yd - dy_t) / s
    return x, y


def lidar_3d_to_camera(P_L, R_CL, t_CL, K, k1, k2, p1, p2) :
    # points in camera coordinate frame
    P_c = R_CL @ P_L + t_CL

    # condition for z - 1.must be greater than zero. and skip
    # too close values to zero too
    Z = P_c[2]
    if ( Z == 0 or np.abs(Z) > 1e-9):
        ValueError("Cant be zero bro")

    # normalize
    P_norm = np.array([P_c[0]/P_c[2], P_c[1]/P_c[2], 1])

    # now, distort
    xd, yd = distort_forward(P_norm[0], P_norm[1], k1, k2,
                    p1, p2)

    # now, project it using K matrix
    uv_cam = K @ np.array([xd, yd])
    return uv_cam