Note
Click here 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
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 )