1# Open3D: www.open3d.org
2# The MIT License (MIT)
3# See license file or visit www.open3d.org for details
4
5# following code is tested with OpenCV 3.2.0 and Python2.7
6# how to install opencv
7# conda create --prefix py27opencv python=2.7
8# source activate py27opencv
9# conda install -c conda-forge opencv
10# conda install -c conda-forge openblas (if openblas conflicts)
11import numpy as np
12import cv2
13from open3d import *
14from matplotlib import pyplot as plt # for visualizing feature matching
15import copy
16
17
18def pose_estimation(source_rgbd_image, target_rgbd_image,
19        pinhole_camera_intrinsic, debug_draw_correspondences):
20    success = False
21    trans = np.identity(4)
22
23    # transform double array to unit8 array
24    color_cv_s = np.uint8(np.asarray(source_rgbd_image.color)*255.0)
25    color_cv_t = np.uint8(np.asarray(target_rgbd_image.color)*255.0)
26
27    orb = cv2.ORB_create(scaleFactor=1.2, nlevels = 8, edgeThreshold = 31,
28            firstLevel = 0, WTA_K = 2, scoreType=cv2.ORB_HARRIS_SCORE,
29            nfeatures = 100, patchSize = 31) # to save time
30    [kp_s, des_s] = orb.detectAndCompute(color_cv_s, None)
31    [kp_t, des_t] = orb.detectAndCompute(color_cv_t, None)
32    if len(kp_s) is 0 or len(kp_t) is 0:
33        return success, trans
34
35    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
36    matches = bf.match(des_s,des_t)
37
38    pts_s = []
39    pts_t = []
40    for match in matches:
41        pts_t.append(kp_t[match.trainIdx].pt)
42        pts_s.append(kp_s[match.queryIdx].pt)
43    pts_s = np.asarray(pts_s)
44    pts_t = np.asarray(pts_t)
45    # inlier points after initial BF matching
46    if debug_draw_correspondences:
47        draw_correspondences(np.asarray(source_rgbd_image.color),
48                np.asarray(target_rgbd_image.color), pts_s, pts_t,
49                np.ones(pts_s.shape[0]), "Initial BF matching")
50
51    focal_input = (pinhole_camera_intrinsic.intrinsic_matrix[0,0] +
52            pinhole_camera_intrinsic.intrinsic_matrix[1,1]) / 2.0
53    pp_x = pinhole_camera_intrinsic.intrinsic_matrix[0,2]
54    pp_y = pinhole_camera_intrinsic.intrinsic_matrix[1,2]
55
56    # Essential matrix is made for masking inliers
57    pts_s_int = np.int32(pts_s + 0.5)
58    pts_t_int = np.int32(pts_t + 0.5)
59    [E, mask] = cv2.findEssentialMat(pts_s_int, pts_t_int, focal=focal_input,
60            pp=(pp_x, pp_y), method=cv2.RANSAC, prob=0.999, threshold=1.0)
61    if mask is None:
62        return success, trans
63
64    # inlier points after 5pt algorithm
65    if debug_draw_correspondences:
66        draw_correspondences(np.asarray(source_rgbd_image.color),
67                np.asarray(target_rgbd_image.color),
68                pts_s, pts_t, mask, "5-pt RANSAC")
69
70    # make 3D correspondences
71    depth_s = np.asarray(source_rgbd_image.depth)
72    depth_t = np.asarray(target_rgbd_image.depth)
73    pts_xyz_s = np.zeros([3, pts_s.shape[0]])
74    pts_xyz_t = np.zeros([3, pts_s.shape[0]])
75    cnt = 0
76    for i in range(pts_s.shape[0]):
77        if mask[i]:
78            xyz_s = get_xyz_from_pts(pts_s[i,:],
79                    depth_s, pp_x, pp_y, focal_input)
80            pts_xyz_s[:,cnt] = xyz_s
81            xyz_t = get_xyz_from_pts(pts_t[i,:],
82                    depth_t, pp_x, pp_y, focal_input)
83            pts_xyz_t[:,cnt] = xyz_t
84            cnt = cnt + 1
85    pts_xyz_s = pts_xyz_s[:,:cnt]
86    pts_xyz_t = pts_xyz_t[:,:cnt]
87
88    success, trans, inlier_id_vec = estimate_3D_transform_RANSAC(
89            pts_xyz_s, pts_xyz_t)
90
91    if debug_draw_correspondences:
92        pts_s_new = np.zeros(shape=(len(inlier_id_vec),2))
93        pts_t_new = np.zeros(shape=(len(inlier_id_vec),2))
94        mask = np.ones(len(inlier_id_vec))
95        cnt = 0
96        for i in inlier_id_vec:
97            u_s,v_s = get_uv_from_xyz(pts_xyz_s[0,i], pts_xyz_s[1,i],
98                    pts_xyz_s[2,i], pp_x, pp_y, focal_input)
99            u_t,v_t = get_uv_from_xyz(pts_xyz_t[0,i], pts_xyz_t[1,i],
100                    pts_xyz_t[2,i], pp_x, pp_y, focal_input)
101            pts_s_new[cnt,:] = [u_s,v_s]
102            pts_t_new[cnt,:] = [u_t,v_t]
103            cnt = cnt + 1
104        draw_correspondences(np.asarray(source_rgbd_image.color),
105                np.asarray(target_rgbd_image.color),
106                pts_s_new, pts_t_new, mask, "5-pt RANSAC + 3D Rigid RANSAC")
107    return success, trans
108
109
110def draw_correspondences(img_s, img_t, pts_s, pts_t, mask, title):
111    ha,wa = img_s.shape[:2]
112    hb,wb = img_t.shape[:2]
113    total_width = wa+wb
114    new_img = np.zeros(shape=(ha, total_width))
115    new_img[:ha,:wa]=img_s
116    new_img[:hb,wa:wa+wb]=img_t
117
118    fig = plt.figure()
119    fig.canvas.set_window_title(title)
120    for i in range(pts_s.shape[0]):
121        if mask[i]:
122            sx = pts_s[i,0]
123            sy = pts_s[i,1]
124            tx = pts_t[i,0] + wa
125            ty = pts_t[i,1]
126            plt.plot([sx,tx], [sy,ty], color=np.random.random(3)/2+0.5, lw=1.0)
127    plt.imshow(new_img)
128    plt.pause(0.5)
129    plt.close()
130
131
132def estimate_3D_transform_RANSAC(pts_xyz_s, pts_xyz_t):
133    max_iter = 1000
134    max_distance = 0.05
135    n_sample = 5
136    n_points = pts_xyz_s.shape[1]
137    Transform_good = np.identity(4)
138    max_inlier = n_sample
139    inlier_vec_good = []
140    success = False
141
142    if n_points < n_sample:
143        return False, np.identity(4), []
144
145    for i in range(max_iter):
146
147        # sampling
148        rand_idx = np.random.randint(n_points, size=n_sample)
149        sample_xyz_s = pts_xyz_s[:,rand_idx]
150        sample_xyz_t = pts_xyz_t[:,rand_idx]
151        R_approx, t_approx = estimate_3D_transform(sample_xyz_s, sample_xyz_t)
152
153        # evaluation
154        diff_mat = pts_xyz_t - (np.matmul(R_approx, pts_xyz_s) +
155                np.tile(t_approx, [1, n_points]))
156        diff = [np.linalg.norm(diff_mat[:,i]) for i in range(n_points)]
157        n_inlier = len([1 for diff_iter in diff if diff_iter < max_distance])
158
159        # note: diag(R_approx) > 0 prevents ankward transformation between
160        # RGBD pair of relatively small amount of baseline.
161        if (n_inlier > max_inlier) and (np.linalg.det(R_approx) != 0.0) and \
162                (R_approx[0,0] > 0 and R_approx[1,1] > 0 and R_approx[2,2] > 0):
163            Transform_good[:3,:3] = R_approx
164            Transform_good[:3,3] = [t_approx[0],t_approx[1],t_approx[2]]
165            max_inlier = n_inlier
166            inlier_vec = [id_iter for diff_iter, id_iter \
167                    in zip(diff, range(n_points)) \
168                    if diff_iter < max_distance]
169            inlier_vec_good = inlier_vec
170            success = True
171
172    return success, Transform_good, inlier_vec_good
173
174
175# singular value decomposition approach
176# based on the description in the sec 3.1.2 in
177# http://graphics.stanford.edu/~smr/ICP/comparison/eggert_comparison_mva97.pdf
178def estimate_3D_transform(input_xyz_s, input_xyz_t):
179    # compute H
180    xyz_s = copy.copy(input_xyz_s)
181    xyz_t = copy.copy(input_xyz_t)
182    n_points = xyz_s.shape[1]
183    mean_s = np.mean(xyz_s, axis=1)
184    mean_t = np.mean(xyz_t, axis=1)
185    mean_s.shape = (3,1)
186    mean_t.shape = (3,1)
187    xyz_diff_s = xyz_s - np.tile(mean_s, [1, n_points])
188    xyz_diff_t = xyz_t - np.tile(mean_t, [1, n_points])
189    H = np.matmul(xyz_diff_s,xyz_diff_t.transpose())
190    # solve system
191    U, s, V = np.linalg.svd(H)
192    R_approx = np.matmul(V.transpose(), U.transpose())
193    if np.linalg.det(R_approx) < 0.0:
194        det = np.linalg.det(np.matmul(U,V))
195        D = np.identity(3)
196        D[2,2] = det
197        R_approx = np.matmul(U,np.matmul(D,V))
198    t_approx = mean_t - np.matmul(R_approx, mean_s)
199    return R_approx, t_approx
200
201
202def get_xyz_from_pts(pts_row, depth, px, py, focal):
203    u = pts_row[0]
204    v = pts_row[1]
205    u0 = int(u)
206    v0 = int(v)
207    height = depth.shape[0]
208    width = depth.shape[1]
209    # bilinear depth interpolation
210    if u0 > 0 and u0 < width-1 and v0 > 0 and v0 < height-1:
211        up = pts_row[0] - u0
212        vp = pts_row[1] - v0
213        d0 = depth[v0, u0]
214        d1 = depth[v0, u0+1]
215        d2 = depth[v0+1, u0]
216        d3 = depth[v0+1, u0]
217        d = (1-vp) * (d1 * up + d0 * (1-up)) + vp * (d3 * up + d2 * (1-up))
218        return get_xyz_from_uv(u, v, d, px, py, focal)
219    else:
220        return [0,0,0]
221
222
223def get_xyz_from_uv(u, v, d, px, py, focal):
224    if focal != 0:
225        x = (u - px) / focal * d
226        y = (v - py) / focal * d
227    else:
228        x = 0
229        y = 0
230    return np.array([x, y, d]).transpose()
231
232
233def get_uv_from_xyz(x, y, z, px, py, focal):
234    if z != 0:
235        u = focal * x / z + px
236        v = focal * y / z + py
237    else:
238        u = 0
239        v = 0
240    return u, v
241