added stereo_match.py sample
authorAlexander Mordvintsev <no@email>
Sun, 26 Jun 2011 21:34:39 +0000 (21:34 +0000)
committerAlexander Mordvintsev <no@email>
Sun, 26 Jun 2011 21:34:39 +0000 (21:34 +0000)
samples/python2/stereo_match.py [new file with mode: 0644]

diff --git a/samples/python2/stereo_match.py b/samples/python2/stereo_match.py
new file mode 100644 (file)
index 0000000..842298e
--- /dev/null
@@ -0,0 +1,73 @@
+'''\r
+Simple example of stereo image matching and point cloud generation.\r
+\r
+Resulting .ply file cam be easily viewed using MeshLab (http://meshlab.sourceforge.net/)\r
+'''\r
+\r
+import numpy as np\r
+import cv2, cv\r
+\r
+ply_header = '''ply\r
+format ascii 1.0\r
+element vertex %(vert_num)d\r
+property float x\r
+property float y\r
+property float z\r
+property uchar red\r
+property uchar green\r
+property uchar blue\r
+end_header\r
+'''\r
+\r
+def write_ply(fn, verts, colors):\r
+    verts = verts.reshape(-1, 3)\r
+    colors = colors.reshape(-1, 3)\r
+    verts = np.hstack([verts, colors])\r
+    with open(fn, 'w') as f:\r
+        f.write(ply_header % dict(vert_num=len(verts)))\r
+        np.savetxt(f, verts, '%f %f %f %d %d %d')\r
+\r
+\r
+if __name__ == '__main__':\r
+    print 'loading images...'\r
+    imgL = cv2.pyrDown( cv2.imread('../gpu/aloeL.jpg') )  # downscale images for faster processing\r
+    imgR = cv2.pyrDown( cv2.imread('../gpu/aloeR.jpg') )\r
+\r
+    # disparity range is tuned for 'aloe' image pair\r
+    window_size = 3\r
+    min_disp = 16\r
+    num_disp = 112-min_disp\r
+    stereo = cv2.StereoSGBM(minDisparity = min_disp, \r
+        numDisparities = num_disp, \r
+        SADWindowSize = window_size,\r
+        uniquenessRatio = 10,\r
+        speckleWindowSize = 100,\r
+        speckleRange = 32,\r
+        disp12MaxDiff = 1,\r
+        P1 = 8*3*window_size**2,\r
+        P2 = 32*3*window_size**2,\r
+        fullDP = False\r
+    )\r
+\r
+    print 'computing disparity...'\r
+    disp = stereo.compute(imgL, imgR).astype(np.float32) / 16.0\r
+    \r
+    print 'generating 3d point cloud...',\r
+    h, w = imgL.shape[:2]\r
+    f = 0.8*w                          # guess for focal length\r
+    Q = np.float32([[1, 0, 0, -0.5*w],\r
+                    [0,-1, 0,  0.5*h], # turn points 180 deg around x-axis, \r
+                    [0, 0, 0,     -f], # so that y-axis looks up\r
+                    [0, 0, 1,      0]])\r
+    points = cv2.reprojectImageTo3D(disp, Q)\r
+    colors = cv2.cvtColor(imgL, cv.CV_BGR2RGB)\r
+    mask = disp > disp.min()\r
+    out_points = points[mask]\r
+    out_colors = colors[mask]\r
+    out_fn = 'out.ply'\r
+    write_ply('out.ply', out_points, out_colors)\r
+    print '%s saved' % 'out.ply'\r
+\r
+    cv2.imshow('left', imgL)\r
+    cv2.imshow('disparity', (disp-min_disp)/num_disp)\r
+    cv2.waitKey()
\ No newline at end of file