samples/python2/stereo_match.py

   1 #!/usr/bin/env python
   2
   3 '''
   4 Simple example of stereo image matching and point cloud generation.
   5
   6 Resulting .ply file cam be easily viewed using MeshLab ( http://meshlab.sourceforge.net/ )
   7 '''
   8
   9 import numpy as np
  10 import cv2
  11
  12 ply_header = '''ply
  13 format ascii 1.0
  14 element vertex %(vert_num)d
  15 property float x
  16 property float y
  17 property float z
  18 property uchar red
  19 property uchar green
  20 property uchar blue
  21 end_header
  22 '''
  23
  24 def write_ply(fn, verts, colors):
  25     verts = verts.reshape(-1, 3)
  26     colors = colors.reshape(-1, 3)
  27     verts = np.hstack([verts, colors])
  28     with open(fn, 'w') as f:
  29         f.write(ply_header % dict(vert_num=len(verts)))
  30         np.savetxt(f, verts, '%f %f %f %d %d %d')
  31
  32
  33 if __name__ == '__main__':
  34     print 'loading images...'
  35     imgL = cv2.pyrDown( cv2.imread('../gpu/aloeL.jpg') )  # downscale images for faster processing
  36     imgR = cv2.pyrDown( cv2.imread('../gpu/aloeR.jpg') )
  37
  38     # disparity range is tuned for 'aloe' image pair
  39     window_size = 3
  40     min_disp = 16
  41     num_disp = 112-min_disp
  42     stereo = cv2.StereoSGBM(minDisparity = min_disp,
  43         numDisparities = num_disp,
  44         SADWindowSize = window_size,
  45         uniquenessRatio = 10,
  46         speckleWindowSize = 100,
  47         speckleRange = 32,
  48         disp12MaxDiff = 1,
  49         P1 = 8*3*window_size**2,
  50         P2 = 32*3*window_size**2,
  51         fullDP = False
  52     )
  53
  54     print 'computing disparity...'
  55     disp = stereo.compute(imgL, imgR).astype(np.float32) / 16.0
  56
  57     print 'generating 3d point cloud...',
  58     h, w = imgL.shape[:2]
  59     f = 0.8*w                          # guess for focal length
  60     Q = np.float32([[1, 0, 0, -0.5*w],
  61                     [0,-1, 0,  0.5*h], # turn points 180 deg around x-axis,
  62                     [0, 0, 0,     -f], # so that y-axis looks up
  63                     [0, 0, 1,      0]])
  64     points = cv2.reprojectImageTo3D(disp, Q)
  65     colors = cv2.cvtColor(imgL, cv2.COLOR_BGR2RGB)
  66     mask = disp > disp.min()
  67     out_points = points[mask]
  68     out_colors = colors[mask]
  69     out_fn = 'out.ply'
  70     write_ply('out.ply', out_points, out_colors)
  71     print '%s saved' % 'out.ply'
  72
  73     cv2.imshow('left', imgL)
  74     cv2.imshow('disparity', (disp-min_disp)/num_disp)
  75     cv2.waitKey()
  76     cv2.destroyAllWindows()