From a11a53a25acaee2d348ed774850617a82671ba6d 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 the whole image at once This is similar to RasterBand::read_as, but it is reading the whole image at once. Similar to how GDAL has RasterIO on both the dataset and band level. --- src/dataset.rs | 262 +++++++++++++++++++++++++++++++++++++++++++- src/lib.rs | 4 +- src/raster/mod.rs | 5 +- src/raster/types.rs | 12 ++ 4 files changed, 276 insertions(+), 7 deletions(-) diff --git a/src/dataset.rs b/src/dataset.rs index e02487dab..fca18b024 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,10 @@ use crate::{ }; use gdal_sys::{ - self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, OGRErr, OGRGeometryH, OGRLayerH, - OGRwkbGeometryType, + self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, GDALRWFlag, GDALRasterIOExtraArg, + OGRErr, OGRGeometryH, OGRLayerH, OGRwkbGeometryType, }; -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; @@ -282,6 +283,23 @@ impl<'a> Default for LayerOptions<'a> { } } +// This defines multiple ways to layout an image in memory, based on GDAL Python bindings +// which have either 'band' or 'pixel' interleave: +// https://github.com/OSGeo/gdal/blob/301f31b9b74cd67edcdc555f7e7a58db87cbadb2/swig/include/gdal_array.i#L2300 +pub enum ImageInterleaving { + /// This means the image is stored in memory with first the first band, + /// then second band and so on + Band, + /// This means the image is stored in memory with first the value of all bands + /// for the first pixel, then the same for second pixel and so on + Pixel, +} + +pub enum BandSelection { + Subset(Vec), + All, +} + // GDAL Docs state: The returned dataset should only be accessed by one thread at a time. // See: https://gdal.org/api/raster_c_api.html#_CPPv48GDALOpenPKc10GDALAccess // Additionally, VRT Datasets are not safe before GDAL 2.3. @@ -741,6 +759,115 @@ 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 (width, height). GDAL will interpolate data if `window_size` != `buffer_size` + /// * `buffer_size` - the desired size of the 'Buffer' (width, height) + /// * `e_resample_alg` - the resample algorithm used for the interpolation. Default: `NearestNeighbor`. + /// * `interleaving`- The output buffer image layout (see `ImageInterleaving`) + /// * `bands` - A subset of bands to select or BandSelection::All to read all bands + /// + /// # Example + /// + /// ```rust + /// # fn main() -> gdal::errors::Result<()> { + /// use gdal::{Dataset, ImageInterleaving, BandSelection}; + /// 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), ImageInterleaving::Pixel, BandSelection::All)?; + /// 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]); + /// let buf = dataset.read_as::((0, 0), dataset.raster_size(), (size, size), Some(ResampleAlg::Bilinear), ImageInterleaving::Band, BandSelection::All)?; + /// assert_eq!(buf.data, [103, 92, 92, 89, 116, 108, 112, 109, 101, 94, 93, 91, 169, 163, 179, 181]); + /// # Ok(()) + /// # } + /// ``` + pub fn read_as( + &self, + window: (isize, isize), + window_size: (usize, usize), + buffer_size: (usize, usize), + e_resample_alg: Option, + interleaving: ImageInterleaving, + bands: BandSelection, + ) -> Result> { + 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, band_count) = match bands { + BandSelection::Subset(bands) => { + let band_count = bands.len(); + (bands, band_count) + } + BandSelection::All => { + let band_count = self.raster_count() as usize; + let bands = (1_i32..band_count as i32 + 1_i32).collect(); + (bands, band_count) + } + }; + + let pixels = buffer_size.0 * buffer_size.1 * band_count; + let mut data: Vec = Vec::with_capacity(pixels); + let size_t = size_of::() as i64; + + let (pixel_space, line_space, band_space) = match interleaving { + ImageInterleaving::Band => (0, 0, 0), + ImageInterleaving::Pixel => ( + size_t * band_count as i64, + buffer_size.0 as i64 * size_t * band_count as i64, + size_t, + ), + }; + + // Safety: the GDALRasterIOEx 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, + pixel_space, + line_space, + band_space, + 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. @@ -1327,4 +1454,131 @@ mod tests { let mut ds = Dataset::open(fixture("roads.geojson")).unwrap(); assert!(ds.start_transaction().is_err()); } + + #[test] + fn test_dataset_read_as_pixel_interleaving() { + let ds = Dataset::open(fixture("m_3607824_se_17_1_20160620_sub.tif")).unwrap(); + print!("{:?}", ds.raster_size()); + let (width, height) = (4, 7); + let band_count = ds.raster_count() as usize; + + // We compare a single dataset.read_as() to reading band-by-band using + // band.read_as() + let ds_buf = ds + .read_as::( + (0, 0), + (width, height), + (width, height), + Some(ResampleAlg::Bilinear), + ImageInterleaving::Pixel, + BandSelection::All, + ) + .unwrap(); + assert_eq!(ds_buf.size, (width, height, 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), + (width, height), + (width, height), + Some(ResampleAlg::Bilinear), + ) + .unwrap(); + assert_eq!(band_buf.size, (width, height)); + for i in 0..height { + for j in 0..width { + assert_eq!( + band_buf.data[i * width + j], + ds_buf.data[i * width * band_count + j * band_count + band_index], + ); + } + } + } + } + + #[test] + fn test_dataset_read_as_band_interleaving() { + let ds = Dataset::open(fixture("m_3607824_se_17_1_20160620_sub.tif")).unwrap(); + let size: (usize, usize) = (4, 7); + let band_count = ds.raster_count() as usize; + // We compare a single dataset.read_as() to reading band-by-band using + // band.read_as() + let ds_buf = ds + .read_as::( + (0, 0), + size, + size, + Some(ResampleAlg::Bilinear), + ImageInterleaving::Band, + BandSelection::All, + ) + .unwrap(); + assert_eq!(ds_buf.size, (size.0, size.1, 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), size, size, Some(ResampleAlg::Bilinear)) + .unwrap(); + assert_eq!(band_buf.size, size); + assert_eq!( + band_buf.data, + ds_buf.data[band_index * size.0 * size.1..(band_index + 1) * size.0 * size.1] + ); + } + } + + #[test] + fn test_dataset_read_as_band_selection() { + let ds = Dataset::open(fixture("m_3607824_se_17_1_20160620_sub.tif")).unwrap(); + let size: (usize, usize) = (4, 7); + // We compare a single dataset.read_as() to reading band-by-band using + // band.read_as() + let ds_buf = ds + .read_as::( + (0, 0), + size, + size, + Some(ResampleAlg::Bilinear), + ImageInterleaving::Band, + BandSelection::Subset(vec![1, 3]), + ) + .unwrap(); + assert_eq!(ds_buf.size, (size.0, size.1, 2)); + + for (i, band_index) in vec![1, 3].iter().enumerate() { + let band = ds.rasterband(*band_index as isize).unwrap(); + let band_buf = band + .read_as::((0, 0), size, size, Some(ResampleAlg::Bilinear)) + .unwrap(); + assert_eq!(band_buf.size, size); + assert_eq!( + band_buf.data, + ds_buf.data[i * size.0 * size.1..(i + 1) * size.0 * size.1] + ); + } + } + + #[test] + fn test_dataset_read_as_buffer_size() { + let ds = Dataset::open(fixture("m_3607824_se_17_1_20160620_sub.tif")).unwrap(); + let size: (usize, usize) = (4, 7); + let buffer_size: (usize, usize) = (8, 14); + let band_count = ds.raster_count() as usize; + let ds_buf = ds + .read_as::( + (0, 0), + size, + buffer_size, + Some(ResampleAlg::Bilinear), + ImageInterleaving::Band, + BandSelection::All, + ) + .unwrap(); + // We only assert that we get the right buffer size back because checking for explicit + // values is convoluted since GDAL handles the decimation by doing some interpolation + assert_eq!(ds_buf.size, (buffer_size.0, buffer_size.1, band_count)); + } } diff --git a/src/lib.rs b/src/lib.rs index 75de79a80..30d63e031 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -120,8 +120,8 @@ pub mod version; pub mod vsi; pub use dataset::{ - Dataset, DatasetOptions, GdalOpenFlags, GeoTransform, GeoTransformEx, LayerIterator, - LayerOptions, Transaction, + BandSelection, Dataset, DatasetOptions, GdalOpenFlags, GeoTransform, GeoTransformEx, + ImageInterleaving, LayerIterator, LayerOptions, Transaction, }; pub use driver::{Driver, DriverManager}; pub use gcp::{Gcp, GcpRef}; 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 26c1be8ae..0384afde8 100644 --- a/src/raster/types.rs +++ b/src/raster/types.rs @@ -472,6 +472,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`.