#!/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/ )
'''

# Python 2/3 compatibility
from __future__ import print_function

import numpy as np
import cv2 as cv

right = cv.VideoCapture(1)
left = cv.VideoCapture(0)

OffsetX = 0.0
OffsetY = -13.3265386728
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 rotateImage(image, angle):
    (h, w) = image.shape[:2]
    center = (w / 2, h / 2)
    M = cv.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv.warpAffine(image, M, (w, h))
    return rotated

def imageoffset(image, offsetx, offsety):
    img = image
    rows, cols = img.shape[:2]
    M = np.float32([[1, 0, offsetx], [0, 1, offsety]])
    dst = cv.warpAffine(img, M, (cols, rows))
    #cv2.imshow("dst", dst)
    return dst
def rotateImage(image, angle):
    (h, w) = image.shape[:2]
    center = (w / 2, h / 2)
    M = cv.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv.warpAffine(image, M, (w, h))
    return rotated

def imageoffset(image, offsetx, offsety):
    img = image
    rows, cols = img.shape[:2]
    M = np.float32([[1, 0, offsetx], [0, 1, offsety]])
    dst = cv.warpAffine(img, M, (cols, rows))
    #cv2.imshow("dst", dst)
    return dst
def get_image1():
 retval, im = left.read()
 return im
def get_image2():
 retval, im = right.read()
 im1 = rotateImage(im, 180)
 return im1


def write_ply(fn, verts, colors):
    verts = verts.reshape(-1, 3)
    colors = colors.reshape(-1, 3)
    verts = np.hstack([verts, colors])
    with open(fn, 'wb') as f:
        f.write((ply_header % dict(vert_num=len(verts))).encode('utf-8'))
        np.savetxt(f, verts, fmt='%f %f %f %d %d %d ')


if __name__ == '__main__':
    while True:
        imgL = get_image1()
        imgR = get_image2()
        imgR = imageoffset(imgR, OffsetX, OffsetY)

        # disparity range is tuned for 'aloe' image pair
        window_size = 2
        min_disp = (0 * 16)
        num_disp = (12 * 16)-min_disp
        stereo = cv.StereoSGBM_create(minDisparity = min_disp,
            numDisparities = num_disp,
            blockSize = 5,
            P1 = 8*3*window_size**2,
            P2 = 32*3*window_size**2,
            disp12MaxDiff = 10,
            uniquenessRatio = 10,
            speckleWindowSize = 1000,
            speckleRange = 32
        )

        disp = stereo.compute(imgL, imgR).astype(np.float32) / 16.0
        if False != False:
            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 = cv.reprojectImageTo3D(disp, Q)
            colors = cv.cvtColor(imgL, cv.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)
            print('%s saved' % 'out.ply')

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