ENH: Tile warping (#48)

contextily/plotting.py

@@ -2,7 +2,9 @@
 
 import numpy as np
 from . import tile_providers as sources
-from .tile import _calculate_zoom, bounds2img, _sm2ll
+from .tile import _calculate_zoom, bounds2img, _sm2ll, warp_tiles, _warper
+from rasterio.enums import Resampling
+from rasterio.warp import transform_bounds
 from matplotlib import patheffects
 
 INTERPOLATION = 'bilinear'
@@ -14,6 +16,7 @@
 def add_basemap(ax, zoom=ZOOM, url=sources.ST_TERRAIN, 
 		interpolation=INTERPOLATION, attribution = ATTRIBUTION, 
         attribution_size = ATTRIBUTION_SIZE, reset_extent=True,
+        crs=None, resampling=Resampling.bilinear,
         **extra_imshow_args):
     """
     Add a (web/local) basemap to `ax`
@@ -47,6 +50,17 @@ def add_basemap(ax, zoom=ZOOM, url=sources.ST_TERRAIN,
                           [Optional. Default=True] If True, the extent of the
                           basemap added is reset to the original extent (xlim,
                           ylim) of `ax`
+    crs                 : None/str/CRS
+                          [Optional. Default=None] CRS,
+                          expressed in any format permitted by rasterio, to
+                          use for the resulting basemap. If
+                          None (default), no warping is performed and the
+                          original Web Mercator (`EPSG:3857`, 
+                          {'init' :'epsg:3857'}) is used.
+    resampling          : <enum 'Resampling'>
+                          [Optional. Default=Resampling.bilinear] Resampling 
+                          method for executing warping, expressed as a 
+                          `rasterio.enums.Resampling` method
     **extra_imshow_args : dict
                           Other parameters to be passed to `imshow`.
 
@@ -82,20 +96,34 @@ def add_basemap(ax, zoom=ZOOM, url=sources.ST_TERRAIN,
     if url[:4] == 'http':
         # Extent
         left, right, bottom, top = xmin, xmax, ymin, ymax
+        # Convert extent from `crs` into WM for tile query
+        if crs is not None:
+            left, right, bottom, top = _reproj_bb(left, right, bottom, top,
+                                                  crs, {'init' :'epsg:3857'})
         # Zoom
         if isinstance(zoom, str) and (zoom.lower() == 'auto'):
             min_ll = _sm2ll(left, bottom)
             max_ll = _sm2ll(right, top)
             zoom = _calculate_zoom(*min_ll, *max_ll)
         image, extent = bounds2img(left, bottom, right, top,
                                    zoom=zoom, url=url, ll=False)
+        # Warping
+        if crs is not None:
+            image, extent = warp_tiles(image, extent, t_crs=crs,
+                                       resampling=resampling)
     # If local source
     else:
         import rasterio as rio
-        # Read extent
+        # Read file
         raster = rio.open(url)
-        image = np.array([ band for band in raster.read() ])\
-                  .transpose(1, 2, 0)
+        image = np.array([ band for band in raster.read() ])
+        # Warp
+        if (crs is not None) and (raster.crs != crs):
+            image, raster = _warper(image, 
+                                    raster.transform,
+                                    raster.crs, crs,
+                                    resampling)
+        image = image.transpose(1, 2, 0)
         bb = raster.bounds
         extent = bb.left, bb.right, bb.bottom, bb.top
     # Plotting
@@ -110,6 +138,13 @@ def add_basemap(ax, zoom=ZOOM, url=sources.ST_TERRAIN,
 
     return ax
 
+
+def _reproj_bb(left, right, bottom, top,
+               s_crs, t_crs):
+    n_l, n_b, n_r, n_t = transform_bounds(s_crs, t_crs,
+                                          left, bottom, right, top)
+    return n_l, n_r, n_b, n_t
+
 def add_attribution(ax, att=ATTRIBUTION, font_size=ATTRIBUTION_SIZE):
     '''
     Utility to add attribution text

contextily/tile.py

@@ -15,13 +15,17 @@
 import numpy as np
 import rasterio as rio
 from PIL import Image
-from rasterio.transform import from_origin
 from joblib import Memory as _Memory
-
+from rasterio.transform import from_origin
+from rasterio.io import MemoryFile
+from rasterio.vrt import WarpedVRT
+from rasterio.enums import Resampling
 from . import tile_providers as sources
 
 
-__all__ = ['bounds2raster', 'bounds2img', 'howmany']
+__all__ = ['bounds2raster', 'bounds2img', 
+           'warp_tiles', 'warp_img_transform',
+           'howmany']
 
 
 USER_AGENT = 'contextily-' + uuid.uuid4().hex
@@ -205,6 +209,125 @@ def _fetch_tile(tile_url, wait, max_retries):
     return image
 
 
+def warp_tiles(img, extent,
+               t_crs='EPSG:4326',
+               resampling=Resampling.bilinear):
+    '''
+    Reproject (warp) a Web Mercator basemap into any CRS on-the-fly
+
+    NOTE: this method works well with contextily's `bounds2img` approach to
+          raster dimensions (h, w, b)
+    ...
+
+    Arguments
+    ---------
+    img         : ndarray
+                  Image as a 3D array (h, w, b) of RGB values (e.g. as
+                  returned from `contextily.bounds2img`)
+    extent      : tuple
+                  Bounding box [minX, maxX, minY, maxY] of the returned image,
+                  expressed in Web Mercator (`EPSG:3857`)
+    t_crs       : str/CRS
+                  [Optional. Default='EPSG:4326'] Target CRS, expressed in any
+                  format permitted by rasterio. Defaults to WGS84 (lon/lat)
+    resampling  : <enum 'Resampling'>
+                  [Optional. Default=Resampling.bilinear] Resampling method for
+                  executing warping, expressed as a `rasterio.enums.Resampling
+                  method
+
+    Returns
+    -------
+    img         : ndarray
+                  Image as a 3D array (h, w, b) of RGB values (e.g. as
+                  returned from `contextily.bounds2img`)
+    ext         : tuple
+                  Bounding box [minX, maxX, minY, maxY] of the returned (warped)
+                  image
+    '''
+    h, w, b = img.shape
+    # --- https://rasterio.readthedocs.io/en/latest/quickstart.html#opening-a-dataset-in-writing-mode
+    minX, maxX, minY, maxY = extent
+    x = np.linspace(minX, maxX, w)
+    y = np.linspace(minY, maxY, h)
+    resX = (x[-1] - x[0]) / w
+    resY = (y[-1] - y[0]) / h
+    transform = from_origin(x[0] - resX / 2,
+                            y[-1] + resY / 2, resX, resY)
+    # ---
+    w_img, vrt = _warper(img.transpose(2, 0, 1),
+                         transform,
+                         'EPSG:3857', t_crs,
+                         resampling)
+    # ---
+    extent = vrt.bounds.left, vrt.bounds.right, \
+             vrt.bounds.bottom, vrt.bounds.top
+    return w_img.transpose(1, 2, 0), extent
+
+
+def warp_img_transform(img, transform, 
+                       s_crs, t_crs,
+                       resampling=Resampling.bilinear):
+    '''
+    Reproject (warp) an `img` with a given `transform` and `s_crs` into a
+    different `t_crs`
+
+    NOTE: this method works well with rasterio's `.read()` approach to
+          raster's dimensions (b, h, w)
+    ...
+
+    Arguments
+    ---------
+    img         : ndarray
+                  Image as a 3D array (b, h, w) of RGB values (e.g. as
+                  returned from rasterio's `.read()` method)
+    transform   : affine.Affine
+                  Transform of the input image as expressed by `rasterio` and
+                  the `affine` package
+    s_crs       : str/CRS
+                  Source CRS in which `img` is passed, expressed in any format
+                  permitted by rasterio.
+    t_crs       : str/CRS
+                  Target CRS, expressed in any format permitted by rasterio.
+    resampling  : <enum 'Resampling'>
+                  [Optional. Default=Resampling.bilinear] Resampling method for
+                  executing warping, expressed as a `rasterio.enums.Resampling
+                  method
+
+    Returns
+    -------
+    w_img       : ndarray
+                  Warped image as a 3D array (b, h, w) of RGB values (e.g. as
+                  returned from rasterio's `.read()` method)
+    w_transform : affine.Affine
+                  Transform of the input image as expressed by `rasterio` and
+                  the `affine` package
+    '''
+    w_img, vrt = _warper(img, transform, 
+                         s_crs, t_crs,
+                         resampling)
+    return w_img, vrt.transform
+
+
+def _warper(img, transform, 
+            s_crs, t_crs,
+            resampling):
+    '''
+    Warp an image returning it as a virtual file
+    '''
+    b, h, w = img.shape
+    with MemoryFile() as memfile:
+        with memfile.open(driver='GTiff', height=h, width=w, \
+                          count=b, dtype=str(img.dtype.name), \
+                          crs=s_crs, transform=transform) as mraster:
+            for band in range(b):
+                mraster.write(img[band, :, :], band+1)
+            # --- Virtual Warp
+            vrt = WarpedVRT(mraster, crs=t_crs,
+                            resampling=resampling)
+            img = vrt.read()
+    return img, vrt
+
+
 def _retryer(tile_url, wait, max_retries):
     """
     Retry a url many times in attempt to get a tile

tests/test_ctx.py

@@ -61,6 +61,29 @@ def test_bounds2img():
     assert img[100, 200, :].tolist() == [156, 180, 131]
     assert img[200, 100, :].tolist() == [230, 225, 189]
 
+def test_warp_tiles():
+    w, s, e, n = (-106.6495132446289, 25.845197677612305,
+            -93.50721740722656, 36.49387741088867)
+    img, ext = ctx.bounds2img(w, s, e, n, zoom=4, ll=True)
+    wimg, wext = ctx.warp_tiles(img, ext)
+    assert wext == (-112.54394531249996, -90.07903186397023,
+                    21.966726124122374,  41.013065787006276)
+    assert wimg[100, 100, :].tolist() == [228, 221, 184]
+    assert wimg[100, 200, :].tolist() == [213, 219, 177]
+    assert wimg[200, 100, :].tolist() == [133, 130, 109]
+
+def test_warp_img_transform():
+    w, s, e, n = ext = (-106.6495132446289, 25.845197677612305,
+                        -93.50721740722656, 36.49387741088867)
+    _ = ctx.bounds2raster(w, s, e, n, 'test.tif', zoom=4, ll=True)
+    rtr = rio.open('test.tif')
+    img = np.array([ band for band in rtr.read() ])
+    wimg, wext = ctx.warp_img_transform(img, rtr.transform,
+                                        rtr.crs, {'init': 'epsg:4326'})
+    assert wimg[:, 100, 100].tolist() == [228, 221, 184]
+    assert wimg[:, 100, 200].tolist() == [213, 219, 177]
+    assert wimg[:, 200, 100].tolist() == [133, 130, 109]
+
 def test_howmany():
     w, s, e, n = (-106.6495132446289, 25.845197677612305,
             -93.50721740722656, 36.49387741088867)
@@ -192,6 +215,38 @@ def test_add_basemap():
     assert_array_almost_equal(ax.images[0].get_array().mean(),
                               184.10206197102863)
 
+    # Test on-th-fly warping
+    x1, x2 = -105.5, -105.00
+    y1, y2 = 39.56, 40.13
+    f, ax = matplotlib.pyplot.subplots(1)
+    ax.set_xlim(x1, x2)
+    ax.set_ylim(y1, y2)
+    ax = ctx.add_basemap(ax, 
+                         crs={'init': 'epsg:4326'},
+                         attribution=None)
+    assert ax.get_xlim() == (x1, x2)
+    assert ax.get_ylim() == (y1, y2)
+    assert ax.images[0].get_array().sum() == 724238693
+    assert ax.images[0].get_array().shape == (1135, 1183, 3)
+    assert_array_almost_equal(ax.images[0].get_array().mean(),
+                              179.79593258881636)
+    # Test local source warping
+    _ = ctx.bounds2raster(x1, y1, x2, y2, "./test2.tif", ll=True)
+    f, ax = matplotlib.pyplot.subplots(1)
+    ax.set_xlim(x1, x2)
+    ax.set_ylim(y1, y2)
+    ax = ctx.add_basemap(ax, 
+                         url="./test2.tif",
+                         crs={'init': 'epsg:4326'},
+                         attribution=None)
+    assert ax.get_xlim() == (x1, x2)
+    assert ax.get_ylim() == (y1, y2)
+    assert ax.images[0].get_array().sum() == 724238693
+    assert ax.images[0].get_array().shape == (1135, 1183, 3)
+    assert_array_almost_equal(ax.images[0].get_array().mean(),
+                              179.79593258881636)
+
+
 def test_attribution():
     f, ax = matplotlib.pyplot.subplots(1)
     ax = ctx.add_attribution(ax, 'Test')