Merge pull request #18 from dancergraham/feature/fields

feature: Intensity

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "e57-python"
-version = "0.1.0-a5"
+version = "0.1.0-a6"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -10,6 +10,6 @@ crate-type = ["cdylib"]
 
 [dependencies]
 pyo3 = "0.20.0"
-e57 = "0.10.2"
+e57 = "0.10.3"
 numpy = "0.20.0"
 ndarray = "0.15.6"

README.md

@@ -11,7 +11,8 @@ This python library wraps the [rust e57 library](https://github.com/cry-inc/e57)
 - [x] Proof of concept xml reading
 - [x] Read e57 point coordinates to numpy array - see `read_points` method.
 - [x] Read color field to numpy array.
-- [ ] Read intensity and other fields to numpy array.
+- [x] Read intensity to numpy array.
+- [ ] Read other fields to numpy array.
 - [ ] Write to e57 (format ?)
 
 ## Getting Started

src/lib.rs

@@ -12,6 +12,8 @@ pub struct E57 {
     pub points: Py<PyArray<f64, Ix2>>,
     #[pyo3(get)]
     pub color: Py<PyArray<f32, Ix2>>,
+    #[pyo3(get)]
+    pub intensity: Py<PyArray<f32, Ix2>>,
 }
 
 /// Extracts the xml contents from an e57 file.
@@ -34,7 +36,7 @@ fn raw_xml(filepath: &str) -> PyResult<String> {
 
 /// Extracts the point data from an e57 file.
 #[pyfunction]
-unsafe fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<E57> {
+unsafe fn read_points(py: Python<'_>, filepath: &str) -> PyResult<E57> {
     let file = E57Reader::from_file(filepath);
     let mut file = match file {
         Ok(file) => file,
@@ -46,8 +48,9 @@ unsafe fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<E57> {
     };
     let pc = file.pointclouds();
     let pc = pc.first().expect("files contain pointclouds");
+    let mut point_vec = Vec::with_capacity(pc.records as usize * 3);
     let mut color_vec = Vec::with_capacity(pc.records as usize * 3);
-    let mut vec = Vec::with_capacity(pc.records as usize * 3);
+    let mut intensity_vec = Vec::with_capacity(pc.records as usize);
     let mut nrows = 0;
     for pointcloud in file.pointclouds() {
         let mut iter = file
@@ -61,7 +64,7 @@ unsafe fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<E57> {
         for p in iter {
             let p = p.expect("Unable to read next point");
             if let CartesianCoordinate::Valid { x, y, z } = p.cartesian {
-                vec.extend([x, y, z]);
+                point_vec.extend([x, y, z]);
                 nrows += 1
             }
             // if let Some(intensity) = p.intensity{
@@ -70,19 +73,26 @@ unsafe fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<E57> {
             if let Some(color) = p.color {
                 color_vec.extend([color.red, color.green, color.blue])
             }
+            if let Some(intensity) = p.intensity {
+                intensity_vec.push(intensity)
+            }
         }
     }
-    if color_vec.len() == vec.len() {
-        Ok(E57 {
-            points: Py::from(PyArray::from_vec(py, vec).reshape((nrows, 3)).unwrap()),
-            color: Py::from(PyArray::from_vec(py, color_vec).reshape((nrows, 3)).unwrap()),
-        })
-    } else {
-        Ok(E57 {
-            points: Py::from(PyArray::from_vec(py, vec).reshape((nrows, 3)).unwrap()),
-            color: Py::from(PyArray::new(py, (0,3), false)),
-        })
+    let n_points = point_vec.len() / 3;
+    let n_colors = color_vec.len() / 3;
+    let n_intensities = intensity_vec.len();
+    let mut e57 = E57 {
+        points: Py::from(PyArray::from_vec(py, point_vec).reshape((nrows, 3)).unwrap()),
+        color: Py::from(PyArray::new(py, (0, 3), false)),
+        intensity: Py::from(PyArray::new(py, (0, 1), false)),
+    };
+    if n_colors == n_points {
+        e57.color = Py::from(PyArray::from_vec(py, color_vec).reshape((nrows, 3)).unwrap())
+    }
+    if n_intensities == n_points {
+        e57.intensity = Py::from(PyArray::from_vec(py, intensity_vec).reshape((nrows, 1)).unwrap())
     }
+    Ok(e57)
 }
 
 /// e57 pointcloud file reading.

tests/test_e57.py

@@ -29,6 +29,13 @@ def test_read_color():
     assert len(color) == 1_220
 
 
+def test_read_intensity():
+    pointcloud = e57.read_points(r"testdata/pipeSpherical.e57")
+    intensity = pointcloud.intensity
+    assert isinstance(intensity, np.ndarray)
+    assert len(intensity) == 1_220
+
+
 def test_box_dimensions():
     pointcloud: np.ndarray = e57.read_points(r"testdata/bunnyFloat.e57")
     points = pointcloud.points