use gdal::{raster::RasterCreationOption, spatial_ref::SpatialRef, Driver};
use ndarray::{s, Array3, ArrayView2, ArrayView3};
use std::path::Path;
use std::path::PathBuf;
use uuid::Uuid;
#[cfg(feature = "use_opencv")]
use opencv::{
core::{Mat_AUTO_STEP, Scalar},
imgproc,
prelude::*,
};
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Index2d {
pub col: usize,
pub row: usize,
}
#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub(crate) struct Coordinates {
pub(crate) x: f64,
pub(crate) y: f64,
}
#[derive(Debug, Copy, Clone)]
pub struct Rectangle {
pub left: usize,
pub top: usize,
pub right: usize,
pub bottom: usize,
}
#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub struct BlockSize {
pub rows: usize,
pub cols: usize,
}
#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub struct ImageSize {
pub rows: usize,
pub cols: usize,
}
#[derive(Default, Debug, PartialEq, Clone, Copy)]
pub struct GeoTransform {
pub x_ul: f64,
pub x_res: f64,
pub x_rot: f64,
pub y_ul: f64,
pub y_rot: f64,
pub y_res: f64,
}
#[derive(PartialEq, Debug, Clone, Copy)]
pub struct BlockAttributes {
pub block_index: usize,
pub read_window: ReadWindow,
pub overlap_size: usize,
pub geo_transform: GeoTransform,
pub overlap: Overlap,
}
#[derive(PartialEq, Debug, Clone, Copy)]
pub struct Offset {
pub rows: isize,
pub cols: isize,
}
#[derive(PartialEq, Debug, Clone, Copy)]
pub struct Size {
pub rows: isize,
pub cols: isize,
}
#[derive(PartialEq, Debug, Clone, Copy)]
pub struct ReadWindow {
pub offset: Offset,
pub size: Size,
}
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Overlap {
pub left: usize,
pub top: usize,
pub right: usize,
pub bottom: usize,
}
pub fn add_cols_left<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
let n = n.to_owned();
if n > 0 {
let r = original.shape()[1];
let c = original.shape()[2];
let b = original.shape()[0];
let mut result: Array3<T> = ndarray::Array::zeros((b, r, c + n));
result.slice_mut(s![.., .., n..]).assign(&original);
let first_data_cols = result.slice(s![.., .., n..n + n]).to_owned();
result.slice_mut(s![.., .., 0..n]).assign(&first_data_cols);
return result;
} else {
return original.to_owned();
}
}
pub fn add_rows_top<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
let n = n.to_owned();
if n > 0 {
let r = original.shape()[1];
let c = original.shape()[2];
let b = original.shape()[0];
let mut result: Array3<T> = ndarray::Array::zeros((b, r + n, c));
result.slice_mut(s![.., n.., ..]).assign(&original);
let first_data_rows = result.slice(s![.., n..n + n, ..]).to_owned();
result.slice_mut(s![.., 0..n, ..]).assign(&first_data_rows);
return result;
} else {
return original.to_owned();
}
}
pub fn add_cols_right<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
let n = n.to_owned();
if n > 0 {
let r = original.shape()[1];
let c = original.shape()[2];
let b = original.shape()[0];
let mut result: Array3<T> = ndarray::Array::zeros((b, r, c + n));
let idx = -1 * n as isize;
result.slice_mut(s![.., .., ..idx]).assign(&original);
let f = c as isize - n as isize;
let t = c as isize;
let last_data_cols = result.slice(s![.., .., f..t]).to_owned();
result.slice_mut(s![.., .., c..]).assign(&last_data_cols);
return result;
} else {
return original.to_owned();
}
}
pub fn add_rows_bottom<T>(original: &Array3<T>, n: &usize) -> Array3<T>
where
T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
let n = n.to_owned();
if n > 0 {
let r = original.shape()[1];
let c = original.shape()[2];
let b = original.shape()[0];
let mut result: Array3<T> = ndarray::Array::zeros((b, r + n, c));
let idx = -1 * n as isize;
result.slice_mut(s![.., ..idx, ..]).assign(&original);
let f = r as isize - n as isize;
let t = r as isize;
let last_data_row = result.slice(s![.., f..t, ..]).to_owned();
result.slice_mut(s![.., r.., ..]).assign(&last_data_row);
return result;
} else {
return original.to_owned();
}
}
pub fn trimm_array3<'a, T>(array: &'a Array3<T>, overlap_size: usize) -> ArrayView3<'a, T> {
let min_row = overlap_size;
let max_row = array.shape()[1] - overlap_size;
let min_col = overlap_size;
let max_col = array.shape()[2] - overlap_size;
array.slice(s![.., min_row..max_row, min_col..max_col])
}
impl GeoTransform {
pub fn to_array(&self) -> [f64; 6] {
let array: [f64; 6] = [
self.x_ul, self.x_res, self.x_rot, self.y_ul, self.y_rot, self.y_res,
];
array
}
}
pub fn raster_from_size<T>(
file_name: &PathBuf,
geo_transform: GeoTransform,
epsg_code: u32,
block_size: BlockSize,
n_bands: isize,
) where
T: gdal::raster::GdalType + Copy + num_traits::identities::Zero,
{
let size_x = block_size.cols;
let size_y = block_size.rows;
let driver = Driver::get("GTiff").unwrap();
let options = [RasterCreationOption {
key: "COMPRESS",
value: "LZW",
}];
let mut dataset = driver
.create_with_band_type_with_options::<T>(
file_name.to_str().unwrap(),
size_x as isize,
size_y as isize,
n_bands,
&options,
)
.unwrap();
dataset
.set_geo_transform(&geo_transform.to_array())
.unwrap();
let srs = SpatialRef::from_epsg(epsg_code).unwrap();
dataset.set_spatial_ref(&srs).unwrap();
}
pub fn create_temp_file(ext: &str) -> String {
let dir = Path::new("/tmp/");
println!("{}", dir.to_str().unwrap());
let file_name = format!("{}.{}", Uuid::new_v4(), ext);
let file_name = dir.join(file_name);
let file_name_string = file_name.into_os_string().into_string().unwrap();
file_name_string
}
pub fn rect_view<'a, T>(
array: &'a ArrayView2<T>,
rect: Rectangle,
indices: Index2d,
) -> ArrayView2<'a, T> {
let max_row = indices.row + rect.bottom + 1;
let min_row = indices.row - rect.top;
let max_col = indices.col + rect.right + 1;
let min_col = indices.col - rect.left;
array.slice(s![min_row..max_row, min_col..max_col])
}
pub fn transpose<T>(v: Vec<Vec<T>>) -> Vec<Vec<T>>
where
T: Clone,
{
assert!(!v.is_empty());
(0..v[0].len())
.map(|i| v.iter().map(|inner| inner[i].clone()).collect::<Vec<T>>())
.collect()
}
#[cfg(feature = "use_opencv")]
pub fn arrayview2_to_mat<T>(data: ArrayView2<T>) -> Mat
where
T: gdal::raster::GdalType + opencv::prelude::DataType,
{
let rows = data.shape()[0];
let data_slice = data.as_slice().unwrap();
let m = Mat::from_slice(data_slice).unwrap();
let m = m.reshape(0, rows as i32).unwrap();
m
}