stereo_match.py 2.3 KB

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