Skip to content

Commit

Permalink
✨ CogReader ndarray method to decode TIFF into an ndarray Array (#10)
Browse files Browse the repository at this point in the history
* ♻️ Turn image decoding into a Vec code into a method call

Making the decoder.read_image() code an impl method under the CogReader struct. Reworded the previous panic message into an unimplemented error with a more helpful message that dtypes other than float32 are not yet supported.

* ✨ Implement ndarray method in CogReader struct

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 authored Mar 15, 2024
1 parent d02dd13 commit d4786a1
Showing 1 changed file with 39 additions and 11 deletions.
50 changes: 39 additions & 11 deletions src/io/geotiff.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,25 @@ impl<R: Read + Seek> CogReader<R> {
Ok(Self { decoder })
}

/// 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."),
};

// 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
/// image pixel (row, col) coordinates to and from geographic/projected (x, y) coordinates.
///
Expand Down Expand Up @@ -76,17 +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 DecodingResult::F32(img_data) = reader.decoder.read_image()? else {
panic!("Cannot read band data")
};

// 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 @@ -96,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 @@ -131,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 d4786a1

Please sign in to comment.