diff --git a/Cargo.lock b/Cargo.lock index e8d7908..b7dff21 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -31,7 +31,7 @@ dependencies = [ [[package]] name = "e57-python" -version = "0.1.0-a4" +version = "0.1.0-a5" dependencies = [ "e57", "ndarray", diff --git a/Cargo.toml b/Cargo.toml index a9092b2..fa82175 100644 --- a/Cargo.toml +++ b/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" diff --git a/README.md b/README.md index 255996c..707e4af 100644 --- a/README.md +++ b/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 diff --git a/src/lib.rs b/src/lib.rs index ac4f4fe..e819e52 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,6 +12,8 @@ pub struct E57 { pub points: Py>, #[pyo3(get)] pub color: Py>, + #[pyo3(get)] + pub intensity: Py>, } /// Extracts the xml contents from an e57 file. @@ -34,7 +36,7 @@ fn raw_xml(filepath: &str) -> PyResult { /// Extracts the point data from an e57 file. #[pyfunction] -unsafe fn read_points<'py>(py: Python<'py>, filepath: &str) -> PyResult { +unsafe fn read_points(py: Python<'_>, filepath: &str) -> PyResult { 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 { }; 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 { 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 { 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. diff --git a/tests/test_e57.py b/tests/test_e57.py index dd5a064..f34568c 100644 --- a/tests/test_e57.py +++ b/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