#!/usr/bin/env python3
# encoding: utf-8
# Copyright (C) 2016-2021 Space Science and Engineering Center (SSEC),
# University of Wisconsin-Madison.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
# This file is part of the polar2grid software package. Polar2grid takes
# satellite observation data, remaps it, and writes it to a file format for
# input into another program.
# Documentation: http://www.ssec.wisc.edu/software/polar2grid/
#
# Written by David Hoese July 2016
# University of Wisconsin-Madison
# Space Science and Engineering Center
# 1225 West Dayton Street
# Madison, WI 53706
# david.hoese@ssec.wisc.edu
"""Script to add coastlines and borders to a geotiff while also creating a PNG."""
from __future__ import annotations
import argparse
import logging
import os
import sys
from pathlib import Path
import numpy as np
import rasterio
from aggdraw import Font
from PIL import Image, ImageFont
from pycoast import ContourWriterAGG
from pydecorate import DecoratorAGG
from pyresample.utils import get_area_def_from_raster
from trollimage.colormap import Colormap
from polar2grid.utils.config import add_polar2grid_config_paths
LOG = logging.getLogger(__name__)
PYCOAST_DIR = os.environ.get("GSHHS_DATA_ROOT")
[docs]
def _get_rio_colormap(rio_ds, bidx):
try:
return rio_ds.colormap(bidx)
except ValueError:
return None
[docs]
def _get_colorbar_vmin_vmax(arg_min, arg_max, metadata, input_dtype):
scale = metadata.get("scale", metadata.get("scale_factor"))
offset = metadata.get("offset", metadata.get("add_offset"))
dtype_min = float(np.iinfo(input_dtype).min)
dtype_max = float(np.iinfo(input_dtype).max)
if arg_min is None and scale is None:
LOG.warning(
"Colorbar min/max metadata not found and not provided "
"on the command line. Defaulting to data type limits."
)
return dtype_min, dtype_max
if arg_min is not None:
vmin = float(arg_min)
vmax = float(arg_max)
else:
scale, offset = _convert_scale_offset_to_float(scale, offset)
vmin, vmax = _vmin_vmax_from_scale_offset(scale, offset, dtype_min, dtype_max)
return vmin, vmax
[docs]
def _convert_scale_offset_to_float(scale: str, offset: str) -> tuple[float, float]:
scale = float(scale)
offset = float(offset)
if np.isnan(scale) or np.isnan(offset):
raise ValueError(
"Can't automatically set colorbar limits with "
"geotiff metadata as scale/offset are set to "
"NaN. This indicates a non-linear enhancement or "
"RGB/A composite that was enhanced. These cases "
"cannot be represented properly by a colorbar."
)
return scale, offset
[docs]
def _vmin_vmax_from_scale_offset(scale: float, offset: float, dtype_min: int, dtype_max: int) -> tuple[float, float]:
delta = dtype_max - dtype_min
vmin = offset
vmax = delta * scale + offset
# floating point error made it not an integer
if abs(vmin - np.round(vmin, 0)) <= 0.001:
vmin = np.round(vmin, 0)
if abs(vmax - np.round(vmax, 0)) <= 0.001:
vmax = np.round(vmax, 0)
return vmin, vmax
[docs]
def _apply_decorator_alignment(dc, align):
if align == "top":
dc.align_top()
dc.write_horizontally()
elif align == "bottom":
dc.align_bottom()
dc.write_horizontally()
elif align == "left":
dc.align_left()
dc.write_vertically()
elif align == "right":
dc.align_right()
dc.write_vertically()
[docs]
def _add_colorbar_to_image(img, font, tick_color, align, **kwargs):
dc = DecoratorAGG(img)
_apply_decorator_alignment(dc, align)
dc.add_scale(
font=font,
line=tick_color,
**kwargs,
)
[docs]
def get_parser():
parser = argparse.ArgumentParser(
description="Add overlays to a GeoTIFF file and save as a PNG file.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
_add_coastlines_arguments(parser)
_add_rivers_arguments(parser)
_add_grid_arguments(parser)
_add_borders_arguments(parser)
_add_colorbar_arguments(parser)
_add_global_arguments(parser)
parser.add_argument("input_tiff", nargs="+", help="Input geotiff(s) to process")
return parser
[docs]
def _add_coastlines_arguments(parser: argparse.ArgumentParser) -> None:
group = parser.add_argument_group("coastlines")
group.add_argument("--add-coastlines", action="store_true", help="Add coastlines")
group.add_argument(
"--coastlines-resolution",
choices="clihf",
default="i",
help="Resolution of coastlines to add (crude, low, intermediate, high, full)",
)
group.add_argument(
"--coastlines-level",
choices=range(1, 7),
type=int,
default=4,
help="Level of detail from the selected resolution dataset",
)
group.add_argument(
"--coastlines-outline",
default=["yellow"],
nargs="*",
help="Color of coastline lines (color name or 3 RGB integers)",
)
group.add_argument("--coastlines-fill", default=None, nargs="*", help="Color of land")
group.add_argument("--coastlines-width", default=1.0, type=float, help="Width of coastline lines")
[docs]
def _add_rivers_arguments(parser: argparse.ArgumentParser) -> None:
group = parser.add_argument_group("rivers")
group.add_argument("--add-rivers", action="store_true", help="Add rivers grid")
group.add_argument(
"--rivers-resolution",
choices="clihf",
default="c",
help="Resolution of rivers to add (crude, low, intermediate, high, full)",
)
group.add_argument(
"--rivers-level", choices=range(0, 11), type=int, default=5, help="Level of detail for river lines"
)
group.add_argument(
"--rivers-outline", default=["blue"], nargs="*", help="Color of river lines (color name or 3 RGB integers)"
)
group.add_argument("--rivers-width", default=1.0, type=float, help="Width of rivers lines")
[docs]
def _add_grid_arguments(parser: argparse.ArgumentParser) -> None:
group = parser.add_argument_group("grid")
group.add_argument("--add-grid", action="store_true", help="Add lat/lon grid")
group.add_argument("--grid-no-text", dest="grid_text", action="store_false", help="Add labels to lat/lon grid")
group.add_argument("--grid-text-size", default=32, type=int, help="Lat/lon grid text font size")
group.add_argument("--grid-font", default="Vera.ttf", help="Path to TTF font (package provided or custom path)")
group.add_argument(
"--grid-fill", nargs="*", default=["cyan"], help="Color of grid text (color name or 3 RGB integers)"
)
group.add_argument(
"--grid-outline", nargs="*", default=["cyan"], help="Color of grid lines (color name or 3 RGB integers)"
)
group.add_argument(
"--grid-minor-outline", nargs="*", default=["cyan"], help="Color of tick lines (color name or 3 RGB integers)"
)
group.add_argument(
"--grid-D", nargs=2, default=(10.0, 10.0), type=float, help="Degrees between grid lines (lon, lat)"
)
group.add_argument(
"--grid-d", nargs=2, default=(2.0, 2.0), type=float, help="Degrees between tick lines (lon, lat)"
)
group.add_argument(
"--grid-lon-placement", choices=["tl", "lr", "lc", "cc"], default="tb", help="Longitude label placement"
)
group.add_argument(
"--grid-lat-placement", choices=["tl", "lr", "lc", "cc"], default="lr", help="Latitude label placement"
)
group.add_argument("--grid-width", default=1.0, type=float, help="Width of grid lines")
[docs]
def _add_borders_arguments(parser: argparse.ArgumentParser) -> None:
group = parser.add_argument_group("borders")
group.add_argument("--add-borders", action="store_true", help="Add country and/or region borders")
group.add_argument(
"--borders-resolution",
choices="clihf",
default="i",
help="Resolution of borders to add (crude, low, intermediate, high, full)",
)
group.add_argument(
"--borders-level", choices=range(1, 4), default=2, type=int, help="Level of detail for border lines"
)
group.add_argument(
"--borders-outline", default=["white"], nargs="*", help="Color of border lines (color name or 3 RGB integers)"
)
group.add_argument("--borders-width", default=1.0, type=float, help="Width of border lines")
[docs]
def _add_colorbar_arguments(parser: argparse.ArgumentParser) -> None:
group = parser.add_argument_group("colorbar")
group.add_argument("--add-colorbar", action="store_true", help="Add colorbar on top of image")
group.add_argument(
"--colorbar-colormap-file",
help=argparse.SUPPRESS,
# help="Specify the colormap file that was used to "
# "colorize the provided RGB geotiff. Only used if "
# "the provided geotiff is RGB/A. Otherwise the "
# "geotiff is expected to include the colormap as "
# "a geotiff color table.",
)
group.add_argument("--colorbar-width", type=int, help="Number of pixels wide")
group.add_argument("--colorbar-height", type=int, help="Number of pixels high")
group.add_argument(
"--colorbar-extend", action="store_true", help="Extend colorbar to full width/height of the image"
)
group.add_argument(
"--colorbar-tick-marks", type=float, default=5.0, help="Major tick and tick label interval in data units"
)
group.add_argument("--colorbar-minor-tick-marks", type=float, default=1.0, help="Minor tick interval in data units")
group.add_argument("--colorbar-text-size", default=32, type=int, help="Tick label font size")
group.add_argument(
"--colorbar-text-color", nargs="*", default=["black"], help="Color of tick text (color name or 3 RGB integers)"
)
group.add_argument("--colorbar-font", default="Vera.ttf", help="Path to TTF font (package provided or custom path)")
group.add_argument(
"--colorbar-align",
choices=["left", "top", "right", "bottom"],
default="bottom",
help="Which side of the image to place the colorbar",
)
group.add_argument("--colorbar-vertical", action="store_true", help="DEPRECATED")
group.add_argument(
"--colorbar-min",
type=float,
help="Minimum data value of the colorbar."
" Defaults to 'min_in' of input metadata or"
" minimum value of the data otherwise.",
)
group.add_argument(
"--colorbar-max",
type=float,
help="Maximum data value of the colorbar."
" Defaults to 'max_in' of input metadata or"
" maximum value of the data otherwise.",
)
group.add_argument("--colorbar-units", help="Units marker to include in the colorbar text")
group.add_argument("--colorbar-title", help="Title shown with the colorbar")
[docs]
def _add_global_arguments(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--shapes-dir",
default=PYCOAST_DIR,
help="Specify alternative directory for coastline shape files (default: GSHSS_DATA_ROOT)",
)
parser.add_argument(
"--cache-dir",
default=None,
help="Specify directory where cached coastline output can be stored and accessed in later "
"executions. The cache will never be cleared by this script. Caching depends on the grid "
"of the image and the decorations added to the image.",
)
parser.add_argument(
"--cache-regenerate",
action="store_true",
help="Force regeneration of any cached overlays. Requires '--cache-dir'.",
)
parser.add_argument(
"-o",
"--output",
dest="output_filename",
nargs="+",
help="Specify the output filename (default replace '.tif' with '.png')",
)
parser.add_argument(
"-v",
"--verbose",
dest="verbosity",
action="count",
default=0,
help="each occurrence increases verbosity 1 level through ERROR-WARNING-INFO-DEBUG (default INFO)",
)
[docs]
def main(argv=sys.argv[1:]):
parser = get_parser()
args = parser.parse_args(argv)
levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]
logging.basicConfig(level=levels[min(3, args.verbosity)])
add_polar2grid_config_paths()
if args.output_filename is None:
args.output_filename = [x[:-3] + "png" for x in args.input_tiff]
elif len(args.output_filename) != len(args.input_tiff):
LOG.error("Output filenames must be equal to number of input tiffs")
return -1
if not (args.add_borders or args.add_coastlines or args.add_grid or args.add_rivers or args.add_colorbar):
LOG.error("Please specify one of the '--add-X' options to modify the image")
return -1
if args.cache_dir and not os.path.isdir(args.cache_dir):
LOG.info(f"Creating cache directory: {args.cache_dir}")
os.makedirs(args.cache_dir, exist_ok=True)
# we may be dealing with large images that look like decompression bombs
# let's turn off the check for the image size in PIL/Pillow
Image.MAX_IMAGE_PIXELS = None
# gather all options into a single dictionary that we can pass to pycoast
pycoast_options = _args_to_pycoast_dict(args)
colorbar_kwargs = _args_to_colorbar_kwargs(args) if args.add_colorbar else {}
for input_tiff, output_filename in zip(args.input_tiff, args.output_filename, strict=True):
_process_one_image(input_tiff, output_filename, pycoast_options, args.shapes_dir, colorbar_kwargs)
return 0
[docs]
def _args_to_pycoast_dict(args):
opts = {}
if args.add_coastlines:
opts["coasts"] = _args_to_coastlines_dict(args)
if args.add_rivers:
opts["rivers"] = _args_to_rivers_dict(args)
if args.add_borders:
opts["borders"] = _args_to_borders_dict(args)
if args.add_grid:
opts["grid"] = _args_to_grid_dict(args)
if args.cache_dir:
opts["cache"] = {
# add "add_coastlines" prefix to cached image name
"file": os.path.join(args.cache_dir, "add_coastlines"),
"regenerate": args.cache_regenerate,
}
return opts
[docs]
def _args_to_coastlines_dict(args):
outline = (
args.coastlines_outline[0]
if len(args.coastlines_outline) == 1
else tuple(int(x) for x in args.coastlines_outline)
)
if args.coastlines_fill:
fill = (
args.coastlines_fill[0] if len(args.coastlines_fill) == 1 else tuple(int(x) for x in args.coastlines_fill)
)
else:
fill = None
coasts_dict = {
"resolution": args.coastlines_resolution,
"level": args.coastlines_level,
"width": args.coastlines_width,
"outline": outline,
"fill": fill,
}
return coasts_dict
[docs]
def _args_to_rivers_dict(args):
outline = args.rivers_outline[0] if len(args.rivers_outline) == 1 else tuple(int(x) for x in args.rivers_outline)
rivers_dict = {
"resolution": args.rivers_resolution,
"level": args.rivers_level,
"width": args.rivers_width,
"outline": outline,
}
return rivers_dict
[docs]
def _args_to_borders_dict(args):
outline = args.borders_outline[0] if len(args.borders_outline) == 1 else tuple(int(x) for x in args.borders_outline)
borders_dict = {
"resolution": args.borders_resolution,
"level": args.borders_level,
"width": args.borders_width,
"outline": outline,
}
return borders_dict
[docs]
def _args_to_grid_dict(args):
outline = args.grid_outline[0] if len(args.grid_outline) == 1 else tuple(int(x) for x in args.grid_outline)
minor_outline = (
args.grid_minor_outline[0]
if len(args.grid_minor_outline) == 1
else tuple(int(x) for x in args.grid_minor_outline)
)
fill = args.grid_fill[0] if len(args.grid_fill) == 1 else tuple(int(x) for x in args.grid_fill)
font_path = find_font(args.grid_font, args.grid_text_size)
font = Font(outline, font_path, size=args.grid_text_size)
grid_dict = {
"lon_major": args.grid_D[0],
"lat_major": args.grid_D[1],
"lon_minor": args.grid_d[0],
"lat_minor": args.grid_d[1],
"font": font,
"fill": fill,
"outline": outline,
"minor_outline": minor_outline,
"write_text": args.grid_text,
"width": args.grid_width,
"lon_placement": args.grid_lon_placement,
"lat_placement": args.grid_lat_placement,
}
return grid_dict
[docs]
def _args_to_colorbar_kwargs(args):
font_color = args.colorbar_text_color
font_color = font_color[0] if len(font_color) == 1 else tuple(int(x) for x in font_color)
font_path = find_font(args.colorbar_font, args.colorbar_text_size)
# this actually needs an aggdraw font
font = Font(font_color, font_path, size=args.colorbar_text_size)
if args.colorbar_width is None or args.colorbar_height is None:
if args.colorbar_width is not None or args.colorbar_height is not None:
LOG.warning(
"'--colorbar-width' and '--colorbar-height' were not both specified. " "Forcing '--colorbar-extend'."
)
args.colorbar_extend = True
colorbar_kwargs = {
"font": font,
"tick_color": font_color,
"cmin": args.colorbar_min,
"cmax": args.colorbar_max,
"align": args.colorbar_align,
"extend": args.colorbar_extend,
"tick_marks": args.colorbar_tick_marks,
"minor_tick_marks": args.colorbar_minor_tick_marks,
"title": args.colorbar_title,
"unit": args.colorbar_units,
}
if args.colorbar_width:
colorbar_kwargs["width"] = args.colorbar_width
if args.colorbar_height:
colorbar_kwargs["height"] = args.colorbar_height
return colorbar_kwargs
[docs]
def find_font(font_name, size):
try:
font = ImageFont.truetype(font_name, size)
return font.path
except IOError as err:
font_path = Path(__file__).parent / "fonts" / font_name
if not font_path.is_file():
raise ValueError(f"Font path does not exist: {font_path}") from err
return str(font_path)
[docs]
def _process_one_image(
input_tiff: str, output_filename: str, pycoast_options: dict, shapes_dir: str, colorbar_kwargs: dict
) -> None:
LOG.info("Creating {} from {}".format(output_filename, input_tiff))
img = Image.open(input_tiff)
all_tiff_tags = img.tag_v2
img_bands = img.getbands()
num_bands = len(img_bands)
# P = palette which we assume to be an RGBA colormap
img = img.convert("RGBA" if num_bands in (2, 4) or "P" in img_bands else "RGB")
if pycoast_options:
area_id = os.path.splitext(input_tiff[0])[0]
area_def = get_area_def_from_raster(input_tiff, area_id=area_id)
cw = ContourWriterAGG(shapes_dir)
cw.add_overlay_from_dict(pycoast_options, area_def, background=img)
if colorbar_kwargs:
colorbar_kwargs = colorbar_kwargs.copy()
cmin = colorbar_kwargs.pop("cmin")
cmax = colorbar_kwargs.pop("cmax")
cmap = _get_colormap_object(input_tiff, num_bands, cmin, cmax)
_add_colorbar_to_image(img, colormap=cmap, **colorbar_kwargs)
kwargs = {}
if output_filename.endswith(".tif") or output_filename.endswith(".tiff"):
# Copy geotiff/tiff tags if output is TIFF
geotiff_tags = {tag_num: tag_val for tag_num, tag_val in all_tiff_tags.items() if tag_num > 30000}
kwargs = {"tiffinfo": geotiff_tags}
img.save(output_filename, **kwargs)
[docs]
def _get_colormap_object(input_tiff, num_bands, cmin, cmax):
rio_ds = rasterio.open(input_tiff)
input_dtype = np.dtype(rio_ds.meta["dtype"])
colormap_csv = rio_ds.tags().get("colormap")
rio_ct = _get_rio_colormap(rio_ds, 1)
metadata = rio_ds.tags()
cmap = _convert_table_to_cmap_or_default_bw(input_dtype, rio_ct, num_bands)
if num_bands in (3, 4) and colormap_csv is None:
raise ValueError("RGB and RGBA geotiffs must have a colormap " "specified with '--colorbar-colormap-file'.")
if num_bands in (3, 4) or colormap_csv is not None:
cmap = Colormap.from_string(colormap_csv)
vmin, vmax = _get_colorbar_vmin_vmax(cmin, cmax, metadata, input_dtype)
cmap = cmap.set_range(vmin, vmax, inplace=False)
return cmap
[docs]
def _convert_table_to_cmap_or_default_bw(band_dtype, band_ct, band_count):
max_val = np.iinfo(band_dtype).max
# if we have an alpha band then include the entire colormap
# otherwise assume it is using 0 as a fill value
start_idx = 1 if band_count == 1 else 0
if band_ct is None:
# all black colormap
# don't assume anything about the colors or scaling of the image if the
# user didn't provide enough information
color_iter = ((idx / float(max_val), (0.0, 0.0, 0.0, 1.0)) for idx in range(max_val))
color_iter = list(color_iter)
else:
color_iter = ((idx / float(max_val), color) for idx, color in sorted(band_ct.items())[start_idx:])
color_iter = ((idx, tuple(x / float(max_val) for x in color)) for idx, color in color_iter)
cmap = Colormap(*color_iter)
return cmap
if __name__ == "__main__":
sys.exit(main())
