raster.h

Go to the documentation of this file.

/******************************************************************************
 *
 * Project: sgs
 * Purpose: GDALDataset wrapper for raster operations
 * Author: Joseph Meyer
 * Date: June, 2025
 *
 ******************************************************************************/

 
#pragma once
 
#include <filesystem>
#include <iostream>
 
#include <gdal_priv.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
 
#include <utils/helper.h>
 
//used as cutoff for max band allowed in memory
#define GIGABYTE 1073741824
 
namespace sgs {
namespace raster {
 
namespace py = pybind11;
using namespace pybind11::literals;

class GDALRasterWrapper {
    private:
    GDALDatasetUniquePtr p_dataset;
 
    std::vector<void *> rasterBandPointers;
    std::vector<bool> rasterBandRead;
 
    std::vector<void *> displayRasterBandPointers;
    std::vector<bool> displayRasterBandRead;
    int displayRasterWidth = -1;
    int displayRasterHeight = -1;
 
    double geotransform[6];
    std::string crs = "";
    std::string proj = "";
 
    std::string tempDir = "";
 
    bool destroyed = false;
    bool externalRasterData = false;

    void readRasterBand(int width, int height, int band) {  
        GDALDataType type = this->getRasterBandType(band);
        size_t size = this->getRasterBandTypeSize(band);
        size_t max = std::numeric_limits<size_t>::max();
        
        //perform size checks
        if (max / size < static_cast<size_t>(width) ||
            max / (size * static_cast<size_t>(width)) < static_cast<size_t>(height)) {
                throw std::runtime_error("raster too large to fit in memory.");
        }
 
        if (static_cast<size_t>(height) * static_cast<size_t>(width) * size > GIGABYTE) {
                throw std::runtime_error("sgs does not allow allocation of a raster into memory for direct pixel access purposes if it would be larger than 1 gigabyte.");
        }
 
        //allocate data
        void *p_data = VSIMalloc3(height, width, size);
 
        //perform raster read on current band
        CPLErr err = this->p_dataset->GetRasterBand(band + 1)->RasterIO(
            GF_Read,            //GDALRWFlag eRWFlag
            0,              //int nXOff
            0,              //int nYOff
            this->getWidth(),       //int nXSize
            this->getHeight(),     //int nYSize
            p_data,             //void *pData
            width,              //int nBufXSize
            height,             //int nBufYSize
            type,               //GDALDataType eBufType 
            0,              //int nPixelSpace
            0               //int nLineSpace
        );
        if (err) {
            throw std::runtime_error("error reading raster band from dataset.");
        }
 
        //update dislpay information as required
        if (width != this->getWidth() || height != this->getHeight()) {
            this->displayRasterBandRead[band] = true;
            this->displayRasterBandPointers[band] = p_data;
        }
        else {
            this->rasterBandRead[band] = true;
            this->rasterBandPointers[band] = p_data;
        }
    }

    template <typename T> 
    py::buffer getBuffer(size_t size, void *p_buffer, int width, int height) {
        //see https://pybind11.readthedocs.io/en/stable/advanced/pycpp/numpy.html#memory-view
        return py::memoryview::from_buffer(
            (T*)p_buffer,       //buffer
            {height, width},    //shape
            {size * width, size}    //stride
        );
    }   

    void createFromDataset(GDALDataset *p_dataset) {
        this->p_dataset = GDALDatasetUniquePtr(p_dataset);
 
        //geotransform
        CPLErr cplerr = this->p_dataset->GetGeoTransform(this->geotransform);
        if (cplerr) {
            throw std::runtime_error("error getting geotransform from dataset.");
        }
 
        //crs
        if(std::string(this->p_dataset->GetProjectionRef()).length() == 0) {
            throw std::runtime_error("raster dataset does not have a projection definition.");
        }
        this->crs = std::string(OGRSpatialReference(this->p_dataset->GetProjectionRef()).GetName());
 
        //proj
        char *p_proj;
        OGRErr ogrerr = OGRSpatialReference(this->p_dataset->GetProjectionRef()).exportToPrettyWkt(&p_proj);
        if (ogrerr) {
            throw std::runtime_error("error getting projection as WKT from dataset.");
        }
        this->proj = std::string(p_proj);
        CPLFree(p_proj);
 
        //initialize (but don't read) raster band pointers
        this->rasterBandPointers = std::vector<void *>(this->getBandCount(), nullptr);
        this->rasterBandRead = std::vector<bool>(this->getBandCount(), false);
        this->displayRasterBandPointers = std::vector<void *>(this->getBandCount(), nullptr);
        this->displayRasterBandRead = std::vector<bool>(this->getBandCount(), false);   
    }
 
