Merge pull request #1461 from luxonis/feature/point_cloud_data_bindings_update

Feature/point cloud data bindings update

Author: aljazdu

bindings/python/src/pipeline/datatype/PointCloudDataBindings.cpp

@@ -117,6 +117,39 @@ void bind_pointclouddata(pybind11::module& m, void* pCallstack) {
         .def("getTimestamp", &PointCloudData::Buffer::getTimestamp, DOC(dai, Buffer, getTimestamp))
         .def("getTimestampDevice", &PointCloudData::Buffer::getTimestampDevice, DOC(dai, Buffer, getTimestampDevice))
         .def("getSequenceNum", &PointCloudData::Buffer::getSequenceNum, DOC(dai, Buffer, getSequenceNum))
+        .def("setPoints", [](py::object& obj, py::array_t<float>& arr) {
+            if (arr.ndim() != 2 || arr.shape(1) != 3) {
+                throw std::runtime_error("Input must be a 2D numpy array of points with the shape of (N, 3)");
+            }
+            dai::PointCloudData& data = obj.cast<dai::PointCloudData&>();
+            std::vector<Point3f> points;
+            points.reserve(arr.shape(0));
+            auto ra = arr.unchecked();
+            for(int i = 0; i < arr.shape(0); i++) {
+                points.emplace_back(ra(i, 0), ra(i, 1), ra(i, 2));
+            }
+            data.setPoints(points);
+        })
+        .def("setPointsRGB", [](py::object& obj, py::array_t<float>& points, py::array_t<uint8_t>& colors) {
+            if (points.ndim() != 2 || points.shape(1) != 3) {
+                throw std::runtime_error("Points input must be a 2D numpy array of points with the shape of (N, 3)");
+            }
+            if (colors.ndim() != 2 || colors.shape(1) != 4) {
+                throw std::runtime_error("Colors input must be a 2D numpy array of colors with the shape of (N, 4)");
+            }
+            if (points.shape(0) != colors.shape(0)) {
+                throw std::runtime_error("Points and Colors must have the same number of rows");
+            }
+            dai::PointCloudData& data = obj.cast<dai::PointCloudData&>();
+            std::vector<Point3fRGBA> pointsRGBA;
+            pointsRGBA.reserve(points.shape(0));
+            auto ra1 = points.unchecked();
+            auto ra2 = colors.unchecked();
+            for(int i = 0; i < points.shape(0); i++) {
+                pointsRGBA.emplace_back(ra1(i, 0), ra1(i, 1), ra1(i, 2), ra2(i, 0), ra2(i, 1), ra2(i, 2), ra2(i, 3));
+            }
+            data.setPointsRGB(pointsRGBA);
+        })
         .def("setWidth", &PointCloudData::setWidth, DOC(dai, PointCloudData, setWidth))
         .def("setHeight", &PointCloudData::setHeight, DOC(dai, PointCloudData, setHeight))
         .def("setSize",

bindings/python/tests/CMakeLists.txt

@@ -13,6 +13,7 @@ set(PYBIND11_TEST_FILES
     "utf8_support_test.py"
     # "dai_path_conversion_test.py" requires device
     "nndata_tensor_test.py"
+    "pcldata_set_points_test.py"
     "messsage_queue_test.py"
     "imgframe_test.py"
     "inherited_messages_test.py"

bindings/python/tests/pcldata_set_points_test.py

@@ -0,0 +1,33 @@
+import depthai as dai
+import numpy as np
+
+def test_pcldata_set_points():
+    pointCloudData = dai.PointCloudData()
+
+    # Points
+    points = np.random.rand(1000, 3).astype(np.float32)
+    pointCloudData.setPoints(points)
+
+    assert(pointCloudData.isColor() == False)
+    assert(isinstance(pointCloudData.getPoints(), np.ndarray))
+    assert(pointCloudData.getPoints().dtype == points.dtype)
+    assert(np.array_equal(pointCloudData.getPoints(), points))
+
+    # RGB Points
+    colors = np.random.rand(1000, 4).astype(np.uint8)
+    pointCloudData.setPointsRGB(points, colors)
+
+    assert(pointCloudData.isColor() == True)
+    assert(isinstance(pointCloudData.getPoints(), np.ndarray))
+    assert(isinstance(pointCloudData.getPointsRGB()[0], np.ndarray))
+    assert(isinstance(pointCloudData.getPointsRGB()[1], np.ndarray))
+    assert(pointCloudData.getPoints().dtype == points.dtype)
+    assert(pointCloudData.getPointsRGB()[0].dtype == points.dtype)
+    assert(pointCloudData.getPointsRGB()[1].dtype == colors.dtype)
+    assert(np.array_equal(pointCloudData.getPoints(), points))
+    assert(np.array_equal(pointCloudData.getPointsRGB()[0], points))
+    assert(np.array_equal(pointCloudData.getPointsRGB()[1], colors))
+
+
+if __name__ == '__main__':
+  test_pcldata_set_points()

examples/cpp/RGBD/CMakeLists.txt

@@ -17,4 +17,8 @@ endif()
 
 if(DEPTHAI_ENABLE_REMOTE_CONNECTION)
     dai_add_example(visualizer_rgbd_autocreate visualizer_rgbd_autocreate.cpp ON OFF)
+endif()
+
+if(DEPTHAI_ENABLE_REMOTE_CONNECTION)
+    dai_add_example(rgbd_pcl_processing rgbd_pcl_processing.cpp ON OFF)
 endif()

examples/cpp/RGBD/rgbd_pcl_processing.cpp

@@ -0,0 +1,83 @@
+
+#include <csignal>
+#include <depthai/remote_connection/RemoteConnection.hpp>
+#include <iostream>
+
+#include "depthai/depthai.hpp"
+
+
+class CustomPCLProcessingNode : public dai::NodeCRTP<dai::node::ThreadedHostNode, CustomPCLProcessingNode> {
+   public:
+    constexpr static const char* NAME = "CustomPCLProcessingNode";
+    constexpr static const float thresholdDistance = 3000.0f;
+
+   public:
+    Input inputPCL{*this, {.name = "inPCL", .types = {{dai::DatatypeEnum::PointCloudData, true}}}};
+    Output outputPCL{*this, {.name = "outPCL", .types = {{dai::DatatypeEnum::PointCloudData, true}}}};
+
+    void run() override {
+        while(isRunning()) {
+            auto pclIn = inputPCL.get<dai::PointCloudData>();
+            auto pclOut = std::make_shared<dai::PointCloudData>();
+            if(pclIn != nullptr) {
+                const auto& pclData = pclIn->getPointsRGB();
+                std::vector<dai::Point3fRGBA> updatedPoints;
+                const float thresholdSquared = thresholdDistance * thresholdDistance;
+                updatedPoints.reserve(pclData.size());
+                for (const auto& point : pclData) {
+                    const float distance = point.x*point.x + point.y*point.y + point.z*point.z;
+                    if (distance <= thresholdSquared) {
+                        updatedPoints.emplace_back(point.x, point.y, point.z, 
+                                                   point.r, point.g, point.b, point.a);
+                    }
+                }
+                updatedPoints.shrink_to_fit();
+                pclOut->setPointsRGB(updatedPoints);
+                outputPCL.send(pclOut);
+            }
+        }
+    }
+};
+
+// Signal handling for clean shutdown
+static bool isRunning = true;
+void signalHandler(int signum) {
+    isRunning = false;
+}
+
+int main() {
+    using namespace std;
+    // Default port values
+    int webSocketPort = 8765;
+    int httpPort = 8080;
+
+    // Register signal handler
+    std::signal(SIGINT, signalHandler);
+
+    // Create RemoteConnection
+    dai::RemoteConnection remoteConnector(dai::RemoteConnection::DEFAULT_ADDRESS, webSocketPort, true, httpPort);
+    // Create pipeline
+    dai::Pipeline pipeline;
+    auto rgbd = pipeline.create<dai::node::RGBD>()->build(true, dai::node::StereoDepth::PresetMode::DEFAULT);
+    auto customNode = pipeline.create<CustomPCLProcessingNode>();
+
+    rgbd->pcl.link(customNode->inputPCL);
+
+    remoteConnector.addTopic("pcl", rgbd->pcl);
+    remoteConnector.addTopic("processed_pcl", customNode->outputPCL);
+    pipeline.start();
+    remoteConnector.registerPipeline(pipeline);
+    auto device = pipeline.getDefaultDevice();
+    device->setIrLaserDotProjectorIntensity(0.7);
+    // Main loop
+    while(isRunning && pipeline.isRunning()) {
+        int key = remoteConnector.waitKey(1);
+        if(key == 'q') {
+            std::cout << "Got 'q' key from the remote connection!" << std::endl;
+            break;
+        }
+    }
+
+    std::cout << "Pipeline stopped." << std::endl;
+    return 0;
+}

examples/python/CMakeLists.txt

@@ -239,6 +239,9 @@ dai_set_example_test_labels(visualizer_rgbd_autocreate ondevice rvc2 usb rvc4 ci
 add_python_example(rgbd_o3d RGBD/rgbd_o3d.py)
 dai_set_example_test_labels(rgbd_o3d ondevice rvc2_all rvc4 ci nowindows)
 
+add_python_example(rgbd_pcl_processing RGBD/rgbd_pcl_processing.py)
+dai_set_example_test_labels(rgbd_pcl_processing ondevice rvc2_all rvc4 ci nowindows)
+
 add_python_example(person_crop_out Misc/DigitalZoom/person_crop_out.py)
 dai_set_example_test_labels(person_crop_out ondevice rvc2_all rvc4 rvc4rgb ci)
 

examples/python/RGBD/rgbd_pcl_processing.py

@@ -0,0 +1,55 @@
+import depthai as dai
+import numpy as np
+from argparse import ArgumentParser
+
+
+class CustomPCLProcessingNode(dai.node.ThreadedHostNode):
+    def __init__(self):
+        dai.node.ThreadedHostNode.__init__(self)
+        self.inputPCL = self.createInput()
+        self.outputPCL = self.createOutput()
+        self.thresholdDistance = 3000.0
+
+    def run(self):
+        while self.isRunning():
+            try:
+                inPointCloud = self.inputPCL.get()
+            except dai.MessageQueue.QueueException:
+                return # Pipeline closed
+            if inPointCloud is not None:
+                outPointCloud = dai.PointCloudData()
+                points, colors = inPointCloud.getPointsRGB()
+
+                mask = np.linalg.norm(points, axis=1) < self.thresholdDistance
+                updatedPoints = points[mask]
+                updatedColors = colors[mask]
+
+                outPointCloud.setPointsRGB(updatedPoints, updatedColors)
+                self.outputPCL.send(outPointCloud)
+
+parser = ArgumentParser()
+parser.add_argument("--webSocketPort", type=int, default=8765)
+parser.add_argument("--httpPort", type=int, default=8080)
+args = parser.parse_args()
+
+with dai.Pipeline() as p:
+    remoteConnector = dai.RemoteConnection(
+        webSocketPort=args.webSocketPort, httpPort=args.httpPort
+    )
+    rgbd = p.create(dai.node.RGBD).build(True, dai.node.StereoDepth.PresetMode.DEFAULT)
+    customNode = p.create(CustomPCLProcessingNode)
+
+    # Link rgbd.pcl to the input of CustomPCLProcessingNode
+    rgbd.pcl.link(customNode.inputPCL)
+
+    remoteConnector.addTopic("pcl", rgbd.pcl, "common")
+    remoteConnector.addTopic("processed_pcl", customNode.outputPCL, "common")
+
+    p.start()
+    remoteConnector.registerPipeline(p)
+
+    while p.isRunning():
+        key = remoteConnector.waitKey(1)
+        if key == ord("q"):
+            print("Got q key from the remote connection!")
+            break