Skip to content

Commit

Permalink
Add a Dataset::read_as method that reads a whole image at once
Browse files Browse the repository at this point in the history
This is similar to how GDAL has ReadAsArray on both the band and the
whole dataset.
  • Loading branch information
julienr committed Feb 26, 2023
1 parent 84fbaf3 commit 0a689a2
Show file tree
Hide file tree
Showing 3 changed files with 147 additions and 4 deletions.
134 changes: 131 additions & 3 deletions src/dataset.rs
Original file line number Diff line number Diff line change
@@ -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},
Expand All @@ -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::{
Expand All @@ -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;
Expand Down Expand Up @@ -702,6 +704,94 @@ impl Dataset {
Ok(self.child_layer(c_layer))
}

/// Read a [`Buffer<T>`] 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::<u8>((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<T: Copy + GdalType>(
&self,
window: (isize, isize),
window_size: (usize, usize),
buffer_size: (usize, usize),
e_resample_alg: Option<ResampleAlg>,
) -> Result<Buffer3D<T>> {
let band_count = self.raster_count() as usize;
let pixels = buffer_size.0 * buffer_size.1 * band_count;
let mut data: Vec<T> = 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<i32> = (1_i32..band_count as i32 + 1_i32).collect();

let size_t = size_of::<T>() 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.
Expand Down Expand Up @@ -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::<u8>(
(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::<u8>(
(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],
);
}
}
}
}
}
5 changes: 4 additions & 1 deletion src/raster/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 12 additions & 0 deletions src/raster/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -460,6 +460,18 @@ impl<T: GdalType> Buffer<T> {

pub type ByteBuffer = Buffer<u8>;

#[derive(Debug)]
pub struct Buffer3D<T: GdalType> {
pub size: (usize, usize, usize),
pub data: Vec<T>,
}

impl<T: GdalType> Buffer3D<T> {
pub fn new(size: (usize, usize, usize), data: Vec<T>) -> Buffer3D<T> {
Buffer3D { size, data }
}
}

/// Extra options used to read a raster.
///
/// For documentation, see `gdal_sys::GDALRasterIOExtraArg`.
Expand Down

0 comments on commit 0a689a2

Please sign in to comment.