    public:
    GDALRasterWrapper(std::string filename, std::string projDBPath) {
        //set proj.db search path to search for the proj.db file which is included in sgs package
        char **paths = nullptr;
        paths = CSLAddString(paths, projDBPath.c_str());
        OSRSetPROJSearchPaths(paths);
        CSLDestroy(paths);
        
        //must register drivers before trying to open a dataset
        GDALAllRegister();
    
        //dataset
        GDALDataset *p_dataset = GDALDataset::FromHandle(GDALOpen(filename.c_str(), GA_ReadOnly));
        if (!p_dataset) {
            throw std::runtime_error("file given does not result in a valid dataset, check to ensure file path is accurate.");
        }
 
        this->createFromDataset(p_dataset);
    }   

    GDALRasterWrapper(GDALDataset *p_dataset, std::vector<void *> bands) {
        this->createFromDataset(p_dataset);
        this->rasterBandPointers = bands;
        this->rasterBandRead = std::vector<bool>(bands.size(), true);
    }

    GDALRasterWrapper(GDALDataset *p_dataset) {
        this->createFromDataset(p_dataset);
    }

    GDALRasterWrapper(py::buffer buffer, std::vector<double> geotransform, std::string projection, std::vector<double> nanVals, std::vector<std::string> names, std::string projDBPath) {
        //set proj.db search path to search for the proj.db file which is included in sgs package
        char **paths = nullptr;
        paths = CSLAddString(paths, projDBPath.c_str());
        OSRSetPROJSearchPaths(paths);
        CSLDestroy(paths);
        
        py::buffer_info info = buffer.request();
 
        //get height width and band count from pybuffer
        int width, height;
        size_t bandCount, bandSize;
        if (info.ndim == 3) {
            bandCount = info.shape[0];
            height = info.shape[1];
            width = info.shape[2];
            bandSize = info.strides[0];
        }
        else if (info.ndim == 2) {
            bandCount = 1;
            height = info.shape[0];
            width = info.shape[1];
        }
        else {
            throw std::runtime_error("dimension of numpy array must be 2 or 3");
        }       
 
        if (bandCount != names.size()) {
            throw std::runtime_error("band names array does not have the same number of bands as the py buffer");
        }
 
        //get data type from pybuffer
        GDALDataType type;
        size_t size;
        if (info.format == py::format_descriptor<int8_t>::format()) {
            type = GDT_Int8;
            size = sizeof(int8_t);
        }
        else if (info.format == py::format_descriptor<int16_t>::format()) {
            type = GDT_Int16;
            size = sizeof(int16_t);
        }
        else if (info.format == py::format_descriptor<uint16_t>::format()) {
            type = GDT_UInt16;
            size = sizeof(uint16_t);
        }
        else if (info.format == py::format_descriptor<int32_t>::format()) {
            type = GDT_Int32;
            size = sizeof(int32_t);
        }   
        else if (info.format == py::format_descriptor<uint32_t>::format()) {
            type = GDT_UInt32;
            size = sizeof(uint32_t);
        }
        else if (info.format == py::format_descriptor<float>::format()) {
            type = GDT_Float32;
            size = sizeof(float);
        }
        else if (info.format == py::format_descriptor<double>::format()) {
            type = GDT_Float64;
            size = sizeof(double);
        }
        else {
            throw std::runtime_error("data type of array must be one of int8, int16, uint16, int32, uint32, float32, or float64.");
        }
 
        GDALAllRegister();
        GDALDataset *p_dataset = helper::createVirtualDataset("MEM", width, height, geotransform.data(), projection);
        std::vector<void *> bands(bandCount);   
        for (size_t i = 0; i < bandCount; i++) {
            helper::RasterBandMetaData band;
            band.p_buffer = (void *)((size_t)info.ptr + (i * bandSize));
            band.type = type;
                band.size = size;
            band.nan = nanVals[i];
            band.name = names[i];
            helper::addBandToMEMDataset(p_dataset, band);    
            bands[i] = band.p_buffer;
 
        }
 
        this->createFromDataset(p_dataset);
        this->rasterBandPointers = bands;
        this->rasterBandRead = std::vector<bool>(bandCount, true);
        this->externalRasterData = true;
    }

