eorst/rasterdataset/
rasterization.rs

1//! Rasterization methods for RasterDataset.
2//!
3//! This module contains methods for burning vector values into raster blocks.
4
5use crate::blocks::RasterBlock;
6use crate::core_types::RasterType;
7use crate::gdal_utils::create_rayon_pool;
8use crate::parallel_writer;
9use crate::rasterdataset::RasterDataset;
10
11use anyhow::Result as GdalResult;
12use gdal::raster::RasterizeOptions;
13use gdal::spatial_ref::CoordTransform;
14use gdal::spatial_ref::SpatialRef;
15use gdal::vector::{Layer, LayerAccess};
16use gdal::{Dataset, DriverManager};
17use kdam::par_tqdm;
18use ndarray::Array3;
19use num_traits::NumCast;
20use rayon::prelude::*;
21use std::collections::HashMap;
22use std::path::Path;
23
24fn get_layer<'a>(dataset: &'a Dataset, layer_name: Option<&str>) -> GdalResult<Layer<'a>> {
25    match layer_name {
26        Some(name) => dataset.layer_by_name(name).map_err(|e| anyhow::anyhow!(e)),
27        None => dataset.layer(0).map_err(|e| anyhow::anyhow!(e)),
28    }
29}
30
31impl<R> RasterDataset<R>
32where
33    R: RasterType,
34{
35    /// Burns vector values into the raster dataset.
36    ///
37    /// Rasterizes geometries from a vector file by processing blocks in parallel
38    /// and writing directly to the output GeoTIFF via [`ParallelGeoTiffWriter`],
39    /// avoiding intermediate files and a mosaic step.
40    pub fn rasterize(
41        &self,
42        vector_fn: &Path,
43        column_name: Option<&str>,
44        map: Option<HashMap<String, i32>>,
45        layer_name: Option<&str>,
46        n_cpus: usize,
47        rasterize_options: Option<RasterizeOptions>,
48        out_file: &Path,
49    ) {
50        log::debug!("Starting rasterize.");
51        gdal::config::set_config_option("OSR_DEFAULT_AXIS_MAPPING_STRATEGY", "TRADITIONAL_GIS_ORDER").unwrap();
52
53        if column_name.is_none() {
54            log::warn!("No --column specified, using feature FID as burn value");
55        }
56
57        let n_bands = self.metadata.shape.layers * self.metadata.shape.times;
58
59        // Pre-create output GeoTIFF
60        parallel_writer::create_output_geotiff::<R>(
61            out_file,
62            &self.metadata.geo_transform,
63            self.metadata.epsg_code as u32,
64            self.metadata.shape.cols,
65            self.metadata.shape.rows,
66            n_bands,
67            R::zero(),
68        )
69        .expect("Failed to create output GeoTIFF");
70
71        let writer = parallel_writer::ParallelGeoTiffWriter::new(
72            out_file.to_path_buf(),
73            self.metadata.geo_transform,
74            self.metadata.epsg_code as u32,
75            self.metadata.shape.cols,
76            self.metadata.shape.rows,
77            n_bands,
78        );
79
80        let pool = create_rayon_pool(n_cpus);
81
82        pool.install(|| {
83            par_tqdm!(self.blocks.to_owned().into_par_iter()).for_each(
84                |raster_block: RasterBlock<R>| {
85                    let dataset = Dataset::open(vector_fn).expect("unable to open dataset");
86                    let mut layer = get_layer(&dataset, layer_name).unwrap();
87
88                    // Compute block bounding box in raster CRS
89                    let x_ul = raster_block.geo_transform.x_ul;
90                    let y_ul = raster_block.geo_transform.y_ul;
91                    let x2 = x_ul
92                        + (raster_block.geo_transform.x_res
93                            * raster_block.read_window.size.cols as f64);
94                    let y2 = y_ul
95                        + (raster_block.geo_transform.y_res
96                            * raster_block.read_window.size.rows as f64);
97                    let (xmin, xmax) = if x_ul < x2 { (x_ul, x2) } else { (x2, x_ul) };
98                    let (ymin, ymax) = if y_ul < y2 { (y_ul, y2) } else { (y2, y_ul) };
99
100                    let src = SpatialRef::from_epsg(raster_block.epsg_code as u32)
101                        .expect("Invalid raster EPSG code");
102                    let dst = layer
103                        .spatial_ref()
104                        .expect("Vector layer has no CRS defined");
105
106                    // Create and transform spatial filter polygon
107                    let mut filter_geom = gdal::vector::Geometry::from_wkt(&format!(
108                        "POLYGON (({} {}, {} {}, {} {}, {} {}, {} {}))",
109                        xmin, ymin, xmin, ymax, xmax, ymax, xmax, ymin, xmin, ymin
110                    ))
111                    .unwrap();
112                    filter_geom.set_spatial_ref(src.clone());
113                    let filter_geom = if src != dst {
114                        let tx = CoordTransform::new(&src, &dst)
115                            .expect("Failed to create coordinate transform");
116                        filter_geom
117                            .transform(&tx)
118                            .expect("Failed to transform bbox into layer CRS")
119                    } else {
120                        filter_geom
121                    };
122                    layer.set_spatial_filter(&filter_geom);
123
124                    // Collect geometries and burn values from features in this block
125                    let mut geometries = Vec::new();
126                    let mut burn_values: Vec<R> = Vec::new();
127                    let available_fields: Vec<String> = if column_name.is_some() {
128                        layer.defn().fields().map(|f| f.name()).collect()
129                    } else {
130                        Vec::new()
131                    };
132
133                    for feature in layer.features() {
134                        let g = feature.geometry().unwrap().to_owned();
135                        let g = g.buffer(0., 8).unwrap();
136                        let g = if src != dst {
137                            let tx = CoordTransform::new(&dst, &src)
138                                .expect("Failed to create coordinate transform");
139                            g.transform(&tx)
140                                .expect("Failed to transform geom into raster CRS")
141                        } else {
142                            g
143                        };
144
145                        if let Some(col) = column_name {
146                            let field_index = match feature.field_index(col) {
147                                Ok(idx) => idx,
148                                Err(_) => panic!(
149                                    "Column '{}' not found in vector layer. Available fields: {}",
150                                    col,
151                                    available_fields.join(", ")
152                                ),
153                            };
154                            if let Some(field_value) =
155                                feature.field(field_index).expect("Invalid column name!")
156                            {
157                                let v = match field_value {
158                                    gdal::vector::FieldValue::IntegerValue(i) => {
159                                        NumCast::from(i).unwrap_or(0)
160                                    }
161                                    gdal::vector::FieldValue::RealValue(r) => {
162                                        NumCast::from(r).unwrap()
163                                    }
164                                    gdal::vector::FieldValue::StringValue(s) => match &map {
165                                        Some(hmap) => match hmap.get(&s) {
166                                            Some(v) => *v,
167                                            None => panic!(
168                                                "String '{}' not found in lookup table",
169                                                s
170                                            ),
171                                        },
172                                        None => panic!(
173                                            "Field '{}' is a string ('{}'), but no string→number map was supplied",
174                                            col, s
175                                        ),
176                                    },
177                                    _ => panic!("The field to rasterize has to be numeric"),
178                                };
179                                burn_values.push(NumCast::from(v).expect("Cant cast!"));
180                            } else {
181                                panic!("{}", format!("Field {col} does not exist"));
182                            }
183                        } else {
184                            let fid = feature.fid().expect("Feature has no FID") as i32;
185                            burn_values.push(NumCast::from(fid).expect("Cant cast!"));
186                        }
187                        geometries.push(g);
188                    }
189
190                    // Create an in-memory raster for this block, rasterize into it
191                    let block_rows = raster_block.read_window.size.rows as usize;
192                    let block_cols = raster_block.read_window.size.cols as usize;
193
194                    let mem_driver = DriverManager::get_driver_by_name("MEM")
195                        .expect("GDAL MEM driver not available");
196                    let mut mem_dataset = mem_driver
197                        .create_with_band_type::<R, _>("", block_cols, block_rows, 1)
198                        .expect("Failed to create in-memory dataset");
199                    mem_dataset
200                        .set_geo_transform(&raster_block.geo_transform.to_array())
201                        .expect("Failed to set geo transform");
202                    let srs = SpatialRef::from_epsg(raster_block.epsg_code as u32)
203                        .expect("Invalid EPSG code");
204                    mem_dataset
205                        .set_spatial_ref(&srs)
206                        .expect("Failed to set CRS");
207
208                    let opts = rasterize_options.unwrap_or_else(|| RasterizeOptions::default());
209                    let burn_values_f64: Vec<f64> = burn_values
210                        .iter()
211                        .map(|v| NumCast::from(*v).unwrap_or(0.0))
212                        .collect();
213                    gdal::raster::rasterize(
214                        &mut mem_dataset,
215                        &[1],
216                        &geometries,
217                        burn_values_f64.as_slice(),
218                        Some(opts),
219                    )
220                    .expect("Rasterize failed");
221
222                    // Read rasterized data back from in-memory dataset
223                    let band = mem_dataset.rasterband(1).expect("No band 1");
224                    let buffer: gdal::raster::Buffer<R> = band
225                        .read_as::<R>(
226                            (0, 0),
227                            (block_cols, block_rows),
228                            (block_cols, block_rows),
229                            None,
230                        )
231                        .expect("Failed to read rasterized data");
232                    let (_shape, data) = buffer.into_shape_and_vec();
233
234                    let array = Array3::from_shape_vec((1, block_rows, block_cols), data)
235                        .expect("Shape mismatch reading rasterized data");
236
237                    // Write to output via parallel writer (no mosaic step needed)
238                    parallel_writer::write_block(
239                        &writer,
240                        array.view(),
241                        raster_block.read_window.clone(),
242                    )
243                    .expect("Failed to write block to output");
244                },
245            )
246        });
247    }
248
249    /// Burns vector values into the raster dataset, outputting a Cloud Optimized GeoTIFF.
250    ///
251    /// COG variant of [`rasterize`](Self::rasterize). Follows the same rasterization
252    /// logic but produces a proper COG with overviews and IFD reordering.
253    pub fn rasterize_cog(
254        &self,
255        vector_fn: &Path,
256        column_name: Option<&str>,
257        map: Option<HashMap<String, i32>>,
258        layer_name: Option<&str>,
259        n_cpus: usize,
260        rasterize_options: Option<RasterizeOptions>,
261        out_file: &Path,
262        config: &crate::types::OutputConfig,
263    ) {
264        use crate::types::OutputFormat;
265
266        match config.format {
267            OutputFormat::GeoTiff => {
268                self.rasterize(vector_fn, column_name, map, layer_name, n_cpus, rasterize_options, out_file);
269            }
270            OutputFormat::GeoTiffOverviews => {
271                self.rasterize(vector_fn, column_name, map, layer_name, n_cpus, rasterize_options, out_file);
272
273                let n_bands = self.metadata.shape.layers * self.metadata.shape.times;
274                let writer = parallel_writer::ParallelGeoTiffWriter::new(
275                    out_file.to_path_buf(),
276                    self.metadata.geo_transform,
277                    self.metadata.epsg_code as u32,
278                    self.metadata.shape.cols,
279                    self.metadata.shape.rows,
280                    n_bands,
281                );
282                writer.build_overviews(&config.overview_resampling, &config.overview_levels).ok();
283            }
284            OutputFormat::COG => {
285                let intermediate = std::path::PathBuf::from(crate::gdal_utils::create_temp_file("tif"));
286
287                self.rasterize(vector_fn, column_name, map, layer_name, n_cpus, rasterize_options, &intermediate);
288
289                crate::gdal_utils::translate_to_cog(
290                    &intermediate, out_file,
291                    &config.compression,
292                    &config.overview_resampling,
293                ).ok();
294                std::fs::remove_file(&intermediate).ok();
295            }
296        }
297    }
298}
eorst/rasterdataset/rasterization.rs

eorst/rasterdataset/
rasterization.rs
