eorst 1.0.1

Earth Observation and Remote Sensing Toolkit - library for raster processing pipelines
//! I/O methods for RasterDataset.
//!
//! This module contains methods for reading blocks and writing windows.

use crate::core_types::{RasterData, RasterType};
use crate::types::BlockSize;
use crate::rasterdataset::RasterDataset;
use crate::gdal_utils::{open_for_update, read_raster_band, raster_from_size, write_bands_to_file};

use ndarray::{s, Array3};

impl<R> RasterDataset<R>
where
    R: RasterType,
{
    /// Writes a 3D array to a window in the output file.
    #[doc(hidden)]
    pub fn write_window3<T>(
        &self,
        block_index: usize,
        data: Array3<T>,
        out_fn: &std::path::PathBuf,
        na_value: T,
    ) where
        T: RasterType,
    {
        let trimmed = crate::array_ops::trimm_array3(&data, self.blocks[block_index].overlap_size);

        let block_geotransform = self.blocks[block_index].geo_transform;
        let block_size: BlockSize = BlockSize {
            rows: trimmed.shape()[1],
            cols: trimmed.shape()[2],
        };
        let n_bands = data.shape()[0] as isize;
        let epsg_code = self.metadata.epsg_code;
        let na_value = T::from(na_value).unwrap();
        raster_from_size::<T>(
            out_fn,
            block_geotransform,
            epsg_code,
            block_size,
            n_bands,
            na_value,
        );
        let out_ds = open_for_update(out_fn);
        write_bands_to_file(&out_ds, trimmed, (0, 0), (block_size.cols, block_size.rows));
    }

    /// Reads a block of raster data.
    #[doc(hidden)]
    pub fn read_block<T>(&self, block_id: usize) -> RasterData<T>
    where
        T: RasterType,
    {
        let block = &self.blocks[block_id];
        let read_window = block.read_window;

        let rows = read_window.size.rows as usize;
        let cols = read_window.size.cols as usize;
        let data_shape = (
            self.metadata.shape.times,
            self.metadata.shape.layers,
            rows,
            cols,
        );

        let mut data: RasterData<T> = RasterData::zeros(data_shape);

        for layer in &self.metadata.layers {
            let raster_path = std::path::Path::new(&layer.source);
            let window = (read_window.offset.cols, read_window.offset.rows);
            let window_size = (cols, rows);

            let layer_data = read_raster_band::<T>(raster_path, 1, window, window_size);

            let slice = s![
                layer.time_pos,
                layer.layer_pos,
                ..,
                ..
            ];

            data.slice_mut(slice).assign(&layer_data);
        }
        data
    }

    /// Reads a specific layer from a block.
    #[doc(hidden)]
    pub fn read_block_layer_idx<T>(&self, block_id: usize, layer_idx: usize) -> RasterData<T>
    where
        T: RasterType,
    {
        let block = &self.blocks[block_id];
        let read_window = block.read_window;

        let rows = read_window.size.rows as usize;
        let cols = read_window.size.cols as usize;
        let data_shape = (
            self.metadata.shape.times,
            self.metadata.shape.layers,
            rows,
            cols,
        );

        let mut data: RasterData<T> = RasterData::zeros(data_shape);

        let layer = &self.metadata.layers[layer_idx];
        let raster_path = std::path::Path::new(&layer.source);
        let window = (read_window.offset.cols, read_window.offset.rows);
        let window_size = (cols, rows);

        let layer_data = read_raster_band::<T>(raster_path, 1, window, window_size);

        let slice = s![
            layer.time_pos,
            layer.layer_pos,
            ..,
            ..
        ];

        data.slice_mut(slice).assign(&layer_data);

        data
    }
}