Merge pull request strawlab#61 from larsmans/buffer-api

nzjrs · nzjrs · commit c60e59f890b3 · 2015-01-20T15:15:41.000+01:00
Buffer api: easier NumPy interop
diff --git a/pcl/_pcl.pxd b/pcl/_pcl.pxd
@@ -7,6 +7,10 @@ cimport pcl_defs as cpp
 cdef class PointCloud:
     cdef cpp.PointCloudPtr_t thisptr_shared
 
+    # Buffer protocol support.
+    cdef Py_ssize_t _shape[2]
+    cdef Py_ssize_t _view_count
+
     cdef inline cpp.PointCloud[cpp.PointXYZ] *thisptr(self) nogil:
         # Shortcut to get raw pointer to underlying PointCloud<PointXYZ>.
         return self.thisptr_shared.get()
diff --git a/pcl/_pcl.pyx b/pcl/_pcl.pyx
@@ -10,6 +10,9 @@ cimport pcl_defs as cpp
 
 cimport cython
 from cython.operator import dereference as deref
+
+from cpython cimport Py_buffer
+
 from libcpp.string cimport string
 from libcpp cimport bool
 from libcpp.vector cimport vector
@@ -125,6 +128,20 @@ cdef class SegmentationNormal:
     def set_axis(self, double ax, double ay, double az):
         mpcl_sacnormal_set_axis(deref(self.me),ax,ay,az)
 
+
+# Empirically determine strides, for buffer support.
+# XXX Is there a more elegant way to get these?
+cdef Py_ssize_t _strides[2]
+cdef PointCloud _pc_tmp = PointCloud(np.array([[1, 2, 3],
+                                               [4, 5, 6]], dtype=np.float32))
+cdef cpp.PointCloud[cpp.PointXYZ] *p = _pc_tmp.thisptr()
+_strides[0] = (  <Py_ssize_t><void *>cpp.getptr(p, 1)
+               - <Py_ssize_t><void *>cpp.getptr(p, 0))
+_strides[1] = (  <Py_ssize_t><void *>&(cpp.getptr(p, 0).y)
+               - <Py_ssize_t><void *>&(cpp.getptr(p, 0).x))
+_pc_tmp = None
+
+
 cdef class PointCloud:
     """Represents a cloud of points in 3-d space.
 
@@ -137,6 +154,8 @@ cdef class PointCloud:
     def __cinit__(self, init=None):
         cdef PointCloud other
 
+        self._view_count = 0
+
         sp_assign(self.thisptr_shared, new cpp.PointCloud[cpp.PointXYZ]())
 
         if init is None:
@@ -170,6 +189,32 @@ cdef class PointCloud:
     def __repr__(self):
         return "<PointCloud of %d points>" % self.size
 
+    # Buffer protocol support. Taking a view locks the pointcloud for
+    # resizing, because that can move it around in memory.
+    def __getbuffer__(self, Py_buffer *buffer, int flags):
+        # TODO parse flags
+        cdef Py_ssize_t npoints = self.thisptr().size()
+
+        if self._view_count == 0:
+            self._view_count += 1
+            self._shape[0] = npoints
+            self._shape[1] = 3
+
+        buffer.buf = <char *>&(cpp.getptr_at(self.thisptr(), 0).x)
+        buffer.format = 'f'
+        buffer.internal = NULL
+        buffer.itemsize = sizeof(float)
+        buffer.len = npoints * 3 * sizeof(float)
+        buffer.ndim = 2
+        buffer.obj = self
+        buffer.readonly = 0
+        buffer.shape = self._shape
+        buffer.strides = _strides
+        buffer.suboffsets = NULL
+
+    def __releasebuffer__(self, Py_buffer *buffer):
+        self._view_count -= 1
+
     # Pickle support. XXX this copies the entire pointcloud; it would be nice
     # to have an asarray member that returns a view, or even better, implement
     # the buffer protocol (https://docs.python.org/c-api/buffer.html).
@@ -246,6 +291,9 @@ cdef class PointCloud:
         return self.to_array().tolist()
 
     def resize(self, cnp.npy_intp x):
+        if self._view_count > 0:
+            raise ValueError("can't resize PointCloud while there are"
+                             " arrays/memoryviews referencing it")
         self.thisptr().resize(x)
 
     def get_point(self, cnp.npy_intp row, cnp.npy_intp col):
diff --git a/readme.rst b/readme.rst
@@ -29,7 +29,7 @@ point types
  * registration (ICP, GICP, ICP_NL)
 
 The code tries to follow the Point Cloud API, and also provides helper function
-for interacting with numpy. For example (from tests/test.py)
+for interacting with NumPy. For example (from tests/test.py)
 
 .. code-block:: python
 
@@ -52,6 +52,20 @@ or, for smoothing
     fil.set_std_dev_mul_thresh (1.0)
     fil.filter().to_file("inliers.pcd")
 
+Point clouds can be viewed as NumPy arrays, so modifying them is possible
+using all the familiar NumPy functionality:
+
+.. code-block:: python
+
+    import numpy as np
+    import pcl
+    p = pcl.PointCloud(10)  # "empty" point cloud
+    a = np.asarray(p)       # NumPy view on the cloud
+    a[:] = 0                # fill with zeros
+    print(p[3])             # prints (0.0, 0.0, 0.0)
+    a[:, 0] = 1             # set x coordinates to 1
+    print(p[3])             # prints (1.0, 0.0, 0.0)
+
 More samples can be found in the `examples directory <https://github.com/strawlab/python-pcl/tree/master/examples>`_,
 and in the `unit tests <https://github.com/strawlab/python-pcl/blob/master/tests/test.py>`_.
 
@@ -60,11 +74,11 @@ This work was supported by `Strawlab <http://strawlab.org/>`_.
 Requirements
 ------------
 
-This release has been tested on Ubuntu 13.10 with
+This release has been tested on Linux Mint 17 with
 
- * Python 2.7.5
- * pcl 1.7.1
- * Cython 0.21
+ * Python 2.7.6
+ * pcl 1.7.2
+ * Cython 0.21.2
 
 and CentOS 6.5 with
 
diff --git a/tests/test.py b/tests/test.py
@@ -7,7 +7,7 @@
 import pcl
 import numpy as np
 
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_almost_equal, assert_array_equal
 
 
 _data = [(i,2*i,3*i+0.2) for i in range(5)]
@@ -38,7 +38,7 @@ def setUp(self):
         self.a = np.array(np.mat(_DATA, dtype=np.float32))
         self.p = pcl.PointCloud(self.a)
 
-    def testFromNumy(self):
+    def testFromNumpy(self):
         for i,d in enumerate(_data):
             pt = self.p[i]
             assert np.allclose(pt, _data[i])
@@ -47,6 +47,13 @@ def testToNumpy(self):
         a = self.p.to_array()
         self.assertTrue(np.alltrue(a == self.a))
 
+    def test_asarray(self):
+        p = pcl.PointCloud(self.p)      # copy
+        old0 = p[0]
+        a = np.asarray(p)               # view
+        a[:] += 6
+        assert_array_almost_equal(p[0], a[0])
+
     def test_pickle(self):
         """Test pickle support."""
         # In this testcase because picking reduces to pickling NumPy arrays.
