stereo_match.py 2.18 KB
Newer Older
wester committed
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
#!/usr/bin/env python

'''
Simple example of stereo image matching and point cloud generation.

Resulting .ply file cam be easily viewed using MeshLab ( http://meshlab.sourceforge.net/ )
'''

import numpy as np
import cv2

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])
wester committed
28 29 30
    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')
wester committed
31 32 33


if __name__ == '__main__':
wester committed
34 35 36
    print 'loading images...'
    imgL = cv2.pyrDown( cv2.imread('../gpu/aloeL.jpg') )  # downscale images for faster processing
    imgR = cv2.pyrDown( cv2.imread('../gpu/aloeR.jpg') )
wester committed
37 38 39 40 41

    # disparity range is tuned for 'aloe' image pair
    window_size = 3
    min_disp = 16
    num_disp = 112-min_disp
wester committed
42
    stereo = cv2.StereoSGBM(minDisparity = min_disp,
wester committed
43
        numDisparities = num_disp,
wester committed
44
        SADWindowSize = window_size,
wester committed
45 46
        uniquenessRatio = 10,
        speckleWindowSize = 100,
wester committed
47 48 49 50 51
        speckleRange = 32,
        disp12MaxDiff = 1,
        P1 = 8*3*window_size**2,
        P2 = 32*3*window_size**2,
        fullDP = False
wester committed
52 53
    )

wester committed
54
    print 'computing disparity...'
wester committed
55 56
    disp = stereo.compute(imgL, imgR).astype(np.float32) / 16.0

wester committed
57
    print 'generating 3d point cloud...',
wester committed
58 59 60 61 62 63 64 65 66 67 68 69 70
    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]
    out_fn = 'out.ply'
    write_ply('out.ply', out_points, out_colors)
wester committed
71
    print '%s saved' % 'out.ply'
wester committed
72 73 74 75 76

    cv2.imshow('left', imgL)
    cv2.imshow('disparity', (disp-min_disp)/num_disp)
    cv2.waitKey()
    cv2.destroyAllWindows()