    ~GDALRasterWrapper() {
        if (destroyed) {
            return;
        }
 
        for (int i = 0; i < this->getBandCount(); i++) {
            //if the raster data is coming from a numpy array (this->externalRasterData true), then
            //the memory will be cleaned up by Pythons garbage collector
            if (this->rasterBandRead[i] && !this->externalRasterData) {
                CPLFree(this->rasterBandPointers[i]);
            }
 
            if (this->displayRasterBandRead[i]) {
                CPLFree(this->displayRasterBandPointers[i]);
            }
        }
 
        GDALClose(GDALDataset::ToHandle(this->p_dataset.release()));
 
        if (this->tempDir != "") {
            std::filesystem::path temp = this->tempDir;
            std::filesystem::remove_all(temp);
        }   
    }

    void close(void) {
        for (int i = 0; i < this->getBandCount(); i++) {
            //if the raster data is coming from a numpy array (this->externalRasterData true), then
            //the memory will be cleaned up by Pythons garbage collector
            if (this->rasterBandRead[i]) {
                CPLFree(this->rasterBandPointers[i]);
            }
 
            if (this->displayRasterBandRead[i]) {
                CPLFree(this->displayRasterBandPointers[i]);
            }
        }
 
        GDALClose(GDALDataset::ToHandle(this->p_dataset.release()));
 
        if (this->tempDir != "") {
            std::filesystem::path temp = this->tempDir;
            std::filesystem::remove_all(temp);
        }
 
        destroyed = true;   
    }

    GDALDataset *getDataset() {
        return this->p_dataset.get();
    }

    std::string getDriver() {
        return std::string(this->p_dataset->GetDriverName()) 
            + "/" 
            + std::string(this->p_dataset->GetDriver()->GetMetadataItem(GDAL_DMD_LONGNAME));
    }

    std::string getFullProjectionInfo() {
        return this->proj;
    }

    std::string getCRS(){
        return this->crs;
    }

    int getWidth() {
        return this->p_dataset->GetRasterXSize();
    }

    int getHeight(){
        return this->p_dataset->GetRasterYSize();
    }

    int getBandCount() {
        return this->p_dataset->GetRasterCount();
    }

    double getXMax() {
        int width = this->p_dataset->GetRasterXSize();
        int height = this->p_dataset->GetRasterYSize();
        return std::max(
            this->geotransform[0], 
            this->geotransform[0] + this->geotransform[1] * width + this->geotransform[2] * height
        );
    }

    double getXMin() {
        int width = this->p_dataset->GetRasterXSize();
        int height = this->p_dataset->GetRasterYSize();
        return std::min(
            this->geotransform[0], 
            this->geotransform[0] + this->geotransform[1] * width + this->geotransform[2] * height
        );
    }

    double getYMax() {
        int width = this->p_dataset->GetRasterXSize();
        int height = this->p_dataset->GetRasterYSize();
        return std::max(
            this->geotransform[3],
            this->geotransform[3] + this->geotransform[4] * width + this->geotransform[5] * height
        );
    }

    double getYMin() {
        int width = this->p_dataset->GetRasterXSize();
        int height = this->p_dataset->GetRasterYSize();
        return std::min(
            this->geotransform[3],
            this->geotransform[3] + this->geotransform[4] * width + this->geotransform[5] * height
        );
    }

    double getPixelWidth() {
        return std::abs(this->geotransform[5]);
    }

    double getPixelHeight() {
        return std::abs(this->geotransform[1]);
    }

    std::vector<std::string> getBands(){
        std::vector<std::string> retval;
 
        for( auto&& p_band : this->p_dataset->GetBands() ){
            retval.push_back(p_band->GetDescription());
        }
 
        return retval;
    }   

    double *getGeotransform() {
        return this->geotransform;
    }

    double getBandNoDataValue(int band){
        GDALRasterBand *p_band = this->getRasterBand(band);
        return p_band->GetNoDataValue();
    }

