#!/usr/bin/env python
import freenect
import cv
import numpy as np

cv.NamedWindow('Depth')
cv.NamedWindow('Video')
print('Press ESC in window to stop')


def pretty_depth(depth):
    np.clip(depth, 0, 2**10 - 1, depth)
    depth >>= 2
    return depth.astype(np.uint8)


def get_depth():
    data, _ = freenect.sync_get_depth()
    data = pretty_depth(data)
    image = cv.CreateImageHeader((data.shape[1], data.shape[0]),
                                 cv.IPL_DEPTH_8U,
                                 1)
    cv.SetData(image, data.tostring(),
               data.dtype.itemsize * data.shape[1])
    return image


def get_video():
    data, _ = freenect.sync_get_video()
    data = data[:, :, ::-1].astype(np.uint8)  # RGB -> BGR
    image = cv.CreateImageHeader((data.shape[1], data.shape[0]),
                                 cv.IPL_DEPTH_8U,
                                 3)
    cv.SetData(image, data.tostring(),
               data.dtype.itemsize * 3 * data.shape[1])
    return image


while 1:
    cv.ShowImage('Depth', get_depth())
    cv.ShowImage('Video', get_video())
    if cv.WaitKey(10) == 27:
        break