python/src/raster.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "bindings.h" | ||
| 2 | #include "common.h" | ||
| 3 | |||
| 4 | #include "terraingen/raster/gaussian_blur.h" | ||
| 5 | #include "terraingen/raster/geo_transform.h" | ||
| 6 | #include "terraingen/raster/geotiff/geotiff.h" | ||
| 7 | #include "terraingen/raster/grid.h" | ||
| 8 | #include "terraingen/raster/raster.h" | ||
| 9 | #include "terraingen/raster/raster_stats.h" | ||
| 10 | #include "terraingen/raster/resample.h" | ||
| 11 | |||
| 12 | #include <nanobind/ndarray.h> | ||
| 13 | #include <nanobind/operators.h> | ||
| 14 | |||
| 15 | #include <cstdio> | ||
| 16 | #include <string> | ||
| 17 | #include <vector> | ||
| 18 | |||
| 19 | namespace terraingen::python { | ||
| 20 | |||
| 21 | namespace { | ||
| 22 | |||
| 23 | using RasterF = Raster<float>; | ||
| 24 | using NumpyView = nb::ndarray<nb::numpy, float, nb::ndim<2>, nb::c_contig>; | ||
| 25 | using ArrayApiView = nb::ndarray<nb::array_api, float, nb::ndim<2>, nb::c_contig>; | ||
| 26 | |||
| 27 | /// @brief A view of the raster `self` that keeps it alive. | ||
| 28 | 467 | template <typename View> View MakeView(nb::handle self) { | |
| 29 | 467 | RasterF& raster = nb::cast<RasterF&>(self); | |
| 30 | 934 | return View(raster.Values().data(), {raster.Height(), raster.Width()}, self); | |
| 31 | } | ||
| 32 | |||
| 33 | 2 | std::string Repr(const GeoTransform& t) { | |
| 34 | char buffer[256]; | ||
| 35 | 2 | std::snprintf(buffer, sizeof(buffer), "GeoTransform(%.17g, %.17g, %.17g, %.17g, %.17g, %.17g)", | |
| 36 | 2 | t.originX, t.pixelWidth, t.rowRotation, t.originY, t.colRotation, t.pixelHeight); | |
| 37 | 4 | return buffer; | |
| 38 | } | ||
| 39 | |||
| 40 | 44 | void BindGeoTransform(nb::module_& m) { | |
| 41 | 44 | nb::class_<GeoTransform>(m, "GeoTransform", | |
| 42 | "Affine pixel-to-world mapping matching GDAL's six-coefficient " | ||
| 43 | "geotransform. Pixel (0, 0) is the outer corner of the first pixel.") | ||
| 44 | 88 | .def(nb::init<>()) | |
| 45 | 44 | .def_rw("origin_x", &GeoTransform::originX) | |
| 46 | 44 | .def_rw("pixel_width", &GeoTransform::pixelWidth) | |
| 47 | 44 | .def_rw("row_rotation", &GeoTransform::rowRotation) | |
| 48 | 44 | .def_rw("origin_y", &GeoTransform::originY) | |
| 49 | 44 | .def_rw("col_rotation", &GeoTransform::colRotation) | |
| 50 | 44 | .def_rw("pixel_height", &GeoTransform::pixelHeight) | |
| 51 | 44 | .def_static("north_up", &GeoTransform::NorthUp, "origin_x"_a, "origin_y"_a, "pixel_size"_a, | |
| 52 | "North-up transform with square pixels.") | ||
| 53 | 88 | .def_static("from_gdal", &GeoTransform::FromGdal, "coefficients"_a, | |
| 54 | "Build from GDAL's coefficient order.") | ||
| 55 | 44 | .def( | |
| 56 | "to_gdal", | ||
| 57 | ✗ | [](const GeoTransform& t) { | |
| 58 | 2 | const auto c = t.ToGdal(); | |
| 59 | 4 | return nb::make_tuple(c[0], c[1], c[2], c[3], c[4], c[5]); | |
| 60 | }, | ||
| 61 | "Coefficients in GDAL's order.") | ||
| 62 | 44 | .def("pixel_to_world", &GeoTransform::PixelToWorld, "col"_a, "row"_a, | |
| 63 | "Map pixel coordinates to world coordinates (x, y).") | ||
| 64 | 44 | .def("world_to_pixel", &GeoTransform::WorldToPixel, "x"_a, "y"_a, | |
| 65 | "Map world coordinates to pixel coordinates (col, row). Raises ValueError if the " | ||
| 66 | "transform is not invertible.") | ||
| 67 | 44 | .def(nb::self == nb::self) | |
| 68 | 88 | .def("__repr__", &Repr); | |
| 69 | 44 | } | |
| 70 | |||
| 71 | 14 | RasterF RasterFromArray(nb::ndarray<const float, nb::ndim<2>, nb::device::cpu> array, | |
| 72 | const GeoTransform& transform) { | ||
| 73 | 14 | RasterF raster(array.shape(1), array.shape(0)); | |
| 74 | const auto view = array.view(); | ||
| 75 | 152 | for (std::size_t row = 0; row < raster.Height(); ++row) { | |
| 76 | 2969 | for (std::size_t col = 0; col < raster.Width(); ++col) { | |
| 77 | 2831 | raster(col, row) = view(row, col); | |
| 78 | } | ||
| 79 | } | ||
| 80 | 14 | raster.SetTransform(transform); | |
| 81 | 28 | return raster; | |
| 82 | } | ||
| 83 | |||
| 84 | 3 | nb::object RasterArray(nb::handle self, nb::object dtype, nb::object copy) { | |
| 85 | 3 | nb::object values = nb::cast(MakeView<NumpyView>(self)); | |
| 86 | 3 | if (!dtype.is_none()) { | |
| 87 | 2 | values = values.attr("astype")(dtype); | |
| 88 | } | ||
| 89 | 4 | if (!copy.is_none() && nb::cast<bool>(copy)) { | |
| 90 | 2 | values = values.attr("copy")(); | |
| 91 | } | ||
| 92 | 3 | return values; | |
| 93 | } | ||
| 94 | |||
| 95 | 44 | void BindArrayProtocols(nb::class_<RasterF>& raster) { | |
| 96 | 176 | raster.def("__array__", &RasterArray, "dtype"_a = nb::none(), "copy"_a = nb::none()) | |
| 97 | 44 | .def( | |
| 98 | "__dlpack__", | ||
| 99 | 47 | [](nb::handle self, nb::kwargs kwargs) { | |
| 100 | 9 | return nb::cast(MakeView<ArrayApiView>(self)).attr("__dlpack__")(**kwargs); | |
| 101 | }, | ||
| 102 | 44 | "kwargs"_a) | |
| 103 | 44 | .def("__dlpack_device__", [](nb::handle self) { | |
| 104 | 3 | return nb::cast(MakeView<ArrayApiView>(self)).attr("__dlpack_device__")(); | |
| 105 | }); | ||
| 106 | 44 | } | |
| 107 | |||
| 108 | 44 | void BindRasterClass(nb::module_& m) { | |
| 109 | nb::class_<RasterF> raster(m, "Raster", | ||
| 110 | "Single-band float32 raster, stored row-major with row 0 at the " | ||
| 111 | "top. Shares memory with NumPy and any DLPack consumer without " | ||
| 112 | 44 | "copying."); | |
| 113 | 132 | raster.def(nb::init<std::size_t, std::size_t, float>(), "width"_a, "height"_a, "fill"_a = 0.0f) | |
| 114 | 132 | .def_static("from_array", &RasterFromArray, "array"_a, "transform"_a = GeoTransform{}, | |
| 115 | "Copy a 2D array of shape (height, width) from NumPy, DLPack or the buffer " | ||
| 116 | "protocol.") | ||
| 117 | 44 | .def_prop_ro("width", &RasterF::Width) | |
| 118 | 88 | .def_prop_ro("height", &RasterF::Height) | |
| 119 | 44 | .def_prop_ro("shape", | |
| 120 | 98 | [](const RasterF& r) { return nb::make_tuple(r.Height(), r.Width()); }) | |
| 121 | 44 | .def_prop_ro("values", &MakeView<NumpyView>, | |
| 122 | "Writeable NumPy view of shape (height, width) that keeps the raster alive.") | ||
| 123 | 44 | .def_prop_rw("transform", &RasterF::Transform, &RasterF::SetTransform) | |
| 124 | 44 | .def_prop_rw("nodata", &RasterF::NoData, &RasterF::SetNoData) | |
| 125 | 88 | .def("cell_center", &RasterF::CellCenter, "col"_a, "row"_a, | |
| 126 | "World coordinates of a cell's centre.") | ||
| 127 | 44 | .def("__repr__", [](const RasterF& r) { | |
| 128 | 2 | return "Raster(width=" + std::to_string(r.Width()) + | |
| 129 | 3 | ", height=" + std::to_string(r.Height()) + ")"; | |
| 130 | }); | ||
| 131 | 44 | BindArrayProtocols(raster); | |
| 132 | 44 | } | |
| 133 | |||
| 134 | 44 | void BindStats(nb::module_& m) { | |
| 135 | 44 | nb::class_<RasterStats>(m, "RasterStats", "Statistics of a raster's valid cells.") | |
| 136 | 44 | .def_ro("min", &RasterStats::min) | |
| 137 | 44 | .def_ro("max", &RasterStats::max) | |
| 138 | 44 | .def_ro("mean", &RasterStats::mean) | |
| 139 | 44 | .def_ro("valid_count", &RasterStats::validCount) | |
| 140 | 44 | .def("__repr__", [](const RasterStats& s) { | |
| 141 | 2 | return "RasterStats(min=" + std::to_string(s.min) + ", max=" + std::to_string(s.max) + | |
| 142 | 4 | ", mean=" + std::to_string(s.mean) + | |
| 143 | 3 | ", valid_count=" + std::to_string(s.validCount) + ")"; | |
| 144 | }); | ||
| 145 | 44 | m.def("compute_stats", &ComputeStats<float>, "raster"_a, | |
| 146 | "Statistics over cells that are finite and not the nodata value."); | ||
| 147 | 44 | } | |
| 148 | |||
| 149 | 53 | RasterGrid CentredGrid(std::size_t width, std::size_t height, double pixelSize, | |
| 150 | const Vec2d& centre) { | ||
| 151 | 53 | const double halfWidth = 0.5 * static_cast<double>(width) * pixelSize; | |
| 152 | 53 | const double halfHeight = 0.5 * static_cast<double>(height) * pixelSize; | |
| 153 | return RasterGrid{ | ||
| 154 | width, height, | ||
| 155 | 53 | GeoTransform::NorthUp(centre.x - halfWidth, centre.y + halfHeight, pixelSize)}; | |
| 156 | } | ||
| 157 | |||
| 158 | 44 | void BindGrid(nb::module_& m) { | |
| 159 | ✗ | nb::class_<RasterGrid>(m, "RasterGrid", | |
| 160 | "A `width` x `height` grid of cells placed by a geotransform.") | ||
| 161 | 44 | .def( | |
| 162 | "__init__", | ||
| 163 | ✗ | [](RasterGrid* self, std::size_t width, std::size_t height, GeoTransform transform) { | |
| 164 | ✗ | new (self) RasterGrid{width, height, transform}; | |
| 165 | ✗ | }, | |
| 166 | 44 | "width"_a, "height"_a, "transform"_a) | |
| 167 | 44 | .def_static("centred", &CentredGrid, "width"_a, "height"_a, "pixel_size"_a, | |
| 168 | 88 | "centre"_a = Vec2d{}, "A north-up grid of square cells centred on `centre`.") | |
| 169 | 88 | .def_static("from_raster", &RasterGrid::FromRaster, "raster"_a, | |
| 170 | "The grid `raster` covers: its size and transform.") | ||
| 171 | 44 | .def_rw("width", &RasterGrid::width) | |
| 172 | 44 | .def_rw("height", &RasterGrid::height) | |
| 173 | 44 | .def_rw("transform", &RasterGrid::transform) | |
| 174 | 44 | .def_prop_ro("bounds", &RasterGrid::Bounds) | |
| 175 | 44 | .def_prop_ro("area", &RasterGrid::Area, "Area in square world units.") | |
| 176 | 44 | .def_prop_ro("pixel_size", &RasterGrid::PixelSize, "Cell size in world units.") | |
| 177 | 44 | .def("surface_area", &RasterGrid::SurfaceArea, "surface"_a, | |
| 178 | "The area the grid covers on `surface`, in square metres.") | ||
| 179 | 44 | .def("spacing", &RasterGrid::Spacing, "surface"_a, | |
| 180 | "The sample spacing in metres on `surface`.") | ||
| 181 | 132 | .def("make_raster", &RasterGrid::MakeRaster, "fill"_a = 0.0f, | |
| 182 | "A raster on this grid filled with `fill`."); | ||
| 183 | 44 | } | |
| 184 | |||
| 185 | 44 | void BindResampling(nb::module_& m) { | |
| 186 | 88 | nb::enum_<ResampleMethod>(m, "ResampleMethod", "How a raster is interpolated.") | |
| 187 | 44 | .value("NEAREST", ResampleMethod::Nearest) | |
| 188 | 44 | .value("BILINEAR", ResampleMethod::Bilinear) | |
| 189 | .value("CUBIC", ResampleMethod::Cubic); | ||
| 190 | 44 | m.def("gaussian_blur", nb::overload_cast<const RasterF&, double>(&GaussianBlur), "raster"_a, | |
| 191 | 44 | "sigma_pixels"_a, | |
| 192 | "The raster blurred by a Gaussian of `sigma_pixels` cells; NaN cells are left out of " | ||
| 193 | "the weights."); | ||
| 194 | 44 | m.def("gaussian_blur", nb::overload_cast<const RasterF&, double, double>(&GaussianBlur), | |
| 195 | 44 | "raster"_a, "sigma_columns"_a, "sigma_rows"_a); | |
| 196 | 88 | m.def("resample", &Resample, "raster"_a, "pixel_size"_a, "method"_a = ResampleMethod::Cubic, | |
| 197 | "The raster resampled to square cells of `pixel_size` over the same extent."); | ||
| 198 | 88 | m.def("resample_to", &ResampleTo, "raster"_a, "grid"_a, "method"_a = ResampleMethod::Cubic, | |
| 199 | "The raster resampled onto `grid`."); | ||
| 200 | 44 | m.def("sample_raster", &SampleRaster, "raster"_a, "x"_a, "y"_a, | |
| 201 | 88 | "method"_a = ResampleMethod::Cubic, | |
| 202 | "The raster interpolated at world position (x, y); NaN outside it."); | ||
| 203 | 44 | } | |
| 204 | |||
| 205 | 44 | void BindTiffEnums(nb::module_& m) { | |
| 206 | 88 | nb::enum_<TiffCompression>(m, "TiffCompression", "TIFF compression.") | |
| 207 | 44 | .value("NONE", TiffCompression::None) | |
| 208 | 44 | .value("LZW", TiffCompression::Lzw) | |
| 209 | 44 | .value("PACKBITS", TiffCompression::PackBits, "Read only.") | |
| 210 | .value("DEFLATE", TiffCompression::Deflate); | ||
| 211 | 88 | nb::enum_<TiffSampleFormat>(m, "TiffSampleFormat", "How a TIFF's samples are stored.") | |
| 212 | 44 | .value("UNSIGNED_INTEGER", TiffSampleFormat::UnsignedInteger) | |
| 213 | 44 | .value("SIGNED_INTEGER", TiffSampleFormat::SignedInteger) | |
| 214 | .value("FLOAT", TiffSampleFormat::Float); | ||
| 215 | 44 | } | |
| 216 | |||
| 217 | 44 | void BindGeoKeys(nb::module_& m) { | |
| 218 | 44 | nb::class_<GeoKeys>(m, "GeoKeys", "A GeoTIFF's projection keys (its CRS), as read or to write.") | |
| 219 | 44 | .def(nb::init<>(), "No keys: no CRS.") | |
| 220 | 44 | .def_prop_ro("empty", &GeoKeys::Empty) | |
| 221 | 44 | .def_prop_ro("is_geographic", &GeoKeys::IsGeographic, | |
| 222 | "Whether the CRS is geographic (degrees of longitude and latitude).") | ||
| 223 | 44 | .def_prop_ro("pixel_is_point", &GeoKeys::PixelIsPoint) | |
| 224 | 44 | .def(nb::self == nb::self); | |
| 225 | 44 | m.def("geographic_sphere_keys", &GeographicSphereKeys, "radius"_a, | |
| 226 | "GeoKeys for longitude and latitude on a sphere of `radius` metres, such as a moon or " | ||
| 227 | "planet."); | ||
| 228 | 88 | nb::class_<BandScaling>(m, "BandScaling", "A band's scale and offset: value * scale + offset.") | |
| 229 | 44 | .def_ro("scale", &BandScaling::scale) | |
| 230 | .def_ro("offset", &BandScaling::offset); | ||
| 231 | 44 | } | |
| 232 | |||
| 233 | 44 | void BindGeoTiffInfo(nb::module_& m) { | |
| 234 | 88 | nb::class_<GeoTiffInfo>(m, "GeoTiffInfo", "What a GeoTIFF's header says.") | |
| 235 | 44 | .def_ro("width", &GeoTiffInfo::width) | |
| 236 | 44 | .def_ro("height", &GeoTiffInfo::height) | |
| 237 | 44 | .def_ro("bands", &GeoTiffInfo::bands) | |
| 238 | 44 | .def_ro("bits_per_sample", &GeoTiffInfo::bitsPerSample) | |
| 239 | 44 | .def_ro("sample_format", &GeoTiffInfo::sampleFormat) | |
| 240 | 44 | .def_ro("compression", &GeoTiffInfo::compression) | |
| 241 | 44 | .def_ro("predictor", &GeoTiffInfo::predictor) | |
| 242 | 44 | .def_ro("tiled", &GeoTiffInfo::tiled) | |
| 243 | 44 | .def_ro("big_tiff", &GeoTiffInfo::bigTiff) | |
| 244 | 44 | .def_ro("big_endian", &GeoTiffInfo::bigEndian) | |
| 245 | 44 | .def_ro("transform", &GeoTiffInfo::transform) | |
| 246 | 44 | .def_ro("georeferenced", &GeoTiffInfo::georeferenced) | |
| 247 | 44 | .def_ro("nodata", &GeoTiffInfo::noData) | |
| 248 | 44 | .def_ro("band_scaling", &GeoTiffInfo::bandScaling) | |
| 249 | .def_ro("geo_keys", &GeoTiffInfo::geoKeys); | ||
| 250 | 44 | m.def("read_geotiff_info", &ReadGeoTiffInfo, "path"_a, | |
| 251 | "Reads a GeoTIFF's header: size, bands, sample format, compression, georeferencing and " | ||
| 252 | "CRS keys."); | ||
| 253 | 44 | } | |
| 254 | |||
| 255 | 44 | void BindGeoTiffIo(nb::module_& m) { | |
| 256 | auto options = nb::class_<GeoTiffWriteOptions>( | ||
| 257 | m, "GeoTiffWriteOptions", | ||
| 258 | "How GeoTIFFs are written: `compression` (none, LZW or Deflate), square tiles of " | ||
| 259 | 44 | "`tile_size` (a positive multiple of 16) and `big_tiff`."); | |
| 260 | 44 | BindKwargsInit<GeoTiffWriteOptions>(options); | |
| 261 | 44 | options.def_rw("compression", &GeoTiffWriteOptions::compression) | |
| 262 | 44 | .def_rw("tile_size", &GeoTiffWriteOptions::tileSize) | |
| 263 | .def_rw("big_tiff", &GeoTiffWriteOptions::bigTiff); | ||
| 264 | using WriteBands = void (*)(const std::filesystem::path&, const std::vector<const RasterF*>&, | ||
| 265 | const GeoKeys&, const GeoTiffWriteOptions&); | ||
| 266 | using WriteOne = void (*)(const std::filesystem::path&, const RasterF&, const GeoKeys&, | ||
| 267 | const GeoTiffWriteOptions&); | ||
| 268 | 44 | m.def("write_geotiff", static_cast<WriteBands>(&WriteGeoTiff), "path"_a, "bands"_a, | |
| 269 | 132 | "geo_keys"_a = GeoKeys{}, "options"_a = GeoTiffWriteOptions{}, | |
| 270 | "Writes one or more rasters as the bands of a float32 GeoTIFF, with `geo_keys` as its " | ||
| 271 | "CRS."); | ||
| 272 | 44 | m.def("write_geotiff", static_cast<WriteOne>(&WriteGeoTiff), "path"_a, "raster"_a, | |
| 273 | 132 | "geo_keys"_a = GeoKeys{}, "options"_a = GeoTiffWriteOptions{}); | |
| 274 | 44 | m.def( | |
| 275 | "read_geotiff", | ||
| 276 | ✗ | [](const std::filesystem::path& path, std::size_t band) { | |
| 277 | 3 | return ReadGeoTiff(path, band).raster; | |
| 278 | }, | ||
| 279 | 132 | "path"_a, "band"_a = 0, | |
| 280 | "Reads band `band` (from 0) of a GeoTIFF as a Raster with the file's transform and no-data " | ||
| 281 | "value; `read_geotiff_info` gives its CRS keys."); | ||
| 282 | 44 | } | |
| 283 | |||
| 284 | } // namespace | ||
| 285 | |||
| 286 | 44 | void BindRaster(nb::module_& m) { | |
| 287 | 44 | BindGeoTransform(m); | |
| 288 | 44 | BindRasterClass(m); | |
| 289 | 44 | BindStats(m); | |
| 290 | 44 | BindGrid(m); | |
| 291 | 44 | BindResampling(m); | |
| 292 | 44 | nb::exception<GeoTiffError>(m, "GeoTiffError", PyExc_RuntimeError); | |
| 293 | 44 | BindTiffEnums(m); | |
| 294 | 44 | BindGeoKeys(m); | |
| 295 | 44 | BindGeoTiffInfo(m); | |
| 296 | 44 | BindGeoTiffIo(m); | |
| 297 | 44 | } | |
| 298 | |||
| 299 | } // namespace terraingen::python | ||
| 300 |