    py::buffer getRasterBandAsMemView(int width, int height, int band) {
        bool display = (width != this->getWidth() || height != this->getHeight());
        void *p_buffer;
        GDALDataType type = this->getRasterBandType(band);
 
        //allocate raster if required
        if (!display && !this->rasterBandRead[band]) {
            this->readRasterBand(width, height, band);
        }
 
        //(re)allocate display raster if required
        if (display) {
            if (width != this->displayRasterWidth || height != this->displayRasterHeight) {
                free(this->displayRasterBandPointers[band]);
                this->displayRasterBandRead[band] = false;
            }
 
            if (this->displayRasterBandRead[band] == false) {
                this->readRasterBand(width, height, band);
            }
        }
 
        //get the (allocated) data buffer
        p_buffer = (!display) ?
            this->rasterBandPointers[band] :
            this->displayRasterBandPointers[band];
 
        switch(type) {
            case GDT_Int8:
                return getBuffer<int8_t>(sizeof(int8_t), p_buffer, width, height);
            case GDT_UInt16:
                return getBuffer<uint16_t>(sizeof(uint16_t), p_buffer, width, height);
            case GDT_Int16:
                return getBuffer<int16_t>(sizeof(int16_t), p_buffer, width, height);
            case GDT_UInt32:
                return getBuffer<uint32_t>(sizeof(uint32_t), p_buffer, width, height);
            case GDT_Int32:
                return getBuffer<int32_t>(sizeof(int32_t), p_buffer, width, height);
            case GDT_Float32:
                return getBuffer<float>(sizeof(float), p_buffer, width, height);
            case GDT_Float64:
                return getBuffer<double>(sizeof(double), p_buffer, width, height);
            default:
                throw std::runtime_error("raster pixel data type not supported.");
        }
    }

    void releaseBandBuffers(void) {
        for (size_t i = 0; i < this->rasterBandPointers.size(); i++) {
            rasterBandPointers[i] = nullptr;
            rasterBandRead[i] = false;
        }   
    }

    GDALRasterBand *getRasterBand(int band) {
        return this->p_dataset->GetRasterBand(band + 1);    
    }

    void *getRasterBandBuffer(int band) {
        if (!this->rasterBandRead[band]) {
            this->readRasterBand(
                this->getWidth(), 
                this->getHeight(),
                band
            );
        }
 
        return this->rasterBandPointers[band];
    }

    GDALDataType getRasterBandType(int band) {
        GDALRasterBand *p_band = this->p_dataset->GetRasterBand(band + 1);
        return p_band->GetRasterDataType(); 
    }

    size_t getRasterBandTypeSize(int band) {
        switch (this->getRasterBandType(band)) {
            case GDALDataType::GDT_Int8:
                return 1;
            case GDALDataType::GDT_UInt16:
            case GDALDataType::GDT_Int16:
                return 2;
            case GDALDataType::GDT_UInt32:
            case GDALDataType::GDT_Int32:
            case GDALDataType::GDT_Float32:
                return 4;
            case GDALDataType::GDT_Float64:
                return 8;
            default:
                std::string errorMsg = "GDALDataType of band " + std::to_string(band) + " not supported.";
                throw std::runtime_error(errorMsg);
        }
    }

    void write(std::string filename) {
        std::filesystem::path filepath = filename;
        std::string extension = filepath.extension().string();
        
        if (extension != ".tif") {
            throw std::runtime_error("write only supports .tif files right now");
        }
    
        GDALDriver *p_driver = GetGDALDriverManager()->GetDriverByName("GTiff");
    
        GDALClose(p_driver->CreateCopy(filename.c_str(), this->p_dataset.get(), (int)false, nullptr, nullptr, nullptr));
    }

    void setTempDir(std::string tempDir) {
        this->tempDir = tempDir;
    }

    std::string getTempDir() {
        return this->tempDir;
    }

    std::vector<double> getGeotransformArray() {
        std::vector<double> retval(6);
        for (int i = 0; i < 6; i++) {
            retval[i] = this->geotransform[i];
        }
 
        return retval;
    }

    std::string getDataType() {
        GDALDataType type = this->getRasterBandType(0);
        for (int i = 1; i < this->getBandCount(); i++) {
            if (this->getRasterBandType(i) != type)    {
                return "";
            }
        }
 
        switch (type) {
            case GDT_Int8: return "int8";
            case GDT_UInt16: return "uint16";
            case GDT_Int16: return "int16";
            case GDT_UInt32: return "uint32";
            case GDT_Int32: return "int32";
            case GDT_Float32: return "float32";
            case GDT_Float64: return "float64";
            default: throw std::runtime_error("GDAL data type not supported");
        }
    }
};
 
} //namespace raster
} //namespace sgs