112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
import numpy as np
|
|
import cv2
|
|
import os
|
|
|
|
def generate_point_cloud(img_left_path, img_right_path, output_path='out.ply'):
|
|
# Load images
|
|
print('loading images...')
|
|
imgL = cv2.pyrDown(cv2.imread(img_left_path)) # downscale images for faster processing
|
|
imgR = cv2.pyrDown(cv2.imread(img_right_path))
|
|
|
|
# Disparity parameters tuned for 'aloe' image pair
|
|
window_size = 3
|
|
min_disp = 16
|
|
num_disp = 112 - min_disp
|
|
stereo = cv2.StereoSGBM(
|
|
minDisparity=min_disp,
|
|
numDisparities=num_disp,
|
|
SADWindowSize=window_size,
|
|
uniquenessRatio=10,
|
|
speckleWindowSize=100,
|
|
speckleRange=32,
|
|
disp12MaxDiff=1,
|
|
P1=8 * 3 * window_size**2,
|
|
P2=32 * 3 * window_size**2,
|
|
fullDP=False
|
|
)
|
|
|
|
# Compute disparity
|
|
print('computing disparity...')
|
|
disp = stereo.compute(imgL, imgR).astype(np.float32) / 16.0
|
|
|
|
# Generate 3D point cloud
|
|
print('generating 3d point cloud...')
|
|
h, w = imgL.shape[:2]
|
|
f = 0.8 * w # Guess for focal length
|
|
Q = np.float32([[1, 0, 0, -0.5 * w],
|
|
[0, -1, 0, 0.5 * h], # Turn points 180 deg around x-axis
|
|
[0, 0, 0, -f], # so that y-axis looks up
|
|
[0, 0, 1, 0]])
|
|
points = cv2.reprojectImageTo3D(disp, Q)
|
|
colors = cv2.cvtColor(imgL, cv2.COLOR_BGR2RGB)
|
|
mask = disp > disp.min()
|
|
out_points = points[mask]
|
|
out_colors = colors[mask]
|
|
|
|
# Write point cloud to file
|
|
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
|
|
'''
|
|
|
|
def write_ply(fn, verts, colors):
|
|
verts = verts.reshape(-1, 3)
|
|
colors = colors.reshape(-1, 3)
|
|
verts = np.hstack([verts, colors])
|
|
with open(fn, 'w') as f:
|
|
f.write(ply_header % dict(vert_num=len(verts)))
|
|
np.savetxt(f, verts, '%f %f %f %d %d %d')
|
|
|
|
write_ply(output_path, out_points, out_colors)
|
|
print('%s saved' % output_path)
|
|
|
|
# Display results
|
|
cv2.imshow('left', imgL)
|
|
cv2.imshow('disparity', (disp - min_disp) / num_disp)
|
|
cv2.waitKey()
|
|
cv2.destroyAllWindows()
|
|
|
|
def capture_stereo_images(cam1_id, cam2_id, num_images=10, output_folder='captured_images'):
|
|
# Initialize camera objects
|
|
cap1 = cv2.VideoCapture(cam1_id)
|
|
cap2 = cv2.VideoCapture(cam2_id)
|
|
|
|
# Check if cameras are opened successfully
|
|
if not (cap1.isOpened() and cap2.isOpened()):
|
|
print("Error: Could not open cameras")
|
|
return
|
|
|
|
# Create output folder if it doesn't exist
|
|
if not os.path.exists(output_folder):
|
|
os.makedirs(output_folder)
|
|
|
|
# Capture stereo image pairs
|
|
for i in range(num_images):
|
|
ret1, frame1 = cap1.read()
|
|
ret2, frame2 = cap2.read()
|
|
|
|
if not (ret1 and ret2):
|
|
print("Error: Could not read frames")
|
|
break
|
|
|
|
# Save images
|
|
cv2.imwrite(f"{output_folder}/left_{i}.jpg", frame1)
|
|
cv2.imwrite(f"{output_folder}/right_{i}.jpg", frame2)
|
|
|
|
print(f"Captured pair {i + 1}/{num_images}")
|
|
|
|
# Release video captures
|
|
cap1.release()
|
|
cap2.release()
|
|
|
|
|
|
capture_stereo_images(1, 2, num_images=10, output_folder='captured_images')
|
|
|
|
generate_point_cloud('path_to_left_img.jpg', 'path_to_right_img.jpg') |