raster_tools/
utils.rs

1//! Utilities to create, read and write raster datasets.
2
3use gdal::{DatasetOptions, DriverManager};
4use gdal::GdalOpenFlags;
5use rasters::Result;
6use std::fs::File;
7use std::path::Path;
8use std::path::PathBuf;
9
10pub type InputArgs = PathBuf;
11pub struct OutputArgs {
12    pub path: PathBuf,
13    pub driver: String,
14}
15
16use anyhow::Context;
17use gdal::Dataset;
18
19pub fn read_dataset(path: &Path) -> Result<Dataset> {
20    Ok(Dataset::open(&path).with_context(|| format!("reading dataset {}", path.display()))?)
21}
22
23pub fn edit_dataset(path: &Path) -> Result<Dataset> {
24    Ok(Dataset::open_ex(
25        &path,
26        DatasetOptions {
27            open_flags: GdalOpenFlags::GDAL_OF_UPDATE,
28            ..Default::default()
29        },
30    )
31    .with_context(|| format!("editing dataset {}", path.display()))?)
32}
33
34use gdal::raster::GdalType;
35pub fn create_output_raster<T: GdalType>(
36    arg: &OutputArgs,
37    ds: &Dataset,
38    num_bands: isize,
39    no_val: Option<f64>,
40) -> Result<Dataset> {
41    let mut out_ds = {
42        let driver = DriverManager::get_driver_by_name(&arg.driver)?;
43        let (width, height) = ds.raster_size();
44        driver
45            .create_with_band_type::<T, _>(&arg.path, width as isize, height as isize, num_bands)
46            .with_context(|| format!("creating dataset {}", arg.path.display()))?
47    };
48    if let Some(no_val) = no_val {
49        for i in 1..=num_bands {
50            out_ds.rasterband(i)?.set_no_data_value(Some(no_val))?;
51        }
52    }
53    if let Ok(gt) = ds.geo_transform() {
54        out_ds.set_geo_transform(&gt)?;
55    }
56    out_ds.set_projection(&ds.projection())?;
57    Ok(out_ds)
58}
59
60#[cfg(test)]
61mod test {
62    use super::*;
63    use rand::*;
64    use tempdir::TempDir;
65
66    const WIDTH: usize = 16;
67    const HEIGHT: usize = 32;
68
69    #[test]
70    fn create_read_update_ds() -> Result<()> {
71        let driver = DriverManager::get_driver_by_name("GTIFF")?;
72        let tmp_dir = TempDir::new("rasters_test").unwrap();
73        let path = tmp_dir.path().join("foo.tif");
74
75        // Create empty raster
76        {
77            driver.create_with_band_type::<f64, _>(&path, WIDTH as isize, HEIGHT as isize, 1)?;
78        }
79
80        // Create random data
81        let data = {
82            use gdal::raster::Buffer;
83            let mut data: Vec<f64> = Vec::with_capacity(WIDTH * HEIGHT);
84
85            let mut rng = thread_rng();
86            for _ in 0..(WIDTH * HEIGHT) {
87                data.push(rng.gen());
88            }
89            Buffer::new((WIDTH, HEIGHT), data)
90        };
91
92        // Write some dataset
93        {
94            let ds = edit_dataset(&path)?;
95            let mut band = ds.rasterband(1)?;
96            let (width, height) = ds.raster_size();
97
98            assert_eq!(width, WIDTH);
99            assert_eq!(height, HEIGHT);
100            assert_eq!(ds.raster_count(), 1);
101
102            band.write((0, 0), (width, height), &data)?;
103        }
104
105        // Read data
106        {
107            let ds = read_dataset(&path)?;
108            let band = ds.rasterband(1)?;
109            let rdata = band.read_band_as::<f64>()?;
110
111            assert_eq!(rdata.data, data.data);
112        }
113
114        Ok(())
115    }
116}
117
118pub fn write_bin<T: serde::Serialize>(path: &Path, data: &T) -> Result<()> {
119    let file = File::create(path)?;
120    let buf = std::io::BufWriter::with_capacity(0x100000, file);
121    serde_cbor::to_writer(buf, data)?;
122    Ok(())
123}
124
125pub fn read_bin<T: for<'a> serde::Deserialize<'a>>(path: &Path) -> Result<T> {
126    let file = File::open(path)?;
127    let file = unsafe { memmap::MmapOptions::new().map(&file)? };
128    Ok(serde_cbor::from_slice(file.as_ref())?)
129}
130
131use serde::Serialize;
132pub fn write_json<T: Serialize>(path: &Path, json: &T) -> Result<()> {
133    let file = File::create(path)?;
134    let buf = std::io::BufWriter::with_capacity(0x100000, file);
135    Ok(serde_json::to_writer(buf, json)?)
136}
137
138pub fn print_json<T: Serialize>(json: &T) -> Result<()> {
139    let writer = std::io::BufWriter::new(std::io::stdout());
140    Ok(serde_json::to_writer(writer, json)?)
141}
raster_tools/utils.rs

raster_tools/
utils.rs
