eorst 1.0.1

Earth Observation and Remote Sensing Toolkit - library for raster processing pipelines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

//! Block types for raster dataset processing.
//!
//! This module provides types for representing blocks of raster data.

use crate::core_types::RasterType;
use crate::gdal_utils::{open_for_update, raster_from_size, write_bands_to_file};
use crate::metadata::RasterMetadata;
use crate::types::{BlockSize, GeoTransform, Overlap, RasterDataShape, ReadWindow};

use gdal::raster::Buffer;
use log::debug;
use ndarray::{s, Array2, Array3, Array4, Axis};
use std::path::{Path, PathBuf};

impl<U> RasterBlock<U>
where
    U: RasterType,
{
    /// Converts a 2D array into a 3D feature array.
    pub fn into_frc(&self, data: &Array2<U>) -> Array3<U> {
        let n_features = data.shape()[1];
        let n_rows = self.read_window.size.rows as usize - self.overlap.top - self.overlap.bottom;
        let n_cols = self.read_window.size.cols as usize; // - self.overlap.left - self.overlap.right;
        let mut result: Array3<U> = Array3::zeros((n_features, n_rows, n_cols));
        for n_feature in 0..n_features {
            let feature = data.slice(s![.., n_feature]);
            let feature = feature.to_shape((n_rows, n_cols)).unwrap();
            feature.assign_to(result.slice_mut(s![n_feature, .., ..]))
        }

        result
    }

    /// Writes sample data to a vector file.
    pub fn write_samples_feature<T>(&self, data: &Array2<T>, file_name: &PathBuf, na: T)
    where
        T: RasterType,
    {
        let n_bands = data.shape()[1] as isize;
        let block_geotransform = self.geo_transform;
        let epsg_code = self.epsg_code;
        let size_x = self.read_window.size.cols as usize - self.overlap.left - self.overlap.right;
        let size_y = self.read_window.size.rows as usize - self.overlap.top - self.overlap.bottom;
        let block_size: BlockSize = BlockSize {
            rows: size_y,
            cols: size_x,
        };
        let na = T::from(na).unwrap();
        raster_from_size::<T>(
            file_name,
            block_geotransform,
            epsg_code as u32,
            block_size,
            n_bands,
            na,
        );
        let out_ds = open_for_update(Path::new(&file_name));

        // each feature goes to a band
        for feature in 1..n_bands + 1 {
            let mut out_band = out_ds.rasterband(feature as usize).unwrap();
            // out_band.set_no_data_value_u64(Some(na)).unwrap();
            let out_data = data.slice(s![.., feature - 1]);
            let out_data_u8: Vec<T> = out_data.into_iter().map(|v| *v as T).collect();
            let mut data_buffer = Buffer::new((block_size.cols, block_size.rows), out_data_u8);

            out_band
                .write((0, 0), (size_x, size_y), &mut data_buffer)
                .unwrap();
        }
    }

    /// Writes 3D data to a file.
    pub fn write3<T>(&self, data: Array3<T>, out_fn: &PathBuf)
    where
        T: RasterType,
    {
        // create and empty raster with the right size
        debug!("Writing block {:?}", out_fn);
        let block_geotransform = self.geo_transform;
        let block_size: BlockSize = BlockSize {
            rows: self.read_window.size.rows as usize,
            cols: self.read_window.size.cols as usize,
        };
        let n_bands = data.shape()[0] as isize;
        let epsg_code = self.epsg_code;
        let na_value = T::zero();
        raster_from_size::<T>(
            out_fn,
            block_geotransform,
            epsg_code.try_into().unwrap(),
            block_size,
            n_bands,
            na_value,
        );
        let out_ds = open_for_update(out_fn);
        write_bands_to_file(&out_ds, data.view(), (0, 0), (block_size.cols, block_size.rows));
    }

    /// Writes each time step (axis 0) of a 4D array as a separate block file.
    ///
    /// Replaces the duplicated loop pattern found in `apply()`, `apply_with_mask()`
    /// and `mosaic()` in processing.rs.
    pub fn write_time_step_blocks<T>(
        &self,
        data: &Array4<T>,
        tmp_file: &Path,
        file_stem: &str,
        bid: usize,
    ) -> Vec<PathBuf>
    where
        T: RasterType,
    {
        let mut block_fns: Vec<PathBuf> = Vec::new();
        for (tid, result3) in data.axis_iter(Axis(0)).enumerate() {
            let block_fn = tmp_file.with_file_name(format!("{}_{}_{}.tif", file_stem, tid, bid));
            self.write3(result3.to_owned(), &block_fn);
            block_fns.push(block_fn);
        }
        block_fns
    }
}

/// A block representing a partition of a raster dataset for parallel processing.
#[derive(Debug, Clone, PartialEq)]
pub struct RasterBlock<U>
where
    U: RasterType,
{
    /// Index of this block in the dataset
    pub block_index: usize,
    /// Window describing which portion of the dataset this block covers
    pub read_window: ReadWindow,
    /// Size of overlap with neighboring blocks
    pub overlap_size: usize,
    /// Geographic transformation parameters
    pub geo_transform: GeoTransform,
    /// Overlap data with neighboring blocks
    pub overlap: Overlap,
    /// Shape of the raster data
    pub shape: RasterDataShape,
    /// EPSG coordinate reference system code
    pub epsg_code: usize,
    /// Full raster metadata
    pub raster_metadata: RasterMetadata<U>,
}