Note

Go to the end to download the full example code.

# Insert Raster

The RasterElement objects store the Raster data in WKB form. This WKB format is usually fetched from the database but when the data comes from another source it can be hard to format it as as a WKB. This example shows a method to convert input data into a WKB in order to insert it. This example uses SQLAlchemy ORM queries.

Warning

The PixelType values are not always properly translated by the Rasterio library, so exporting a raster and re-importing it using this method will properly import the values but might not keep the same internal types.

     16 import struct
     17 from sys import byteorder
     18
     19 import numpy as np
     20 import pytest
     21 import rasterio
     22 from sqlalchemy import Column
     23 from sqlalchemy import Integer
     24 from sqlalchemy import MetaData
     25 from sqlalchemy import text
     26 from sqlalchemy.orm import declarative_base
     27
     28 from geoalchemy2 import Raster
     29 from geoalchemy2 import RasterElement
     30
     31 # Tests imports
     32 from tests import test_only_with_dialects
     33
     34 metadata = MetaData()
     35 Base = declarative_base(metadata=metadata)
     36
     37
     38 class Ocean(Base):  # type: ignore
     39     __tablename__ = "ocean"
     40     id = Column(Integer, primary_key=True)
     41     rast = Column(Raster)
     42
     43     def __init__(self, rast):
     44         self.rast = rast
     45
     46
     47 _DTYPE = {
     48     "?": [0, "?", 1],
     49     "u1": [2, "B", 1],
     50     "i1": [3, "b", 1],
     51     "B": [4, "B", 1],
     52     "i2": [5, "h", 2],
     53     "u2": [6, "H", 2],
     54     "i4": [7, "i", 4],
     55     "u4": [8, "I", 4],
     56     "f4": [10, "f", 4],
     57     "f8": [11, "d", 8],
     58 }
     59
     60
     61 def write_wkb_raster(dataset):
     62     """Creates a WKB raster from the given raster file with rasterio.
     63     :dataset: Rasterio dataset
     64     :returns: binary: Binary raster in WKB format
     65
     66     This function was imported from
     67     https://github.com/nathancahill/wkb-raster/blob/master/wkb_raster.py
     68     and slightly adapted.
     69     """
     70     # Define format, see https://docs.python.org/3/library/struct.html
     71     format_string = "bHHddddddIHH"
     72
     73     if byteorder == "big":
     74         endian = ">"
     75         endian_byte = 0
     76     elif byteorder == "little":
     77         endian = "<"
     78         endian_byte = 1
     79
     80     # Write the raster header data.
     81     header = b""
     82
     83     transform = dataset.transform.to_gdal()
     84
     85     version = 0
     86     nBands = int(dataset.count)
     87     scaleX = transform[1]
     88     scaleY = transform[5]
     89     ipX = transform[0]
     90     ipY = transform[3]
     91     skewX = 0
     92     skewY = 0
     93     srid = int(dataset.crs.to_string().split("EPSG:")[1])
     94     width = int(dataset.meta.get("width"))
     95     height = int(dataset.meta.get("height"))
     96
     97     if width > 65535 or height > 65535:
     98         raise ValueError("PostGIS does not support rasters with width or height greater than 65535")
     99
    100     fmt = f"{endian}{format_string}"
    101
    102     header = struct.pack(
    103         fmt,
    104         endian_byte,
    105         version,
    106         nBands,
    107         scaleX,
    108         scaleY,
    109         ipX,
    110         ipY,
    111         skewX,
    112         skewY,
    113         srid,
    114         width,
    115         height,
    116     )
    117
    118     bands = []
    119
    120     # Create band header data
    121
    122     # not used - always False
    123     isOffline = False
    124     hasNodataValue = False
    125
    126     if "nodata" in dataset.meta:
    127         hasNodataValue = True
    128
    129     # not used - always False
    130     isNodataValue = False
    131
    132     # unset
    133     reserved = False
    134
    135     # # Based on the pixel type, determine the struct format, byte size and
    136     # # numpy dtype
    137     rasterio_dtype = dataset.meta.get("dtype")
    138     dt_short = np.dtype(rasterio_dtype).str[1:]
    139     pixtype, nodata_fmt, _ = _DTYPE[dt_short]
    140
    141     # format binary -> :b
    142     binary_str = f"{isOffline:b}{hasNodataValue:b}{isNodataValue:b}{reserved:b}{pixtype:b}"
    143     # convert to int
    144     binary_decimal = int(binary_str, 2)
    145
    146     # pack to 1 byte
    147     # 4 bits for ifOffline, hasNodataValue, isNodataValue, reserved
    148     # 4 bit for pixtype
    149     # -> 8 bit = 1 byte
    150     band_header = struct.pack("<b", binary_decimal)
    151
    152     # Write the nodata value
    153     nodata = struct.pack(nodata_fmt, int(dataset.meta.get("nodata") or 0))
    154
    155     for i in range(1, nBands + 1):
    156         band_array = dataset.read(i)
    157
    158         # # Write the pixel values: width * height * size
    159
    160         # numpy tobytes() method instead of packing with struct.pack()
    161         band_binary = band_array.reshape(width * height).tobytes()
    162
    163         bands.append(band_header + nodata + band_binary)
    164
    165     # join all bands
    166     allbands = b""
    167     for b in bands:
    168         allbands += b
    169
    170     wkb = header + allbands
    171
    172     return wkb
    173
    174
    175 @test_only_with_dialects("postgresql")
    176 class TestInsertRaster:
    177     @pytest.fixture(
    178         params=[
    179             "1BB",
    180             "2BUI",
    181             "4BUI",
    182             "8BSI",
    183             "8BUI",
    184             "16BSI",
    185             "16BUI",
    186             "32BSI",
    187             "32BUI",
    188             "32BF",
    189             "64BF",
    190         ]
    191     )
    192     def input_img(self, conn, tmpdir, request):
    193         """Create a TIFF image that will be imported as Raster."""
    194         pixel_type = request.param
    195         conn.execute(text("SET postgis.gdal_enabled_drivers = 'ENABLE_ALL';"))
    196         data = conn.execute(
    197             text(
    198                 f"""SELECT
    199                     ST_AsTIFF(
    200                         ST_AsRaster(
    201                             ST_GeomFromText('POLYGON((0 0,1 1,0 1,0 0))'),
    202                             5,
    203                             6,
    204                             '{pixel_type}'
    205                         ),
    206                         'GTiff'
    207                     );"""
    208             )
    209         ).scalar()
    210         filename = tmpdir / "image.tiff"
    211         with open(filename, "wb") as f:
    212             f.write(data.tobytes())
    213         return filename, pixel_type
    214
    215     def test_insert_raster(self, session, conn, input_img):
    216         """Insert a TIFF image into a raster column."""
    217         filename, pixel_type = input_img
    218         metadata.drop_all(conn, checkfirst=True)
    219         metadata.create_all(conn)
    220
    221         # Load the image and transform it into a WKB
    222         with rasterio.open(str(filename), "r+") as dataset:
    223             dataset.crs = rasterio.crs.CRS.from_epsg(4326)
    224             expected_values = dataset.read()[0]
    225             raw_wkb = write_wkb_raster(dataset)
    226
    227         # Insert a new raster element
    228         polygon_raster = RasterElement(raw_wkb)
    229         o = Ocean(polygon_raster)
    230         session.add(o)
    231         session.flush()
    232
    233         # Check inserted values
    234         new_values = conn.execute(text("SELECT ST_DumpValues(rast, 1, false) FROM ocean;")).scalar()
    235         np.testing.assert_array_equal(
    236             np.array(new_values, dtype=expected_values.dtype), expected_values
    237         )

Gallery generated by Sphinx-Gallery