Merge pull request #16 from dancergraham/feature/color

feature: Now in technicolor! 🏳️‍🌈

Author: dancergraham

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "e57-python"
-version = "0.1.0-a4"
+version = "0.1.0-a5"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

README.md

@@ -10,6 +10,7 @@ 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.
 - [ ] Write to e57 (format ?)
 

src/lib.rs

@@ -2,10 +2,18 @@ use std::fs::File;
 use std::io::BufReader;
 
 use ::e57::{CartesianCoordinate, E57Reader};
-use ndarray::Ix2;
-use numpy::{PyArray};
+use ndarray::{Ix2};
+use numpy::PyArray;
 use pyo3::prelude::*;
 
+#[pyclass]
+pub struct E57 {
+    #[pyo3(get)]
+    pub points: Py<PyArray<f64, Ix2>>,
+    #[pyo3(get)]
+    pub color: Py<PyArray<f32, Ix2>>,
+}
+
 /// Extracts the xml contents from an e57 file.
 #[pyfunction]
 fn raw_xml(filepath: &str) -> PyResult<String> {
@@ -26,7 +34,7 @@ fn raw_xml(filepath: &str) -> PyResult<String> {
 
 /// Extracts the point data from an e57 file.
 #[pyfunction]
-fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<&'py PyArray<f64, Ix2>> {
+unsafe fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<E57> {
     let file = E57Reader::from_file(filepath);
     let mut file = match file {
         Ok(file) => file,
@@ -38,8 +46,8 @@ fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<&'py PyArray<f6
     };
     let pc = file.pointclouds();
     let pc = pc.first().expect("files contain pointclouds");
-    let ncols = 3;
-    let mut vec = Vec::with_capacity(pc.records as usize * ncols);
+    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 nrows = 0;
     for pointcloud in file.pointclouds() {
         let mut iter = file
@@ -56,15 +64,31 @@ fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult<&'py PyArray<f6
                 vec.extend([x, y, z]);
                 nrows += 1
             }
+            // if let Some(intensity) = p.intensity{
+            //     vec.append(intensity as f64)
+            // }
+            if let Some(color) = p.color {
+                color_vec.extend([color.red, color.green, color.blue])
+            }
         }
     }
-
-    Ok(PyArray::from_vec(py, vec).reshape((nrows, ncols)).unwrap())
+    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)),
+        })
+    }
 }
 
 /// e57 pointcloud file reading.
 #[pymodule]
-fn e57(_py: Python, m: &PyModule) -> PyResult<()> {
+fn e57(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    m.add_class::<E57>()?;
     m.add_function(wrap_pyfunction!(raw_xml, m)?)?;
     m.add_function(wrap_pyfunction!(read_points, m)?)?;
     Ok(())

testdata/pipeSpherical.e57

@@ -10,14 +10,30 @@ def test_raw_xml():
 
 def test_read_points():
     pointcloud = e57.read_points(r"testdata/bunnyFloat.e57")
-    assert isinstance(pointcloud, np.ndarray)
-    assert len(pointcloud) == 30_571
+    points = pointcloud.points
+    assert isinstance(points, np.ndarray)
+    assert len(points) == 30_571
+
+
+def test_read_spherical():
+    pointcloud = e57.read_points(r"testdata/pipeSpherical.e57")
+    points = pointcloud.points
+    assert isinstance(points, np.ndarray)
+    assert len(points) == 1_220
+
+
+def test_read_color():
+    pointcloud = e57.read_points(r"testdata/pipeSpherical.e57")
+    color = pointcloud.color
+    assert isinstance(color, np.ndarray)
+    assert len(color) == 1_220
 
 
 def test_box_dimensions():
     pointcloud: np.ndarray = e57.read_points(r"testdata/bunnyFloat.e57")
-    max_coords = pointcloud.max(0, None, False, -np.inf)
-    min_coords = pointcloud.min(0, None, False, np.inf)
+    points = pointcloud.points
+    max_coords = points.max(0, None, False, -np.inf)
+    min_coords = points.min(0, None, False, np.inf)
     X, Y, Z = max_coords - min_coords
     assert X == pytest.approx(0.155698)
     assert Y == pytest.approx(0.14731)
@@ -26,8 +42,8 @@ def test_box_dimensions():
 
 def test_global_box_center():
     pointcloud: np.ndarray = e57.read_points(r"testdata/bunnyFloat.e57")
-    max_coords = pointcloud.max(0, None, False, -np.inf)
-    min_coords = pointcloud.min(0, None, False, np.inf)
+    max_coords = pointcloud.points.max(0, None, False, -np.inf)
+    min_coords = pointcloud.points.min(0, None, False, np.inf)
     X, Y, Z = (max_coords + min_coords) / 2
     assert X == pytest.approx(-0.016840)
     assert Y == pytest.approx(0.113666)