Skip to content

Commit

Permalink
✨ Implement ndarray method in CogReader struct
Browse files Browse the repository at this point in the history
Move the previous tiff-decoder -> Vec -> ndarray code to an impl method under CogReader. Was going to separate the as_vec and ndarray code, but got caught up with move semantics when trying to use the methods in pyo3, so just having an all-in-one ndarray method now. Added a unit test to check that an ndarray is properly returned.
  • Loading branch information
weiji14 committed Mar 14, 2024
1 parent 54858db commit 1a0b405
Showing 1 changed file with 32 additions and 12 deletions.
44 changes: 32 additions & 12 deletions src/io/geotiff.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,23 @@ impl<R: Read + Seek> CogReader<R> {
Ok(Self { decoder })
}

/// Decode GeoTIFF image to a Vec
fn as_vec(&mut self) -> TiffResult<Vec<f32>> {
/// Decode GeoTIFF image to an [`ndarray::Array`]
fn ndarray(&mut self) -> TiffResult<Array2<f32>> {
// Get image dimensions
let (width, height): (u32, u32) = self.decoder.dimensions()?;

// Get image pixel data
let decode_result = self.decoder.read_image()?;
let image_data: Vec<f32> = match decode_result {
DecodingResult::F32(img_data) => img_data,
_ => unimplemented!("Data types other than float32 are not yet supported."),
};
Ok(image_data)

// Put image pixel data into an ndarray
let vec_data = Array2::from_shape_vec((height as usize, width as usize), image_data)
.map_err(|_| TiffFormatError::InvalidDimensions(height, width))?;

Ok(vec_data)
}

/// Affine transformation for 2D matrix extracted from TIFF tag metadata, used to transform
Expand Down Expand Up @@ -86,15 +95,8 @@ pub fn read_geotiff<R: Read + Seek>(stream: R) -> TiffResult<Array2<f32>> {
// Open TIFF stream with decoder
let mut reader = CogReader::new(stream)?;

// Get image dimensions
let (width, height): (u32, u32) = reader.decoder.dimensions()?;

// Get image pixel data
let img_data: Vec<f32> = reader.as_vec()?;

// Put image pixel data into an ndarray
let vec_data = Array2::from_shape_vec((height as usize, width as usize), img_data)
.map_err(|_| TiffFormatError::InvalidDimensions(height, width))?;
// Decode TIFF into ndarray
let vec_data: Array2<f32> = reader.ndarray()?;

Ok(vec_data)
}
Expand All @@ -104,6 +106,7 @@ mod tests {
use std::io::{Cursor, Seek, SeekFrom};

use geo::AffineTransform;
use ndarray::array;
use object_store::parse_url;
use tempfile::tempfile;
use tiff::encoder::{colortype, TiffEncoder};
Expand Down Expand Up @@ -139,6 +142,23 @@ mod tests {
assert_eq!(arr.mean(), Some(14.0));
}

#[tokio::test]
async fn test_cogreader_ndarray() {
let cog_url: &str = "https://github.com/rasterio/rasterio/raw/1.3.9/tests/data/float32.tif";
let tif_url = Url::parse(cog_url).unwrap();
let (store, location) = parse_url(&tif_url).unwrap();

let result = store.get(&location).await.unwrap();
let bytes = result.bytes().await.unwrap();
let stream = Cursor::new(bytes);

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]])
}

#[tokio::test]
async fn test_cogreader_transform() {
let cog_url: &str =
Expand Down

0 comments on commit 1a0b405

Please sign in to comment.