diff --git a/src/lib.rs b/src/lib.rs index a2ef5f8..afe49df 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,9 @@ use ::e57::{CartesianCoordinate, E57Reader}; use ndarray::Ix2; use numpy::{PyArray, PyArrayMethods}; use pyo3::prelude::*; +use image::DynamicImage; +use opencv::core::Mat; +use opencv::prelude::*; #[pyclass] pub struct E57 { @@ -14,6 +17,10 @@ pub struct E57 { pub color: Py>, #[pyo3(get)] pub intensity: Py>, + #[pyo3(get)] + pub images: Vec>>, + #[pyo3(get)] + pub metadata: Vec, } /// Extracts the xml contents from an e57 file. @@ -87,6 +94,8 @@ unsafe fn read_points(py: Python<'_>, filepath: &str) -> PyResult { ), color: Py::from(PyArray::new_bound(py, (0, 3), false)), intensity: Py::from(PyArray::new_bound(py, (0, 1), false)), + images: Vec::new(), + metadata: Vec::new(), }; if n_colors == n_points { e57.color = Py::from( @@ -105,11 +114,44 @@ unsafe fn read_points(py: Python<'_>, filepath: &str) -> PyResult { Ok(e57) } +/// Extracts the image data and metadata from an e57 file. +#[pyfunction] +unsafe fn read_images(py: Python<'_>, filepath: &str) -> PyResult { + let file = E57Reader::from_file(filepath); + let mut file = match file { + Ok(file) => file, + Err(e) => { + return Err(PyErr::new::( + e.to_string(), + )); + } + }; + let mut images = Vec::new(); + let mut metadata = Vec::new(); + for image in file.images() { + let img_data = file.image_data(&image).expect("Unable to read image data"); + let img = DynamicImage::from_bytes(img_data).expect("Unable to create image"); + let mat = Mat::from_slice(&img.to_bytes()).expect("Unable to create Mat"); + let np_array = PyArray::from_vec(py, mat.data_bytes().to_vec()); + images.push(Py::from(np_array)); + metadata.push(image.metadata.clone()); + } + let e57 = E57 { + points: Py::from(PyArray::new_bound(py, (0, 3), false)), + color: Py::from(PyArray::new_bound(py, (0, 3), false)), + intensity: Py::from(PyArray::new_bound(py, (0, 1), false)), + images, + metadata, + }; + Ok(e57) +} + /// e57 pointcloud file reading. #[pymodule] fn e57(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_function(wrap_pyfunction!(raw_xml, m)?)?; m.add_function(wrap_pyfunction!(read_points, m)?)?; + m.add_function(wrap_pyfunction!(read_images, m)?)?; Ok(()) } diff --git a/tests/test_e57.py b/tests/test_e57.py index f5598de..72536af 100644 --- a/tests/test_e57.py +++ b/tests/test_e57.py @@ -150,3 +150,15 @@ def test_raw_xml_just_xml(): assert "Failed to read E57" in str(e) assert "Failed creating paged CRC reader" in str(e) assert raised + + +def test_read_images(): + e57_data = e57.read_images(r"testdata/pipeSpherical.e57") + images = e57_data.images + metadata = e57_data.metadata + assert isinstance(images, list) + assert all(isinstance(img, np.ndarray) for img in images) + assert isinstance(metadata, list) + assert all(isinstance(meta, str) for meta in metadata) + assert len(images) == 6 + assert len(metadata) == 6