♻️ Refactor to return 3D arrays of shape (band, height, width) (#11)

Adding a third dimension (band or channel) to the ndarray output in anticipation of multi-band support. Updated the unit tests to check for the new 3D shape.

README.md

@@ -34,7 +34,7 @@ use std::io::Cursor;
 
 use bytes::Bytes;
 use cog3pio::io::geotiff::read_geotiff;
-use ndarray::Array2;
+use ndarray::Array3;
 use object_store::path::Path;
 use object_store::{parse_url, GetResult, ObjectStore};
 use tokio;
@@ -53,9 +53,9 @@ async fn main() {
         Cursor::new(bytes)
     };
 
-    let arr: Array2<f32> = read_geotiff(stream).unwrap();
-    assert_eq!(arr.dim(), (549, 549));
-    assert_eq!(arr[[500, 500]], 0.13482364);
+    let arr: Array3<f32> = read_geotiff(stream).unwrap();
+    assert_eq!(arr.dim(), (1, 549, 549));
+    assert_eq!(arr[[0, 500, 500]], 0.13482364);
 }
 ```
 
@@ -68,7 +68,7 @@ from cog3pio import read_geotiff
 array: np.ndarray = read_geotiff(
     path="https://github.com/cogeotiff/rio-tiler/raw/6.4.0/tests/fixtures/cog_nodata_nan.tif"
 )
-assert array.shape == (549, 549)
+assert array.shape == (1, 549, 549)  # bands, height, width
 assert array.dtype == "float32"
 ```
 

python/tests/test_io_geotiff.py

@@ -31,7 +31,7 @@ def test_read_geotiff_local(geotiff_path):
     Read a GeoTIFF file from a local file path.
     """
     array = read_geotiff(path=geotiff_path)
-    assert array.shape == (20, 20)
+    assert array.shape == (1, 20, 20)
     assert array.dtype == "float32"
 
 
@@ -43,7 +43,7 @@ def test_read_geotiff_remote():
     array = read_geotiff(
         path="https://github.com/pka/georaster/raw/v0.1.0/data/tiff/float32.tif"
     )
-    assert array.shape == (20, 20)
+    assert array.shape == (1, 20, 20)
     assert array.dtype == "float32"
 
 

src/io/geotiff.rs

@@ -1,7 +1,7 @@
 use std::io::{Read, Seek};
 
 use geo::AffineTransform;
-use ndarray::Array2;
+use ndarray::Array3;
 use tiff::decoder::{Decoder, DecodingResult, Limits};
 use tiff::tags::Tag;
 use tiff::{TiffError, TiffFormatError, TiffResult};
@@ -22,7 +22,7 @@ impl<R: Read + Seek> CogReader<R> {
     }
 
     /// Decode GeoTIFF image to an [`ndarray::Array`]
-    fn ndarray(&mut self) -> TiffResult<Array2<f32>> {
+    fn ndarray(&mut self) -> TiffResult<Array3<f32>> {
         // Get image dimensions
         let (width, height): (u32, u32) = self.decoder.dimensions()?;
 
@@ -34,7 +34,7 @@ impl<R: Read + Seek> CogReader<R> {
         };
 
         // Put image pixel data into an ndarray
-        let vec_data = Array2::from_shape_vec((height as usize, width as usize), image_data)
+        let vec_data = Array3::from_shape_vec((1, height as usize, width as usize), image_data)
             .map_err(|_| TiffFormatError::InvalidDimensions(height, width))?;
 
         Ok(vec_data)
@@ -91,12 +91,12 @@ impl<R: Read + Seek> CogReader<R> {
 }
 
 /// Synchronously read a GeoTIFF file into an [`ndarray::Array`]
-pub fn read_geotiff<R: Read + Seek>(stream: R) -> TiffResult<Array2<f32>> {
+pub fn read_geotiff<R: Read + Seek>(stream: R) -> TiffResult<Array3<f32>> {
     // Open TIFF stream with decoder
     let mut reader = CogReader::new(stream)?;
 
     // Decode TIFF into ndarray
-    let vec_data: Array2<f32> = reader.ndarray()?;
+    let vec_data: Array3<f32> = reader.ndarray()?;
 
     Ok(vec_data)
 }
@@ -106,7 +106,7 @@ mod tests {
     use std::io::{Cursor, Seek, SeekFrom};
 
     use geo::AffineTransform;
-    use ndarray::array;
+    use ndarray::{array, s};
     use object_store::parse_url;
     use tempfile::tempfile;
     use tiff::encoder::{colortype, TiffEncoder};
@@ -135,10 +135,11 @@ mod tests {
 
         // Read a BigTIFF file
         let arr = read_geotiff(file).unwrap();
-        assert_eq!(arr.ndim(), 2);
-        assert_eq!(arr.dim(), (10, 20)); // (height, width)
-        assert_eq!(arr.nrows(), 10); // y-axis
-        assert_eq!(arr.ncols(), 20); // x-axis
+        assert_eq!(arr.ndim(), 3);
+        assert_eq!(arr.dim(), (1, 10, 20)); // (channels, height, width)
+        let first_band = arr.slice(s![0, .., ..]);
+        assert_eq!(first_band.nrows(), 10); // y-axis
+        assert_eq!(first_band.ncols(), 20); // x-axis
         assert_eq!(arr.mean(), Some(14.0));
     }
 
@@ -155,8 +156,8 @@ mod tests {
         let mut reader = CogReader::new(stream).unwrap();
         let array = reader.ndarray().unwrap();
 
-        assert_eq!(array.shape(), [2, 3]);
-        assert_eq!(array, array![[1.41, 1.23, 0.78], [0.32, -0.23, -1.88]])
+        assert_eq!(array.shape(), [1, 2, 3]);
+        assert_eq!(array, array![[[1.41, 1.23, 0.78], [0.32, -0.23, -1.88]]])
     }
 
     #[tokio::test]

src/python/adapters.rs

@@ -1,7 +1,7 @@
 use std::io::Cursor;
 
 use bytes::Bytes;
-use numpy::{PyArray2, ToPyArray};
+use numpy::{PyArray3, ToPyArray};
 use object_store::{parse_url, ObjectStore};
 use pyo3::exceptions::{PyBufferError, PyFileNotFoundError, PyValueError};
 use pyo3::prelude::{pyfunction, pymodule, PyModule, PyResult, Python};
@@ -44,7 +44,7 @@ fn path_to_stream(path: &str) -> PyResult<Cursor<Bytes>> {
     Ok(stream)
 }
 
-/// Read a GeoTIFF file from a path on disk into an ndarray
+/// Read a GeoTIFF file from a path on disk or a url into an ndarray
 ///
 /// Parameters
 /// ----------
@@ -64,7 +64,7 @@ fn path_to_stream(path: &str) -> PyResult<Cursor<Bytes>> {
 /// assert array.shape == (20, 20)
 #[pyfunction]
 #[pyo3(name = "read_geotiff")]
-fn read_geotiff_py<'py>(path: &str, py: Python<'py>) -> PyResult<&'py PyArray2<f32>> {
+fn read_geotiff_py<'py>(path: &str, py: Python<'py>) -> PyResult<&'py PyArray3<f32>> {
     // Parse URL into byte stream
     let stream = path_to_stream(path)?;