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