From 1adfcba9426d0a97ac9b5d8d0d7f5dc6513ddd1a Mon Sep 17 00:00:00 2001
From: Evan Pratten <ewpratten@gmail.com>
Date: Sat, 29 May 2021 16:26:27 -0400
Subject: [PATCH] basic depth calculations

---
 README.md           |  2 +-
 odm/__main__.py     | 35 +++++++++++++++++++++++++++++++++--
 pylibodm/src/lib.rs |  7 +++++++
 requirements.txt    |  3 ++-
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 02bec57..f86de3b 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 # OpenDepthMap
 Point cloud streams from Leap Motion cameras
 ```
-clang libclang-dev libopencv-dev python3-dev python-dev
+clang libclang-dev libopencv-dev python3-dev python-dev python3-opencv
 
 ```
 
diff --git a/odm/__main__.py b/odm/__main__.py
index 9563260..7cde98d 100644
--- a/odm/__main__.py
+++ b/odm/__main__.py
@@ -2,6 +2,8 @@ import argparse
 import sys
 import logging
 import os
+import numpy as np
+import cv2
 
 FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s'
 logging.basicConfig(format=FORMAT)
@@ -11,17 +13,46 @@ logging.getLogger().setLevel(logging.DEBUG)
 sys.path.append(os.getcwd() + "/target/debug")
 import libpylibodm as pylibodm
 
+def handle_image_data(stereo, left_image, right_image):
+
+    # Convert images to something opencv can use
+    left_image_cv = np.frombuffer(bytes(left_image.buffer), np.uint8).reshape(
+        left_image.height, left_image.width)
+    right_image_cv = np.frombuffer(bytes(right_image.buffer), np.uint8).reshape(
+        left_image.height, left_image.width)
+
+    # Compute a disparity map
+    disparity = stereo.compute(left_image_cv, right_image_cv)
+    disparity = cv2.convertScaleAbs(disparity, alpha=1.5)
+
+    cv2.imshow("Left", left_image_cv)
+    cv2.imshow("Right", right_image_cv)
+    cv2.imshow("Disparity", disparity)
+
+
 def main() -> int:
     # Handle program arguments
-    ap = argparse.ArgumentParser(prog='odm.py', description='Stream 3D data from a LeapMotion camera')
+    ap = argparse.ArgumentParser(
+        prog='odm.py', description='Stream 3D data from a LeapMotion camera')
     args = ap.parse_args()
 
     # Connect to the leapmotion device
     print("Connecting to LeapMotion")
-    print(dir(pylibodm))
     pylibodm.connect(4)
 
+    # Set up depth mapping
+    stereo = cv2.StereoBM_create(numDisparities=32, blockSize=11)
+
+    # Handle data
+    while True:
+        frame = pylibodm.get_frame()
+        handle_image_data(stereo, frame.left_camera, frame.right_camera)
+
+        if cv2.waitKey(25) & 0xFF == ord('q'):
+            break
+
     return 0
 
+
 if __name__ == "__main__":
     sys.exit(main())
\ No newline at end of file
diff --git a/pylibodm/src/lib.rs b/pylibodm/src/lib.rs
index ba454a0..d3c37a5 100644
--- a/pylibodm/src/lib.rs
+++ b/pylibodm/src/lib.rs
@@ -26,9 +26,13 @@ use libodm::{image::Image, leapmotion::device::DeviceFrame};
 use pyo3::wrap_pyfunction;
 
 #[pyclass]
+#[derive(Debug, Clone)]
 struct PyImage {
+    #[pyo3(get)]
     pub width: u32,
+    #[pyo3(get)]
     pub height: u32,
+    #[pyo3(get)]
     pub buffer: Vec<u8>,
 }
 
@@ -44,8 +48,11 @@ impl IntoPy<PyImage> for Image<'_> {
 
 #[pyclass]
 struct PyDeviceFrame {
+    #[pyo3(get)]
     pub bytes_per_pixel: u8,
+    #[pyo3(get)]
     pub left_camera: PyImage,
+    #[pyo3(get)]
     pub right_camera: PyImage,
 }
 
diff --git a/requirements.txt b/requirements.txt
index 419c596..7e67d7b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,2 @@
-setuptools-rust>=0.10.2
\ No newline at end of file
+setuptools-rust>=0.10.2
+numpy
\ No newline at end of file