From 0a689a24426e0efb87a68823fddd9223842fbcde Mon Sep 17 00:00:00 2001 From: Julien Rebetez Date: Sun, 19 Feb 2023 21:12:27 +0100 Subject: [PATCH] Add a Dataset::read_as method that reads a whole image at once This is similar to how GDAL has ReadAsArray on both the band and the whole dataset. --- src/dataset.rs | 134 +++++++++++++++++++++++++++++++++++++++++++- src/raster/mod.rs | 5 +- src/raster/types.rs | 12 ++++ 3 files changed, 147 insertions(+), 4 deletions(-) diff --git a/src/dataset.rs b/src/dataset.rs index ea6f478cb..86ccb5822 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -1,6 +1,6 @@ use ptr::null_mut; use std::convert::TryInto; -use std::mem::MaybeUninit; +use std::mem::{size_of, MaybeUninit}; use std::{ ffi::NulError, ffi::{CStr, CString}, @@ -12,6 +12,7 @@ use std::{ use crate::cpl::CslStringList; use crate::errors::*; use crate::raster::RasterCreationOption; +use crate::raster::{Buffer3D, GdalType, RasterIOExtraArg, ResampleAlg}; use crate::utils::{_last_cpl_err, _last_null_pointer_err, _path_to_c_string, _string}; use crate::vector::{sql, Geometry, OwnedLayer}; use crate::{ @@ -20,10 +21,11 @@ use crate::{ }; use gdal_sys::{ - self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, OGRErr, OGRLayerH, OGRwkbGeometryType, + self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, GDALRWFlag, GDALRasterIOExtraArg, + OGRErr, OGRLayerH, OGRwkbGeometryType, }; use gdal_sys::{GDALFlushCache, OGRGeometryH}; -use libc::{c_double, c_int, c_uint}; +use libc::{c_double, c_int, c_uint, c_void}; #[cfg(all(major_ge_3, minor_ge_1))] use crate::raster::Group; @@ -702,6 +704,94 @@ impl Dataset { Ok(self.child_layer(c_layer)) } + /// Read a [`Buffer`] from this dataset, where `T` implements [`GdalType`]. + /// + /// # Arguments + /// * `window` - the window position from top left + /// * `window_size` - the window size (GDAL will interpolate data if `window_size` != `buffer_size`) + /// * `buffer_size` - the desired size of the 'Buffer' + /// * `e_resample_alg` - the resample algorithm used for the interpolation. Default: `NearestNeighbor`. + /// + /// The returned buffer contains the image data in HWC order. + /// + /// # Example + /// + /// ```rust + /// # fn main() -> gdal::errors::Result<()> { + /// use gdal::Dataset; + /// use gdal::raster::ResampleAlg; + /// let dataset = Dataset::open("fixtures/m_3607824_se_17_1_20160620_sub.tif")?; + /// let size = 2; + /// let buf = dataset.read_as::((0, 0), dataset.raster_size(), (size, size), Some(ResampleAlg::Bilinear))?; + /// assert_eq!(buf.size, (size, size, dataset.raster_count() as usize)); + /// assert_eq!(buf.data, [103, 116, 101, 169, 92, 108, 94, 163, 92, 112, 93, 179, 89, 109, 91, 181]); + /// # Ok(()) + /// # } + /// ``` + pub fn read_as( + &self, + window: (isize, isize), + window_size: (usize, usize), + buffer_size: (usize, usize), + e_resample_alg: Option, + ) -> Result> { + let band_count = self.raster_count() as usize; + let pixels = buffer_size.0 * buffer_size.1 * band_count; + let mut data: Vec = Vec::with_capacity(pixels); + + let resample_alg = e_resample_alg.unwrap_or(ResampleAlg::NearestNeighbour); + + let mut options: GDALRasterIOExtraArg = RasterIOExtraArg { + e_resample_alg: resample_alg, + ..Default::default() + } + .into(); + + let options_ptr: *mut GDALRasterIOExtraArg = &mut options; + + let mut bands: Vec = (1_i32..band_count as i32 + 1_i32).collect(); + + let size_t = size_of::() as i64; + // Safety: the GDALDatasetRasterIOEx writes + // exactly pixel elements into the slice, before we + // read from this slice. This paradigm is suggested + // in the rust std docs + // (https://doc.rust-lang.org/std/vec/struct.Vec.html#examples-18) + let rv = unsafe { + gdal_sys::GDALDatasetRasterIOEx( + self.c_dataset, + GDALRWFlag::GF_Read, + window.0 as c_int, + window.1 as c_int, + window_size.0 as c_int, + window_size.1 as c_int, + data.as_mut_ptr() as *mut c_void, + buffer_size.0 as c_int, + buffer_size.1 as c_int, + T::gdal_ordinal(), + band_count as i32, + bands.as_mut_ptr() as *mut c_int, + // We want to read a HWC + size_t * band_count as i64, + buffer_size.1 as i64 * size_t * band_count as i64, + size_t, + options_ptr, + ) + }; + if rv != CPLErr::CE_None { + return Err(_last_cpl_err(rv)); + } + + unsafe { + data.set_len(pixels); + }; + + Ok(Buffer3D { + size: (buffer_size.0, buffer_size.1, band_count), + data, + }) + } + /// Set the [`Dataset`]'s affine transformation; also called a _geo-transformation_. /// /// This is like a linear transformation preserves points, straight lines and planes. @@ -1288,4 +1378,42 @@ mod tests { let mut ds = Dataset::open(fixture("roads.geojson")).unwrap(); assert!(ds.start_transaction().is_err()); } + + #[test] + fn test_dataset_read_as() { + let ds = Dataset::open(fixture("m_3607824_se_17_1_20160620_sub.tif")).unwrap(); + let size: usize = 4; + let band_count = ds.raster_count() as usize; + // Compare reading the whole dataset at once to reading each band individually + let ds_buf = ds + .read_as::( + (0, 0), + ds.raster_size(), + (size, size), + Some(ResampleAlg::Bilinear), + ) + .unwrap(); + assert_eq!(ds_buf.size, (size, size, band_count)); + + for band_index in 0..band_count { + let band = ds.rasterband(band_index as isize + 1).unwrap(); + let band_buf = band + .read_as::( + (0, 0), + ds.raster_size(), + (size, size), + Some(ResampleAlg::Bilinear), + ) + .unwrap(); + assert_eq!(band_buf.size, (size, size)); + for i in 0..size { + for j in 0..size { + assert_eq!( + band_buf.data[i * size + j], + ds_buf.data[i * size * band_count + j * band_count + band_index], + ); + } + } + } + } } diff --git a/src/raster/mod.rs b/src/raster/mod.rs index c92e5757c..e47240eef 100644 --- a/src/raster/mod.rs +++ b/src/raster/mod.rs @@ -90,7 +90,10 @@ pub use rasterband::{ PaletteInterpretation, RasterBand, RgbaEntry, StatisticsAll, StatisticsMinMax, }; pub use rasterize::{rasterize, BurnSource, MergeAlgorithm, OptimizeMode, RasterizeOptions}; -pub use types::{AdjustedValue, Buffer, ByteBuffer, GdalDataType, GdalType, ResampleAlg}; +pub use types::{ + AdjustedValue, Buffer, Buffer3D, ByteBuffer, GdalDataType, GdalType, RasterIOExtraArg, + ResampleAlg, +}; pub use warp::reproject; /// Key/value pair for passing driver-specific creation options to diff --git a/src/raster/types.rs b/src/raster/types.rs index 83d95877e..55e369917 100644 --- a/src/raster/types.rs +++ b/src/raster/types.rs @@ -460,6 +460,18 @@ impl Buffer { pub type ByteBuffer = Buffer; +#[derive(Debug)] +pub struct Buffer3D { + pub size: (usize, usize, usize), + pub data: Vec, +} + +impl Buffer3D { + pub fn new(size: (usize, usize, usize), data: Vec) -> Buffer3D { + Buffer3D { size, data } + } +} + /// Extra options used to read a raster. /// /// For documentation, see `gdal_sys::GDALRasterIOExtraArg`.