├── rio_rgbify ├── scripts │ ├── __init__.py │ └── cli.py ├── __init__.py ├── encoders.py └── mbtiler.py ├── requirements.txt ├── test ├── fixtures │ └── elev.tif ├── expected │ └── elev-rgb.tif ├── test_encoders.py ├── test_mbtiler.py └── test_cli.py ├── codecov.yml ├── requirements-dev.txt ├── tox.ini ├── .travis.yml ├── .gitignore ├── LICENSE ├── .pre-commit-config.yaml ├── setup.py └── README.md /rio_rgbify/scripts/__init__.py: -------------------------------------------------------------------------------- 1 | """rio-rgbify cli.""" 2 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | click 2 | rasterio~=1.0 3 | Pillow 4 | rio-mucho>=0.2.1 5 | mercantile -------------------------------------------------------------------------------- /test/fixtures/elev.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mapbox/rio-rgbify/HEAD/test/fixtures/elev.tif -------------------------------------------------------------------------------- /test/expected/elev-rgb.tif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mapbox/rio-rgbify/HEAD/test/expected/elev-rgb.tif -------------------------------------------------------------------------------- /rio_rgbify/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | __version__ = "0.4.0" 4 | 5 | log = logging.getLogger(__name__) 6 | log.addHandler(logging.NullHandler()) 7 | -------------------------------------------------------------------------------- /codecov.yml: -------------------------------------------------------------------------------- 1 | comment: off 2 | 3 | coverage: 4 | status: 5 | project: 6 | default: 7 | target: auto 8 | threshold: 5 9 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | affine 2 | cligj 3 | coveralls>=0.4 4 | delocate 5 | enum34 6 | hypothesis 7 | numpy>=1.8.0 8 | snuggs>=1.2 9 | pytest 10 | raster-tester 11 | setuptools>=0.9.8 12 | wheel 13 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py37,py36,py27 3 | 4 | 5 | [testenv] 6 | extras = test 7 | commands= 8 | python -m pytest --cov rio_rgbify --cov-report term-missing --ignore=venv 9 | deps= 10 | numpy 11 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | sudo: false 3 | cache: 4 | directories: 5 | - ~/.cache/pip 6 | 7 | matrix: 8 | include: 9 | - python: 2.7 10 | env: TOXENV=py27 11 | - python: 3.6 12 | env: TOXENV=py36 13 | - python: 3.7 14 | env: TOXENV=py37 15 | 16 | before_install: 17 | - python -m pip install -U pip 18 | install: 19 | - python -m pip install codecov pre-commit tox 20 | script: 21 | - if [[ $TRAVIS_PYTHON_VERSION == 3.6 ]]; then pre-commit run --all-files; fi 22 | - tox 23 | after_success: 24 | - if [[ $TRAVIS_PYTHON_VERSION == 3.6 ]]; then coverage xml; codecov; fi 25 | 26 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | lib/ 17 | lib64/ 18 | parts/ 19 | sdist/ 20 | var/ 21 | *.egg-info/ 22 | .installed.cfg 23 | *.egg 24 | 25 | # PyInstaller 26 | # Usually these files are written by a python script from a template 27 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 28 | *.manifest 29 | *.spec 30 | 31 | # Installer logs 32 | pip-log.txt 33 | pip-delete-this-directory.txt 34 | 35 | # Unit test / coverage reports 36 | htmlcov/ 37 | .tox/ 38 | .coverage 39 | .cache 40 | nosetests.xml 41 | coverage.xml 42 | 43 | # Translations 44 | *.mo 45 | *.pot 46 | 47 | # Django stuff: 48 | *.log 49 | 50 | # Sphinx documentation 51 | docs/_build/ 52 | 53 | # PyBuilder 54 | target/ 55 | 56 | # OS X 57 | .DS_Store 58 | 59 | # Pyenv 60 | .python-version 61 | 62 | .*.swp 63 | 64 | docs/notebooks/.ipynb_checkpoints/* 65 | 66 | .pytest_cache 67 | .hypothesis 68 | .coverage* 69 | -------------------------------------------------------------------------------- /test/test_encoders.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | from rio_rgbify.encoders import data_to_rgb, _decode, _range_check 3 | import numpy as np 4 | import pytest 5 | 6 | 7 | def test_encode_data_roundtrip(): 8 | minrand, maxrand = np.sort(np.random.randint(-427, 8848, 2)) 9 | 10 | testdata = np.round((np.sum( 11 | np.dstack( 12 | np.indices((512, 512), 13 | dtype=np.float64)), 14 | axis=2) / (511. + 511.)) * maxrand, 2) + minrand 15 | 16 | baseval = -1000 17 | interval = 0.1 18 | round_digits = 0 19 | 20 | rtripped = _decode(data_to_rgb(testdata.copy(), baseval, interval, round_digits=round_digits), baseval, interval) 21 | 22 | assert testdata.min() == rtripped.min() 23 | assert testdata.max() == rtripped.max() 24 | 25 | 26 | def test_encode_failrange(): 27 | testdata = np.zeros((2)) 28 | 29 | testdata[1] = 256 ** 3 + 1 30 | 31 | with pytest.raises(ValueError): 32 | data_to_rgb(testdata, 0, 1, 0) 33 | 34 | 35 | def test_catch_range(): 36 | assert _range_check(256 ** 3 + 1) 37 | assert not _range_check(256 ** 3 - 1) 38 | 39 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Mapbox 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | 2 | repos: 3 | - 4 | repo: 'https://github.com/ambv/black' 5 | # 18.6b1 6 | rev: ed50737290662f6ef4016a7ea44da78ee1eff1e2 7 | hooks: 8 | - id: black 9 | args: ['--safe'] 10 | language_version: python3.6 11 | - 12 | repo: 'https://github.com/pre-commit/pre-commit-hooks' 13 | # v1.3.0 14 | rev: a6209d8d4f97a09b61855ea3f1fb250f55147b8b 15 | hooks: 16 | - id: flake8 17 | language_version: python3.6 18 | args: [ 19 | # E501 let black handle all line length decisions 20 | # W503 black conflicts with "line break before operator" rule 21 | # E203 black conflicts with "whitespace before ':'" rule 22 | '--ignore=E501,W503,E203'] 23 | - 24 | repo: 'https://github.com/chewse/pre-commit-mirrors-pydocstyle' 25 | # 2.1.1 26 | rev: 22d3ccf6cf91ffce3b16caa946c155778f0cb20f 27 | hooks: 28 | - id: pydocstyle 29 | language_version: python3.6 30 | args: [ 31 | # Check for docstring presence only 32 | '--select=D1', 33 | # Don't require docstrings for tests 34 | '--match=(?!test).*\.py'] 35 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | """rio-rgbify: setup.""" 2 | 3 | from setuptools import setup, find_packages 4 | 5 | # Parse the version from the fiona module. 6 | with open("rio_rgbify/__init__.py") as f: 7 | for line in f: 8 | if line.find("__version__") >= 0: 9 | version = line.split("=")[1].strip() 10 | version = version.strip('"') 11 | version = version.strip("'") 12 | break 13 | 14 | long_description = """""" 15 | 16 | # Runtime requirements. 17 | inst_reqs = ["click", "rasterio~=1.0", "rio-mucho", "Pillow", "mercantile"] 18 | 19 | extra_reqs = { 20 | "test": ["pytest", "pytest-cov", "codecov", "hypothesis", "raster_tester"], 21 | "dev": [ 22 | "pytest", "pytest-cov", "codecov", "hypothesis", "raster_tester", "pre-commit" 23 | ], 24 | } 25 | 26 | setup(name="rio-rgbify", 27 | version=version, 28 | description=u"Encode arbitrary bit depth rasters in pseudo base-256 as RGB", 29 | long_description=long_description, 30 | classifiers=[], 31 | keywords="", 32 | author=u"Damon Burgett", 33 | author_email="damon@mapbox.com", 34 | url="https://github.com/mapbox/rio-rgbify", 35 | license="BSD", 36 | packages=find_packages(exclude=["ez_setup", "examples", "tests"]), 37 | include_package_data=True, 38 | zip_safe=False, 39 | install_requires=inst_reqs, 40 | extras_require=extra_reqs, 41 | entry_points=""" 42 | [rasterio.rio_plugins] 43 | rgbify=rio_rgbify.scripts.cli:rgbify 44 | """) 45 | -------------------------------------------------------------------------------- /rio_rgbify/encoders.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import numpy as np 3 | 4 | 5 | def data_to_rgb(data, baseval, interval, round_digits=0): 6 | """ 7 | Given an arbitrary (rows x cols) ndarray, 8 | encode the data into uint8 RGB from an arbitrary 9 | base and interval 10 | 11 | Parameters 12 | ----------- 13 | data: ndarray 14 | (rows x cols) ndarray of data to encode 15 | baseval: float 16 | the base value of the RGB numbering system. 17 | will be treated as zero for this encoding 18 | interval: float 19 | the interval at which to encode 20 | round_digits: int 21 | erased less significant digits 22 | 23 | Returns 24 | -------- 25 | ndarray: rgb data 26 | a uint8 (3 x rows x cols) ndarray with the 27 | data encoded 28 | """ 29 | data = data.astype(np.float64) 30 | data -= baseval 31 | data /= interval 32 | 33 | data = np.around(data / 2**round_digits) * 2**round_digits 34 | 35 | rows, cols = data.shape 36 | 37 | datarange = data.max() - data.min() 38 | 39 | if _range_check(datarange): 40 | raise ValueError("Data of {} larger than 256 ** 3".format(datarange)) 41 | 42 | rgb = np.zeros((3, rows, cols), dtype=np.uint8) 43 | 44 | rgb[2] = ((data / 256) - (data // 256)) * 256 45 | rgb[1] = (((data // 256) / 256) - ((data // 256) // 256)) * 256 46 | rgb[0] = ((((data // 256) // 256) / 256) - (((data // 256) // 256) // 256)) * 256 47 | 48 | return rgb 49 | 50 | 51 | def _decode(data, base, interval): 52 | """ 53 | Utility to decode RGB encoded data 54 | """ 55 | data = data.astype(np.float64) 56 | return base + (((data[0] * 256 * 256) + (data[1] * 256) + data[2]) * interval) 57 | 58 | 59 | def _range_check(datarange): 60 | """ 61 | Utility to check if data range is outside of precision for 3 digit base 256 62 | """ 63 | maxrange = 256 ** 3 64 | 65 | return datarange > maxrange 66 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rio-rgbify 2 | Encode arbitrary bit depth rasters in pseudo base-256 as RGB 3 | 4 | [![Build Status](https://travis-ci.org/mapbox/rio-rgbify.svg)](https://travis-ci.org/mapbox/rio-rgbify)[![Coverage Status](https://coveralls.io/repos/github/mapbox/rio-rgbify/badge.svg?branch=its-a-setup)](https://coveralls.io/github/mapbox/rio-rgbify) 5 | 6 | ## Installation 7 | 8 | ### From PyPi 9 | ``` 10 | pip install rio-rgbify 11 | ``` 12 | ### Development 13 | ``` 14 | git clone git@github.com:mapbox/rio-rgbify.git 15 | 16 | cd rio-rgbify 17 | 18 | pip install -e '.[test]' 19 | 20 | ``` 21 | 22 | ## CLI usage 23 | 24 | - Input can be any raster readable by `rasterio` 25 | - Output can be any raster format writable by `rasterio` OR 26 | - To create tiles _directly_ from data (recommended), output to an `.mbtiles` 27 | 28 | ``` 29 | Usage: rio rgbify [OPTIONS] SRC_PATH DST_PATH 30 | 31 | Options: 32 | -b, --base-val FLOAT The base value of which to base the output encoding 33 | on [DEFAULT=0] 34 | -i, --interval FLOAT Describes the precision of the output, by 35 | incrementing interval [DEFAULT=1] 36 | -r, --round-digits Less significants encoded bits to be set 37 | to 0. Round the values, but have better 38 | images compression [DEFAULT=0] 39 | --bidx INTEGER Band to encode [DEFAULT=1] 40 | --max-z INTEGER Maximum zoom to tile (.mbtiles output only) 41 | --bounding-tile TEXT Bounding tile '[{x}, {y}, {z}]' to limit output tiles 42 | (.mbtiles output only) 43 | --min-z INTEGER Minimum zoom to tile (.mbtiles output only) 44 | --format [png|webp] Output tile format (.mbtiles output only) 45 | -j, --workers INTEGER Workers to run [DEFAULT=4] 46 | -v, --verbose 47 | --co NAME=VALUE Driver specific creation options.See the 48 | documentation for the selected output driver for more 49 | information. 50 | --help Show this message and exit. 51 | ``` 52 | -------------------------------------------------------------------------------- /test/test_mbtiler.py: -------------------------------------------------------------------------------- 1 | import mercantile 2 | import types 3 | 4 | from hypothesis import given 5 | import hypothesis.strategies as st 6 | import pytest 7 | 8 | import numpy as np 9 | from rasterio import Affine 10 | from rio_rgbify.mbtiler import (_encode_as_webp, _encode_as_png, _make_tiles, _tile_range, RGBTiler) 11 | 12 | 13 | @given( 14 | st.integers( 15 | min_value=0, max_value=(2 ** 10 - 1) 16 | ), 17 | st.integers( 18 | min_value=0, max_value=(2 ** 10 - 1))) 19 | def test_make_tiles_tile_bounds(x, y): 20 | ''' 21 | Test if children tiles from z10 are created correctly 22 | ''' 23 | test_bounds = mercantile.bounds(x, y, 10) 24 | 25 | test_bbox = list(mercantile.xy(test_bounds.west, test_bounds.south)) + list(mercantile.xy(test_bounds.east, test_bounds.north)) 26 | 27 | test_crs = 'epsg:3857' 28 | test_minz = 10 29 | test_maxz = 13 30 | 31 | created_tiles_gen = _make_tiles(test_bbox, test_crs, test_minz, test_maxz) 32 | 33 | assert isinstance(created_tiles_gen, types.GeneratorType) 34 | 35 | created_tiles = list(created_tiles_gen) 36 | 37 | assert len(created_tiles) == 85 38 | 39 | 40 | @given( 41 | st.lists( 42 | elements=st.integers(min_value=0, max_value=99), 43 | max_size=3, min_size=3 44 | ), 45 | st.lists( 46 | elements=st.integers( 47 | min_value=100, max_value=200 48 | ), max_size=3, min_size=3 49 | )) 50 | def test_tile_range(mintile, maxtile): 51 | minx, miny, _ = mintile 52 | maxx, maxy, _ = maxtile 53 | 54 | expected_length = (maxx - minx + 1) * (maxy - miny + 1) 55 | assert expected_length == len(list(_tile_range(mintile, maxtile))) 56 | 57 | 58 | def test_webp_writer(): 59 | test_data = np.zeros((3, 256, 256), dtype=np.uint8) 60 | 61 | test_bytearray = _encode_as_webp(test_data) 62 | 63 | assert len(test_bytearray) == 34 64 | 65 | test_complex_data = test_data.copy() 66 | 67 | test_complex_data[0] += (np.random.rand(256, 256) * 255).astype(np.uint8) 68 | test_complex_data[1] += 10 69 | 70 | test_bytearray_complex = _encode_as_webp(test_complex_data) 71 | 72 | assert len(test_bytearray) < len(test_bytearray_complex) 73 | 74 | 75 | def test_file_writer(): 76 | test_data = np.zeros((3, 256, 256), dtype=np.uint8) 77 | 78 | test_opts = { 79 | 'driver': 'PNG', 80 | 'dtype': 'uint8', 81 | 'height': 512, 82 | 'width': 512, 83 | 'count': 3, 84 | 'crs': 'EPSG:3857' 85 | } 86 | 87 | test_affine = Affine(1, 0, 0, 0, -1, 0) 88 | 89 | test_bytearray = _encode_as_png(test_data, test_opts, test_affine) 90 | 91 | assert len(test_bytearray) == 842 92 | 93 | test_complex_data = test_data.copy() 94 | 95 | test_complex_data[0] += (np.random.rand(256, 256) * 255).astype(np.uint8) 96 | test_complex_data[1] += 10 97 | 98 | test_bytearray_complex = _encode_as_png(test_complex_data, test_opts, test_affine) 99 | 100 | assert len(test_bytearray) < len(test_bytearray_complex) 101 | 102 | 103 | def test_webp_writer_fails_dtype(): 104 | test_data = np.zeros((3, 256, 256), dtype=np.float64) 105 | 106 | with pytest.raises(TypeError): 107 | _encode_as_webp(test_data) 108 | 109 | 110 | def test_png_writer_fails_dtype(): 111 | test_data = np.zeros((3, 256, 256), dtype=np.float64) 112 | 113 | with pytest.raises(TypeError): 114 | _encode_as_png(test_data) 115 | 116 | 117 | def test_RGBtiler_format_fails(): 118 | test_in = 'i/do/not/exist.tif' 119 | test_out = 'nor/do/i.tif' 120 | test_minz = 0 121 | test_maxz = 1 122 | 123 | with pytest.raises(ValueError): 124 | with RGBTiler(test_in, test_out, test_minz, test_maxz, 125 | format='poo') as rtiler: 126 | pass 127 | -------------------------------------------------------------------------------- /rio_rgbify/scripts/cli.py: -------------------------------------------------------------------------------- 1 | """rio_rgbify CLI.""" 2 | 3 | import click 4 | 5 | import rasterio as rio 6 | import numpy as np 7 | from riomucho import RioMucho 8 | import json 9 | from rasterio.rio.options import creation_options 10 | 11 | from rio_rgbify.encoders import data_to_rgb 12 | from rio_rgbify.mbtiler import RGBTiler 13 | 14 | 15 | def _rgb_worker(data, window, ij, g_args): 16 | return data_to_rgb( 17 | data[0][g_args["bidx"] - 1], g_args["base_val"], g_args["interval"], g_args["round_digits"] 18 | ) 19 | 20 | 21 | @click.command("rgbify") 22 | @click.argument("src_path", type=click.Path(exists=True)) 23 | @click.argument("dst_path", type=click.Path(exists=False)) 24 | @click.option( 25 | "--base-val", 26 | "-b", 27 | type=float, 28 | default=0, 29 | help="The base value of which to base the output encoding on [DEFAULT=0]", 30 | ) 31 | @click.option( 32 | "--interval", 33 | "-i", 34 | type=float, 35 | default=1, 36 | help="Describes the precision of the output, by incrementing interval [DEFAULT=1]", 37 | ) 38 | @click.option( 39 | "--round-digits", 40 | "-r", 41 | type=int, 42 | default=0, 43 | help="Less significants encoded bits to be set to 0. Round the values, but have better images compression [DEFAULT=0]", 44 | ) 45 | @click.option("--bidx", type=int, default=1, help="Band to encode [DEFAULT=1]") 46 | @click.option( 47 | "--max-z", 48 | type=int, 49 | default=None, 50 | help="Maximum zoom to tile (.mbtiles output only)", 51 | ) 52 | @click.option( 53 | "--bounding-tile", 54 | type=str, 55 | default=None, 56 | help="Bounding tile '[{x}, {y}, {z}]' to limit output tiles (.mbtiles output only)", 57 | ) 58 | @click.option( 59 | "--min-z", 60 | type=int, 61 | default=None, 62 | help="Minimum zoom to tile (.mbtiles output only)", 63 | ) 64 | @click.option( 65 | "--format", 66 | type=click.Choice(["png", "webp"]), 67 | default="png", 68 | help="Output tile format (.mbtiles output only)", 69 | ) 70 | @click.option("--workers", "-j", type=int, default=4, help="Workers to run [DEFAULT=4]") 71 | @click.option("--verbose", "-v", is_flag=True, default=False) 72 | @click.pass_context 73 | @creation_options 74 | def rgbify( 75 | ctx, 76 | src_path, 77 | dst_path, 78 | base_val, 79 | interval, 80 | round_digits, 81 | bidx, 82 | max_z, 83 | min_z, 84 | bounding_tile, 85 | format, 86 | workers, 87 | verbose, 88 | creation_options, 89 | ): 90 | """rio-rgbify cli.""" 91 | if dst_path.split(".")[-1].lower() == "tif": 92 | with rio.open(src_path) as src: 93 | meta = src.profile.copy() 94 | 95 | meta.update(count=3, dtype=np.uint8) 96 | 97 | for c in creation_options: 98 | meta[c] = creation_options[c] 99 | 100 | gargs = {"interval": interval, "base_val": base_val, "round_digits": round_digits, "bidx": bidx} 101 | 102 | with RioMucho( 103 | [src_path], dst_path, _rgb_worker, options=meta, global_args=gargs 104 | ) as rm: 105 | 106 | rm.run(workers) 107 | 108 | elif dst_path.split(".")[-1].lower() == "mbtiles": 109 | if min_z is None or max_z is None: 110 | raise ValueError("Zoom range must be provided for mbtile output") 111 | 112 | if max_z < min_z: 113 | raise ValueError( 114 | "Max zoom {0} must be greater than min zoom {1}".format(max_z, min_z) 115 | ) 116 | 117 | if bounding_tile is not None: 118 | try: 119 | bounding_tile = json.loads(bounding_tile) 120 | except Exception: 121 | raise TypeError( 122 | "Bounding tile of {0} is not valid".format(bounding_tile) 123 | ) 124 | 125 | with RGBTiler( 126 | src_path, 127 | dst_path, 128 | interval=interval, 129 | base_val=base_val, 130 | round_digits=round_digits, 131 | format=format, 132 | bounding_tile=bounding_tile, 133 | max_z=max_z, 134 | min_z=min_z, 135 | ) as tiler: 136 | tiler.run(workers) 137 | 138 | else: 139 | raise ValueError( 140 | "{} output filetype not supported".format(dst_path.split(".")[-1]) 141 | ) 142 | -------------------------------------------------------------------------------- /test/test_cli.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import click 4 | from click.testing import CliRunner 5 | 6 | import numpy as np 7 | 8 | import rasterio as rio 9 | from rio_rgbify.scripts.cli import rgbify 10 | 11 | from raster_tester.compare import affaux, upsample_array 12 | 13 | 14 | in_elev_src = os.path.join(os.path.dirname(__file__), "fixtures", "elev.tif") 15 | expected_src = os.path.join(os.path.dirname(__file__), "expected", "elev-rgb.tif") 16 | 17 | 18 | def flex_compare(r1, r2, thresh=10): 19 | upsample = 4 20 | r1 = r1[::upsample] 21 | r2 = r2[::upsample] 22 | toAff, frAff = affaux(upsample) 23 | r1 = upsample_array(r1, upsample, frAff, toAff) 24 | r2 = upsample_array(r2, upsample, frAff, toAff) 25 | tdiff = np.abs(r1.astype(np.float64) - r2.astype(np.float64)) 26 | 27 | click.echo( 28 | "{0} values exceed the threshold difference with a max variance of {1}".format( 29 | np.sum(tdiff > thresh), tdiff.max() 30 | ), 31 | err=True, 32 | ) 33 | 34 | return not np.any(tdiff > thresh) 35 | 36 | 37 | def test_cli_good_elev(): 38 | runner = CliRunner() 39 | with runner.isolated_filesystem(): 40 | result = runner.invoke( 41 | rgbify, 42 | [in_elev_src, "rgb.tif", "--interval", 0.001, "--base-val", -100, "-j", 1], 43 | ) 44 | 45 | assert result.exit_code == 0 46 | 47 | with rio.open("rgb.tif") as created: 48 | with rio.open(expected_src) as expected: 49 | carr = created.read() 50 | earr = expected.read() 51 | for a, b in zip(carr, earr): 52 | assert flex_compare(a, b) 53 | 54 | 55 | def test_cli_fail_elev(): 56 | runner = CliRunner() 57 | with runner.isolated_filesystem(): 58 | result = runner.invoke( 59 | rgbify, 60 | [ 61 | in_elev_src, 62 | "rgb.tif", 63 | "--interval", 64 | 0.00000001, 65 | "--base-val", 66 | -100, 67 | "-j", 68 | 1, 69 | ], 70 | ) 71 | assert result.exit_code == 1 72 | assert result.exception 73 | 74 | 75 | def test_mbtiler_webp(): 76 | runner = CliRunner() 77 | with runner.isolated_filesystem(): 78 | out_mbtiles_finer = "output-0-dot-1.mbtiles" 79 | result_finer = runner.invoke( 80 | rgbify, 81 | [ 82 | in_elev_src, 83 | out_mbtiles_finer, 84 | "--interval", 85 | 0.1, 86 | "--min-z", 87 | 10, 88 | "--max-z", 89 | 11, 90 | "--format", 91 | "webp", 92 | "-j", 93 | 1, 94 | ], 95 | ) 96 | assert result_finer.exit_code == 0 97 | 98 | out_mbtiles_coarser = "output-1.mbtiles" 99 | result_coarser = runner.invoke( 100 | rgbify, 101 | [ 102 | in_elev_src, 103 | out_mbtiles_coarser, 104 | "--min-z", 105 | 10, 106 | "--max-z", 107 | 11, 108 | "--format", 109 | "webp", 110 | "-j", 111 | 1, 112 | ], 113 | ) 114 | assert result_coarser.exit_code == 0 115 | 116 | assert os.path.getsize(out_mbtiles_finer) > os.path.getsize(out_mbtiles_coarser) 117 | 118 | 119 | def test_mbtiler_png(): 120 | runner = CliRunner() 121 | with runner.isolated_filesystem(): 122 | out_mbtiles_finer = "output-0-dot-1.mbtiles" 123 | result_finer = runner.invoke( 124 | rgbify, 125 | [ 126 | in_elev_src, 127 | out_mbtiles_finer, 128 | "--interval", 129 | 0.1, 130 | "--min-z", 131 | 10, 132 | "--max-z", 133 | 11, 134 | "--format", 135 | "png", 136 | ], 137 | ) 138 | assert result_finer.exit_code == 0 139 | 140 | out_mbtiles_coarser = "output-1.mbtiles" 141 | result_coarser = runner.invoke( 142 | rgbify, 143 | [ 144 | in_elev_src, 145 | out_mbtiles_coarser, 146 | "--min-z", 147 | 10, 148 | "--max-z", 149 | 11, 150 | "--format", 151 | "png", 152 | "-j", 153 | 1, 154 | ], 155 | ) 156 | assert result_coarser.exit_code == 0 157 | 158 | assert os.path.getsize(out_mbtiles_finer) > os.path.getsize(out_mbtiles_coarser) 159 | 160 | 161 | def test_mbtiler_png_bounding_tile(): 162 | runner = CliRunner() 163 | with runner.isolated_filesystem(): 164 | out_mbtiles_not_limited = "output-not-limited.mbtiles" 165 | result_not_limited = runner.invoke( 166 | rgbify, 167 | [ 168 | in_elev_src, 169 | out_mbtiles_not_limited, 170 | "--min-z", 171 | 12, 172 | "--max-z", 173 | 12, 174 | "--format", 175 | "png", 176 | ], 177 | ) 178 | assert result_not_limited.exit_code == 0 179 | 180 | out_mbtiles_limited = "output-limited.mbtiles" 181 | result_limited = runner.invoke( 182 | rgbify, 183 | [ 184 | in_elev_src, 185 | out_mbtiles_limited, 186 | "--min-z", 187 | 12, 188 | "--max-z", 189 | 12, 190 | "--format", 191 | "png", 192 | "--bounding-tile", 193 | "[654, 1582, 12]", 194 | ], 195 | ) 196 | assert result_limited.exit_code == 0 197 | 198 | assert os.path.getsize(out_mbtiles_not_limited) > os.path.getsize( 199 | out_mbtiles_limited 200 | ) 201 | 202 | result_badtile = runner.invoke( 203 | rgbify, 204 | [ 205 | in_elev_src, 206 | out_mbtiles_limited, 207 | "--min-z", 208 | 12, 209 | "--max-z", 210 | 12, 211 | "--format", 212 | "png", 213 | "--bounding-tile", 214 | "654-1582-12", 215 | ], 216 | ) 217 | assert result_badtile.exit_code == 1 218 | assert "is not valid" in str(result_badtile.exception) 219 | 220 | 221 | def test_mbtiler_webp_badzoom(): 222 | runner = CliRunner() 223 | with runner.isolated_filesystem(): 224 | out_mbtiles = "output.mbtiles" 225 | result = runner.invoke( 226 | rgbify, 227 | [ 228 | in_elev_src, 229 | out_mbtiles, 230 | "--min-z", 231 | 10, 232 | "--max-z", 233 | 9, 234 | "--format", 235 | "webp", 236 | "-j", 237 | 1, 238 | ], 239 | ) 240 | assert result.exit_code == 1 241 | assert result.exception 242 | 243 | 244 | def test_mbtiler_webp_badboundingtile(): 245 | runner = CliRunner() 246 | with runner.isolated_filesystem(): 247 | out_mbtiles = "output.mbtiles" 248 | result = runner.invoke( 249 | rgbify, 250 | [ 251 | in_elev_src, 252 | out_mbtiles, 253 | "--min-z", 254 | 10, 255 | "--max-z", 256 | 9, 257 | "--format", 258 | "webp", 259 | "--bounding-tile", 260 | "654, 1582, 12", 261 | ], 262 | ) 263 | assert result.exit_code == 1 264 | assert result.exception 265 | 266 | 267 | def test_mbtiler_webp_badboundingtile_values(): 268 | runner = CliRunner() 269 | with runner.isolated_filesystem(): 270 | out_mbtiles = "output.mbtiles" 271 | result = runner.invoke( 272 | rgbify, 273 | [ 274 | in_elev_src, 275 | out_mbtiles, 276 | "--min-z", 277 | 10, 278 | "--max-z", 279 | 9, 280 | "--format", 281 | "webp", 282 | "--bounding-tile", 283 | "[654, 1582]", 284 | ], 285 | ) 286 | assert result.exit_code == 1 287 | assert result.exception 288 | 289 | 290 | def test_bad_input_format(): 291 | runner = CliRunner() 292 | with runner.isolated_filesystem(): 293 | out_mbtiles = "output.lol" 294 | result = runner.invoke( 295 | rgbify, 296 | [ 297 | in_elev_src, 298 | out_mbtiles, 299 | "--min-z", 300 | 10, 301 | "--max-z", 302 | 9, 303 | "--format", 304 | "webp", 305 | "-j", 306 | 1, 307 | ], 308 | ) 309 | assert result.exit_code == 1 310 | assert result.exception 311 | -------------------------------------------------------------------------------- /rio_rgbify/mbtiler.py: -------------------------------------------------------------------------------- 1 | from __future__ import with_statement 2 | from __future__ import division 3 | 4 | import os 5 | import sys 6 | import math 7 | import traceback 8 | import itertools 9 | 10 | import mercantile 11 | import rasterio 12 | import numpy as np 13 | import sqlite3 14 | from multiprocessing import Pool 15 | from rasterio._io import virtual_file_to_buffer 16 | from riomucho.single_process_pool import MockTub 17 | 18 | from io import BytesIO 19 | from PIL import Image 20 | 21 | from rasterio import transform 22 | from rasterio.warp import reproject, transform_bounds 23 | 24 | from rasterio.enums import Resampling 25 | 26 | from rio_rgbify.encoders import data_to_rgb 27 | 28 | buffer = bytes if sys.version_info > (3,) else buffer 29 | 30 | work_func = None 31 | global_args = None 32 | src = None 33 | 34 | 35 | def _main_worker(inpath, g_work_func, g_args): 36 | """ 37 | Util for setting global vars w/ a Pool 38 | """ 39 | global work_func 40 | global global_args 41 | global src 42 | work_func = g_work_func 43 | global_args = g_args 44 | 45 | src = rasterio.open(inpath) 46 | 47 | 48 | def _encode_as_webp(data, profile=None, affine=None): 49 | """ 50 | Uses BytesIO + PIL to encode a (3, 512, 512) 51 | array into a webp bytearray. 52 | 53 | Parameters 54 | ----------- 55 | data: ndarray 56 | (3 x 512 x 512) uint8 RGB array 57 | profile: None 58 | ignored 59 | affine: None 60 | ignored 61 | 62 | Returns 63 | -------- 64 | contents: bytearray 65 | webp-encoded bytearray of the provided input data 66 | """ 67 | with BytesIO() as f: 68 | im = Image.fromarray(np.rollaxis(data, 0, 3)) 69 | im.save(f, format="webp", lossless=True) 70 | 71 | return f.getvalue() 72 | 73 | 74 | def _encode_as_png(data, profile, dst_transform): 75 | """ 76 | Uses rasterio's virtual file system to encode a (3, 512, 512) 77 | array as a png-encoded bytearray. 78 | 79 | Parameters 80 | ----------- 81 | data: ndarray 82 | (3 x 512 x 512) uint8 RGB array 83 | profile: dictionary 84 | dictionary of kwargs for png writing 85 | affine: Affine 86 | affine transform for output tile 87 | 88 | Returns 89 | -------- 90 | contents: bytearray 91 | png-encoded bytearray of the provided input data 92 | """ 93 | profile["affine"] = dst_transform 94 | 95 | with rasterio.open("/vsimem/tileimg", "w", **profile) as dst: 96 | dst.write(data) 97 | 98 | contents = bytearray(virtual_file_to_buffer("/vsimem/tileimg")) 99 | 100 | return contents 101 | 102 | 103 | def _tile_worker(tile): 104 | """ 105 | For each tile, and given an open rasterio src, plus a`global_args` dictionary 106 | with attributes of `base_val`, `interval`, `round_digits` and a `writer_func`, 107 | warp a continous single band raster to a 512 x 512 mercator tile, 108 | then encode this tile into RGB. 109 | 110 | Parameters 111 | ----------- 112 | tile: list 113 | [x, y, z] indices of tile 114 | 115 | Returns 116 | -------- 117 | tile, buffer 118 | tuple with the input tile, and a bytearray with the data encoded into 119 | the format created in the `writer_func` 120 | 121 | """ 122 | x, y, z = tile 123 | 124 | bounds = [ 125 | c 126 | for i in ( 127 | mercantile.xy(*mercantile.ul(x, y + 1, z)), 128 | mercantile.xy(*mercantile.ul(x + 1, y, z)), 129 | ) 130 | for c in i 131 | ] 132 | 133 | toaffine = transform.from_bounds(*bounds + [512, 512]) 134 | 135 | out = np.empty((512, 512), dtype=src.meta["dtype"]) 136 | 137 | reproject( 138 | rasterio.band(src, 1), 139 | out, 140 | dst_transform=toaffine, 141 | dst_crs="EPSG:3857", 142 | resampling=Resampling.bilinear, 143 | ) 144 | 145 | out = data_to_rgb(out, global_args["base_val"], global_args["interval"], global_args["round_digits"]) 146 | 147 | return tile, global_args["writer_func"](out, global_args["kwargs"].copy(), toaffine) 148 | 149 | 150 | def _tile_range(min_tile, max_tile): 151 | """ 152 | Given a min and max tile, return an iterator of 153 | all combinations of this tile range 154 | 155 | Parameters 156 | ----------- 157 | min_tile: list 158 | [x, y, z] of minimun tile 159 | max_tile: 160 | [x, y, z] of minimun tile 161 | 162 | Returns 163 | -------- 164 | tiles: iterator 165 | iterator of [x, y, z] tiles 166 | """ 167 | min_x, min_y, _ = min_tile 168 | max_x, max_y, _ = max_tile 169 | 170 | return itertools.product(range(min_x, max_x + 1), range(min_y, max_y + 1)) 171 | 172 | 173 | def _make_tiles(bbox, src_crs, minz, maxz): 174 | """ 175 | Given a bounding box, zoom range, and source crs, 176 | find all tiles that would intersect 177 | 178 | Parameters 179 | ----------- 180 | bbox: list 181 | [w, s, e, n] bounds 182 | src_crs: str 183 | the source crs of the input bbox 184 | minz: int 185 | minumum zoom to find tiles for 186 | maxz: int 187 | maximum zoom to find tiles for 188 | 189 | Returns 190 | -------- 191 | tiles: generator 192 | generator of [x, y, z] tiles that intersect 193 | the provided bounding box 194 | """ 195 | w, s, e, n = transform_bounds(*[src_crs, "EPSG:4326"] + bbox, densify_pts=0) 196 | 197 | EPSILON = 1.0e-10 198 | 199 | w += EPSILON 200 | s += EPSILON 201 | e -= EPSILON 202 | n -= EPSILON 203 | 204 | for z in range(minz, maxz + 1): 205 | for x, y in _tile_range(mercantile.tile(w, n, z), mercantile.tile(e, s, z)): 206 | yield [x, y, z] 207 | 208 | 209 | class RGBTiler: 210 | """ 211 | Takes continous source data of an arbitrary bit depth and encodes it 212 | in parallel into RGB tiles in an MBTiles file. Provided with a context manager: 213 | ``` 214 | with RGBTiler(inpath, outpath, min_z, max_x, **kwargs) as tiler: 215 | tiler.run(processes) 216 | ``` 217 | 218 | Parameters 219 | ----------- 220 | inpath: string 221 | filepath of the source file to read and encode 222 | outpath: string 223 | filepath of the output `mbtiles` 224 | min_z: int 225 | minimum zoom level to tile 226 | max_z: int 227 | maximum zoom level to tile 228 | 229 | Keyword Arguments 230 | ------------------ 231 | baseval: float 232 | the base value of the RGB numbering system. 233 | (will be treated as zero for this encoding) 234 | Default=0 235 | interval: float 236 | the interval at which to encode 237 | Default=1 238 | round_digits: int 239 | Erased less significant digits 240 | Default=0 241 | format: str 242 | output tile image format (png or webp) 243 | Default=png 244 | bounding_tile: list 245 | [x, y, z] of bounding tile; limits tiled output to this extent 246 | 247 | Returns 248 | -------- 249 | None 250 | 251 | """ 252 | 253 | def __init__( 254 | self, 255 | inpath, 256 | outpath, 257 | min_z, 258 | max_z, 259 | interval=1, 260 | base_val=0, 261 | round_digits=0, 262 | bounding_tile=None, 263 | **kwargs 264 | ): 265 | self.run_function = _tile_worker 266 | self.inpath = inpath 267 | self.outpath = outpath 268 | self.min_z = min_z 269 | self.max_z = max_z 270 | self.bounding_tile = bounding_tile 271 | 272 | if not "format" in kwargs: 273 | writer_func = _encode_as_png 274 | self.image_format = "png" 275 | elif kwargs["format"].lower() == "png": 276 | writer_func = _encode_as_png 277 | self.image_format = "png" 278 | elif kwargs["format"].lower() == "webp": 279 | writer_func = _encode_as_webp 280 | self.image_format = "webp" 281 | else: 282 | raise ValueError( 283 | "{0} is not a supported filetype!".format(kwargs["format"]) 284 | ) 285 | 286 | # global kwargs not used if output is webp 287 | self.global_args = { 288 | "kwargs": { 289 | "driver": "PNG", 290 | "dtype": "uint8", 291 | "height": 512, 292 | "width": 512, 293 | "count": 3, 294 | "crs": "EPSG:3857", 295 | }, 296 | "base_val": base_val, 297 | "interval": interval, 298 | "round_digits": round_digits, 299 | "writer_func": writer_func, 300 | } 301 | 302 | def __enter__(self): 303 | return self 304 | 305 | def __exit__(self, ext_t, ext_v, trace): 306 | if ext_t: 307 | traceback.print_exc() 308 | 309 | def run(self, processes=4): 310 | """ 311 | Warp, encode, and tile 312 | """ 313 | 314 | # get the bounding box + crs of the file to tile 315 | with rasterio.open(self.inpath) as src: 316 | bbox = list(src.bounds) 317 | src_crs = src.crs 318 | 319 | # remove the output filepath if it exists 320 | if os.path.exists(self.outpath): 321 | os.unlink(self.outpath) 322 | 323 | # create a connection to the mbtiles file 324 | conn = sqlite3.connect(self.outpath) 325 | cur = conn.cursor() 326 | 327 | # create the tiles table 328 | cur.execute( 329 | "CREATE TABLE tiles " 330 | "(zoom_level integer, tile_column integer, " 331 | "tile_row integer, tile_data blob);" 332 | ) 333 | # create empty metadata 334 | cur.execute("CREATE TABLE metadata (name text, value text);") 335 | 336 | conn.commit() 337 | 338 | # populate metadata with required fields 339 | cur.execute( 340 | "INSERT INTO metadata " "(name, value) " "VALUES ('format', ?);", 341 | (self.image_format,), 342 | ) 343 | 344 | cur.execute("INSERT INTO metadata " "(name, value) " "VALUES ('name', '');") 345 | cur.execute( 346 | "INSERT INTO metadata " "(name, value) " "VALUES ('description', '');" 347 | ) 348 | cur.execute("INSERT INTO metadata " "(name, value) " "VALUES ('version', '1');") 349 | cur.execute( 350 | "INSERT INTO metadata " "(name, value) " "VALUES ('type', 'baselayer');" 351 | ) 352 | 353 | conn.commit() 354 | 355 | if processes == 1: 356 | # use mock pool for profiling / debugging 357 | self.pool = MockTub( 358 | _main_worker, (self.inpath, self.run_function, self.global_args) 359 | ) 360 | else: 361 | self.pool = Pool( 362 | processes, 363 | _main_worker, 364 | (self.inpath, self.run_function, self.global_args), 365 | ) 366 | 367 | # generator of tiles to make 368 | if self.bounding_tile is None: 369 | tiles = _make_tiles(bbox, src_crs, self.min_z, self.max_z) 370 | else: 371 | constrained_bbox = list(mercantile.bounds(self.bounding_tile)) 372 | tiles = _make_tiles(constrained_bbox, "EPSG:4326", self.min_z, self.max_z) 373 | 374 | for tile, contents in self.pool.imap_unordered(self.run_function, tiles): 375 | x, y, z = tile 376 | 377 | # mbtiles use inverse y indexing 378 | tiley = int(math.pow(2, z)) - y - 1 379 | 380 | # insert tile object 381 | cur.execute( 382 | "INSERT INTO tiles " 383 | "(zoom_level, tile_column, tile_row, tile_data) " 384 | "VALUES (?, ?, ?, ?);", 385 | (z, x, tiley, buffer(contents)), 386 | ) 387 | 388 | conn.commit() 389 | 390 | conn.close() 391 | 392 | self.pool.close() 393 | self.pool.join() 394 | 395 | return None 396 | --------------------------------------------------------------------------------