-
Notifications
You must be signed in to change notification settings - Fork 589
/
pointCloudStereoVision.py
138 lines (109 loc) · 3.91 KB
/
pointCloudStereoVision.py
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
138
import cv2
import numpy as np
import glob
from tqdm import tqdm
import PIL.ExifTags
import PIL.Image
from matplotlib import pyplot as plt
#=====================================
# Function declarations
#=====================================
#Function to create point cloud file
def create_output(vertices, colors, filename):
colors = colors.reshape(-1,3)
vertices = np.hstack([vertices.reshape(-1,3),colors])
ply_header = '''ply
format ascii 1.0
element vertex %(vert_num)d
property float x
property float y
property float z
property uchar red
property uchar green
property uchar blue
end_header
'''
with open(filename, 'w') as f:
f.write(ply_header %dict(vert_num=len(vertices)))
np.savetxt(f,vertices,'%f %f %f %d %d %d')
#Function that Downsamples image x number (reduce_factor) of times.
def downsample_image(image, reduce_factor):
for i in range(0,reduce_factor):
#Check if image is color or grayscale
if len(image.shape) > 2:
row,col = image.shape[:2]
else:
row,col = image.shape
image = cv2.pyrDown(image, dstsize= (col//2, row // 2))
return image
#=========================================================
# Stereo 3D reconstruction
#=========================================================
# Camera parameters to undistort and rectify images
cv_file = cv2.FileStorage()
cv_file.open('stereoMap.xml', cv2.FileStorage_READ)
stereoMapL_x = cv_file.getNode('stereoMapL_x').mat()
stereoMapL_y = cv_file.getNode('stereoMapL_y').mat()
stereoMapR_x = cv_file.getNode('stereoMapR_x').mat()
stereoMapR_y = cv_file.getNode('stereoMapR_y').mat()
Q = cv_file.getNode('q').mat()
#print(Q)
imgL = cv2.imread('images/left/imageL0.png', cv2.IMREAD_GRAYSCALE)
imgR = cv2.imread('images/right/imageR0.png', cv2.IMREAD_GRAYSCALE)
# Show the frames
cv2.imshow("frame right", imgR)
cv2.imshow("frame left", imgL)
cv2.waitKey(0)
# Undistort and rectify images
imgR = cv2.remap(imgR, stereoMapR_x, stereoMapR_y, cv2.INTER_LANCZOS4, cv2.BORDER_CONSTANT, 0)
imgL = cv2.remap(imgL, stereoMapL_x, stereoMapL_y, cv2.INTER_LANCZOS4, cv2.BORDER_CONSTANT, 0)
# Show the frames
cv2.imshow("frame right", imgR)
cv2.imshow("frame left", imgL)
#Downsample each image 3 times (because they're too big)
imgL = downsample_image(imgL,2)
imgR = downsample_image(imgR,2)
#stereo = cv2.StereoBM_create(numDisparities=32, blockSize=9)
# For each pixel algorithm will find the best disparity from 0
# Larger block size implies smoother, though less accurate disparity map
#disparity = stereo.compute(imgL, imgR)
#Set disparity parameters
#Note: disparity range is tuned according to specific parameters obtained through trial and error.
win_size = 5
min_disp = -1
max_disp = 63 #min_disp * 9
num_disp = max_disp - min_disp # Needs to be divisible by 16
#Create Block matching object.
stereo = cv2.StereoSGBM_create(minDisparity= min_disp,
numDisparities = num_disp,
blockSize = 5,
uniquenessRatio = 5,
speckleWindowSize = 5,
speckleRange = 5,
disp12MaxDiff = 2,
P1 = 8*3*win_size**2,#8*3*win_size**2,
P2 =32*3*win_size**2) #32*3*win_size**2)
#Compute disparity map
print ("\nComputing the disparity map...")
disparity_map = stereo.compute(imgL, imgR)
#Show disparity map before generating 3D cloud to verify that point cloud will be usable.
plt.imshow(disparity_map,'gray')
plt.show()
#Generate point cloud.
print ("\nGenerating the 3D map...")
#Get new downsampled width and height
h,w = imgR.shape[:2]
#Reproject points into 3D
points_3D = cv2.reprojectImageTo3D(disparity_map, Q)
#Get color points
colors = cv2.cvtColor(imgL, cv2.COLOR_BGR2RGB)
#Get rid of points with value 0 (i.e no depth)
mask_map = disparity_map > disparity_map.min()
#Mask colors and points.
output_points = points_3D[mask_map]
output_colors = colors[mask_map]
#Define name for output file
output_file = 'reconstructed.ply'
#Generate point cloud
print ("\n Creating the output file... \n")
create_output(output_points, output_colors, output_file)