from __future__ import with_statement
from __future__ import division
import os
import sys
import math
import traceback
import itertools
import mercantile
import rasterio
import numpy as np
import sqlite3
from multiprocessing import Pool
from rasterio._io import virtual_file_to_buffer
from riomucho.single_process_pool import MockTub
from io import BytesIO
from PIL import Image
from rasterio import transform
from rasterio.warp import reproject, transform_bounds
from rasterio.enums import Resampling
from rio_rgbify.encoders import data_to_rgb
buffer = bytes if sys.version_info > (3,) else buffer
work_func = None
global_args = None
src = None
def _main_worker(inpath, g_work_func, g_args):
"""
Util for setting global vars w/ a Pool
"""
global work_func
global global_args
global src
work_func = g_work_func
global_args = g_args
src = rasterio.open(inpath)
def _encode_as_webp(data, profile=None, affine=None):
"""
Uses BytesIO + PIL to encode a (3, 512, 512)
array into a webp bytearray.
Parameters
-----------
data: ndarray
(3 x 512 x 512) uint8 RGB array
profile: None
ignored
affine: None
ignored
Returns
--------
contents: bytearray
webp-encoded bytearray of the provided input data
"""
with BytesIO() as f:
im = Image.fromarray(np.rollaxis(data, 0, 3))
im.save(f, format="webp", lossless=True)
return f.getvalue()
def _encode_as_png(data, profile, dst_transform):
"""
Uses rasterio's virtual file system to encode a (3, 512, 512)
array as a png-encoded bytearray.
Parameters
-----------
data: ndarray
(3 x 512 x 512) uint8 RGB array
profile: dictionary
dictionary of kwargs for png writing
affine: Affine
affine transform for output tile
Returns
--------
contents: bytearray
png-encoded bytearray of the provided input data
"""
profile["affine"] = dst_transform
with rasterio.open("/vsimem/tileimg", "w", **profile) as dst:
dst.write(data)
contents = bytearray(virtual_file_to_buffer("/vsimem/tileimg"))
return contents
def _tile_worker(tile):
"""
For each tile, and given an open rasterio src, plus a`global_args` dictionary
with attributes of `base_val`, `interval`, and a `writer_func`,
warp a continous single band raster to a 512 x 512 mercator tile,
then encode this tile into RGB.
Parameters
-----------
tile: list
[x, y, z] indices of tile
Returns
--------
tile, buffer
tuple with the input tile, and a bytearray with the data encoded into
the format created in the `writer_func`
"""
x, y, z = tile
bounds = [
c
for i in (
mercantile.xy(*mercantile.ul(x, y + 1, z)),
mercantile.xy(*mercantile.ul(x + 1, y, z)),
)
for c in i
]
toaffine = transform.from_bounds(*bounds + [512, 512])
out = np.empty((512, 512), dtype=src.meta["dtype"])
reproject(
rasterio.band(src, 1),
out,
dst_transform=toaffine,
dst_crs="epsg:3857",
resampling=Resampling.bilinear,
)
out = data_to_rgb(out, global_args["base_val"], global_args["interval"])
return tile, global_args["writer_func"](out, global_args["kwargs"].copy(), toaffine)
def _tile_range(min_tile, max_tile):
"""
Given a min and max tile, return an iterator of
all combinations of this tile range
Parameters
-----------
min_tile: list
[x, y, z] of minimun tile
max_tile:
[x, y, z] of minimun tile
Returns
--------
tiles: iterator
iterator of [x, y, z] tiles
"""
min_x, min_y, _ = min_tile
max_x, max_y, _ = max_tile
return itertools.product(range(min_x, max_x + 1), range(min_y, max_y + 1))
def _make_tiles(bbox, src_crs, minz, maxz):
"""
Given a bounding box, zoom range, and source crs,
find all tiles that would intersect
Parameters
-----------
bbox: list
[w, s, e, n] bounds
src_crs: str
the source crs of the input bbox
minz: int
minumum zoom to find tiles for
maxz: int
maximum zoom to find tiles for
Returns
--------
tiles: generator
generator of [x, y, z] tiles that intersect
the provided bounding box
"""
w, s, e, n = transform_bounds(*[src_crs, "epsg:4326"] + bbox, densify_pts=0)
EPSILON = 1.0e-10
w += EPSILON
s += EPSILON
e -= EPSILON
n -= EPSILON
for z in range(minz, maxz + 1):
for x, y in _tile_range(mercantile.tile(w, n, z), mercantile.tile(e, s, z)):
yield [x, y, z]
class RGBTiler:
"""
Takes continous source data of an arbitrary bit depth and encodes it
in parallel into RGB tiles in an MBTiles file. Provided with a context manager:
```
with RGBTiler(inpath, outpath, min_z, max_x, **kwargs) as tiler:
tiler.run(processes)
```
Parameters
-----------
inpath: string
filepath of the source file to read and encode
outpath: string
filepath of the output `mbtiles`
min_z: int
minimum zoom level to tile
max_z: int
maximum zoom level to tile
Keyword Arguments
------------------
baseval: float
the base value of the RGB numbering system.
(will be treated as zero for this encoding)
Default=0
interval: float
the interval at which to encode
Default=1
format: str
output tile image format (png or webp)
Default=png
bounding_tile: list
[x, y, z] of bounding tile; limits tiled output to this extent
Returns
--------
None
"""
def __init__(
self,
inpath,
outpath,
min_z,
max_z,
interval=1,
base_val=0,
bounding_tile=None,
**kwargs
):
self.run_function = _tile_worker
self.inpath = inpath
self.outpath = outpath
self.min_z = min_z
self.max_z = max_z
self.bounding_tile = bounding_tile
if not "format" in kwargs:
writer_func = _encode_as_png
self.image_format = "png"
elif kwargs["format"].lower() == "png":
writer_func = _encode_as_png
self.image_format = "png"
elif kwargs["format"].lower() == "webp":
writer_func = _encode_as_webp
self.image_format = "webp"
else:
raise ValueError(
"{0} is not a supported filetype!".format(kwargs["format"])
)
# global kwargs not used if output is webp
self.global_args = {
"kwargs": {
"driver": "PNG",
"dtype": "uint8",
"height": 512,
"width": 512,
"count": 3,
"crs": "EPSG:3857",
},
"base_val": base_val,
"interval": interval,
"writer_func": writer_func,
}
def __enter__(self):
return self
def __exit__(self, ext_t, ext_v, trace):
if ext_t:
traceback.print_exc()
def run(self, processes=4):
"""
Warp, encode, and tile
"""
# get the bounding box + crs of the file to tile
with rasterio.open(self.inpath) as src:
bbox = list(src.bounds)
src_crs = src.crs
# remove the output filepath if it exists
if os.path.exists(self.outpath):
os.unlink(self.outpath)
# create a connection to the mbtiles file
conn = sqlite3.connect(self.outpath)
cur = conn.cursor()
# create the tiles table
cur.execute(
"CREATE TABLE tiles "
"(zoom_level integer, tile_column integer, "
"tile_row integer, tile_data blob);"
)
# create empty metadata
cur.execute("CREATE TABLE metadata (name text, value text);")
conn.commit()
# populate metadata with required fields
cur.execute(
"INSERT INTO metadata " "(name, value) " "VALUES ('format', ?);",
(self.image_format,),
)
cur.execute("INSERT INTO metadata " "(name, value) " "VALUES ('name', '');")
cur.execute(
"INSERT INTO metadata " "(name, value) " "VALUES ('description', '');"
)
cur.execute("INSERT INTO metadata " "(name, value) " "VALUES ('version', '1');")
cur.execute(
"INSERT INTO metadata " "(name, value) " "VALUES ('type', 'baselayer');"
)
conn.commit()
if processes == 1:
# use mock pool for profiling / debugging
self.pool = MockTub(
_main_worker, (self.inpath, self.run_function, self.global_args)
)
else:
self.pool = Pool(
processes,
_main_worker,
(self.inpath, self.run_function, self.global_args),
)
# generator of tiles to make
if self.bounding_tile is None:
tiles = _make_tiles(bbox, src_crs, self.min_z, self.max_z)
else:
constrained_bbox = list(mercantile.bounds(self.bounding_tile))
tiles = _make_tiles(constrained_bbox, "epsg:4326", self.min_z, self.max_z)
for tile, contents in self.pool.imap_unordered(self.run_function, tiles):
x, y, z = tile
# mbtiles use inverse y indexing
tiley = int(math.pow(2, z)) - y - 1
# insert tile object
cur.execute(
"INSERT INTO tiles "
"(zoom_level, tile_column, tile_row, tile_data) "
"VALUES (?, ?, ?, ?);",
(z, x, tiley, buffer(contents)),
)
conn.commit()
conn.close()
self.pool.close()
self.pool.join()
return None
