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
 71     # Define format, see https://docs.python.org/3/library/struct.html
 72     format_string = "bHHddddddIHH"
 73
 74     if byteorder == "big":
 75         endian = ">"
 76         endian_byte = 0
 77     elif byteorder == "little":
 78         endian = "<"
 79         endian_byte = 1
 80
 81     # Write the raster header data.
 82     header = bytes()
 83
 84     transform = dataset.transform.to_gdal()
 85
 86     version = 0
 87     nBands = int(dataset.count)
 88     scaleX = transform[1]
 89     scaleY = transform[5]
 90     ipX = transform[0]
 91     ipY = transform[3]
 92     skewX = 0
 93     skewY = 0
 94     srid = int(dataset.crs.to_string().split("EPSG:")[1])
 95     width = int(dataset.meta.get("width"))
 96     height = int(dataset.meta.get("height"))
 97
 98     if width > 65535 or height > 65535:
 99         raise ValueError("PostGIS does not support rasters with width or height greater than 65535")
100
101     fmt = f"{endian}{format_string}"
102
103     header = struct.pack(
104         fmt,
105         endian_byte,
106         version,
107         nBands,
108         scaleX,
109         scaleY,
110         ipX,
111         ipY,
112         skewX,
113         skewY,
114         srid,
115         width,
116         height,
117     )
118
119     bands = []
120
121     # Create band header data
122
123     # not used - always False
124     isOffline = False
125     hasNodataValue = False
126
127     if "nodata" in dataset.meta:
128         hasNodataValue = True
129
130     # not used - always False
131     isNodataValue = False
132
133     # unset
134     reserved = False
135
136     # # Based on the pixel type, determine the struct format, byte size and
137     # # numpy dtype
138     rasterio_dtype = dataset.meta.get("dtype")
139     dt_short = np.dtype(rasterio_dtype).str[1:]
140     pixtype, nodata_fmt, _ = _DTYPE[dt_short]
141
142     # format binary -> :b
143     binary_str = f"{isOffline:b}{hasNodataValue:b}{isNodataValue:b}{reserved:b}{pixtype:b}"
144     # convert to int
145     binary_decimal = int(binary_str, 2)
146
147     # pack to 1 byte
148     # 4 bits for ifOffline, hasNodataValue, isNodataValue, reserved
149     # 4 bit for pixtype
150     # -> 8 bit = 1 byte
151     band_header = struct.pack("<b", binary_decimal)
152
153     # Write the nodata value
154     nodata = struct.pack(nodata_fmt, int(dataset.meta.get("nodata") or 0))
155
156     for i in range(1, nBands + 1):
157         band_array = dataset.read(i)
158
159         # # Write the pixel values: width * height * size
160
161         # numpy tobytes() method instead of packing with struct.pack()
162         band_binary = band_array.reshape(width * height).tobytes()
163
164         bands.append(band_header + nodata + band_binary)
165
166     # join all bands
167     allbands = bytes()
168     for b in bands:
169         allbands += b
170
171     wkb = header + allbands
172
173     return wkb
174
175
176 @test_only_with_dialects("postgresql")
177 class TestInsertRaster:
178     @pytest.fixture(
179         params=[
180             "1BB",
181             "2BUI",
182             "4BUI",
183             "8BSI",
184             "8BUI",
185             "16BSI",
186             "16BUI",
187             "32BSI",
188             "32BUI",
189             "32BF",
190             "64BF",
191         ]
192     )
193     def input_img(self, conn, tmpdir, request):
194         """Create a TIFF image that will be imported as Raster."""
195         pixel_type = request.param
196         conn.execute(text("SET postgis.gdal_enabled_drivers = 'ENABLE_ALL';"))
197         data = conn.execute(
198             text(
199                 """SELECT
200                     ST_AsTIFF(
201                         ST_AsRaster(
202                             ST_GeomFromText('POLYGON((0 0,1 1,0 1,0 0))'),
203                             5,
204                             6,
205                             '{}'
206                         ),
207                         'GTiff'
208                     );""".format(
209                     pixel_type
210                 )
211             )
212         ).scalar()
213         filename = tmpdir / "image.tiff"
214         with open(filename, "wb") as f:
215             f.write(data.tobytes())
216         return filename, pixel_type
217
218     def test_insert_raster(self, session, conn, input_img):
219         """Insert a TIFF image into a raster column."""
220         filename, pixel_type = input_img
221         metadata.drop_all(conn, checkfirst=True)
222         metadata.create_all(conn)
223
224         # Load the image and transform it into a WKB
225         with rasterio.open(str(filename), "r+") as dataset:
226             dataset.crs = rasterio.crs.CRS.from_epsg(4326)
227             expected_values = dataset.read()[0]
228             raw_wkb = write_wkb_raster(dataset)
229
230         # Insert a new raster element
231         polygon_raster = RasterElement(raw_wkb)
232         o = Ocean(polygon_raster)
233         session.add(o)
234         session.flush()
235
236         # Check inserted values
237         new_values = conn.execute(text("SELECT ST_DumpValues(rast, 1, false) FROM ocean;")).scalar()
238         np.testing.assert_array_equal(
239             np.array(new_values, dtype=expected_values.dtype), expected_values
240         )

Gallery generated by Sphinx-Gallery