-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathbother
executable file
·192 lines (167 loc) · 9.57 KB
/
bother
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""A script to produce heightmaps (primarily for OpenTTD) using real-world
elevation data.
"""
import sys
import logging
import math
import time
import argparse
import os
import os.path
import tempfile
from typing import Optional, Set, Tuple, List
from PIL import Image
from rasterio.io import MemoryFile
from bother_utils.srtm import create_tif_file, clear_cache
from bother_utils.heightmap import (remove_sea, resample, reproject_raster, set_lakes_to_elev, raise_undersea_land,
raise_low_pixels, to_png, crop_modes, crop_image, scale_image, png_to_file)
# EPSG codes
WGS84 = 4326 # Mercator - The default CRS used in the STRM data
PSEUDO_MERCATOR = 3857 # Web Mercator - The projection used by Google Maps, OpenStreetMap, etc.
def error(msg: str):
print(f'ERROR: {msg}', file=sys.stderr)
sys.exit(1)
def check_namespace(ns: argparse.Namespace):
"""Check for errors in the arguments passed."""
if (ns.bounds is None) and (ns.infile_tif is None) and (ns.infile_png is None):
error('Must pass --bounds, --infile-tif or --infile-png.')
elif sum(((ns.bounds is not None), (ns.infile_tif is not None), (ns.infile_png is not None))) > 1:
error('--bounds, --infile-tif and --infile-png are mutually exclusive.')
if (ns.scale_data is not None) and ns.scale_data == 0:
error('0 is invalid value for scaling.')
if ns.crop:
res, mode = ns.crop
if mode.lower() not in crop_modes:
error(f'Mode must be one of {crop_modes}.')
res = res.split('x')
try:
width = int(res[0])
height = int(res[1])
except (IndexError, ValueError):
error('Size for cropped image must be in form "WIDTHxHEIGHT", where WIDTH and HEIGHT are integers.')
if (width <= 0) or (height <= 0):
error(f'Invalid dimensions for cropping: {width}x{height}.')
if ns.scale_image:
res = ns.scale_image.split('x')
try:
width = int(res[0])
height = int(res[1])
except (IndexError, ValueError):
error('Size for scaled image must be in form "WIDTHxHEIGHT", where WIDTH and HEIGHT are integers.')
if (width <= 0) or (height <= 0):
error(f'Invalid dimensions for scaling: {width}x{height}.')
def parse_namespace(ns: argparse.Namespace):
"""Parse the arguments passed and execute request."""
tmp_file = None
if ns.bounds:
if ns.outfile_tif:
to_file = os.path.abspath(ns.outfile_tif)
else:
to_file = os.path.join(tempfile.gettempdir(), f'othg_{time.time()}.tif')
tmp_file = to_file
lat1, lon1, lat2, lon2 = ns.bounds
tif_file = create_tif_file(lon1, lat1, lon2, lat2, to_file)
else:
tif_file = ns.infile_tif
if tif_file:
with open(tif_file, 'rb') as f:
#memfile = handle_nodata(MemoryFile(f))
memfile = MemoryFile(f)
if ns.scale_data is not None:
memfile = resample(memfile, ns.scale_data)
if ns.no_sea:
memfile = remove_sea(memfile)
if ns.epsg and (ns.epsg != WGS84): # The SRTM data already uses WGS84 so no need to reproject to that
memfile = reproject_raster(memfile, dst_crs=f'EPSG:{ns.epsg}')
if ns.lakes:
memfile = set_lakes_to_elev(memfile, ns.lakes)
if ns.raise_undersea is not None:
memfile = raise_undersea_land(memfile, ns.raise_undersea)
if ns.raise_low is not None:
memfile = raise_low_pixels(memfile, ns.raise_low, ns.max_brightness)
im = to_png(memfile, not ns.raise_low, ns.max_brightness)
elif ns.infile_png:
im = Image.open(ns.infile_png)
if ns.crop:
res, mode = ns.crop
res = res.split('x')
width = int(res[0])
height = int(res[1])
im = crop_image(im, width, height, mode)
if ns.scale_image:
res = ns.scale_image.split('x')
width = int(res[0])
height = int(res[1])
im = scale_image(im, width, height)
if ns.outfile.endswith('.png'):
save_to = ns.outfile
else:
save_to = ns.outfile + '.png'
try:
png_to_file(im, save_to)
except FileNotFoundError:
error(f'Could not save to {save_to}. Check that the directory to which you want to save exists.')
if tmp_file:
os.remove(tmp_file)
if ns.clear_cache:
clear_cache()
def get_arg_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description='Generate heightmaps from NASA\'s Shuttle Radar Topography Mission '
'elevation data, primarily for use in OpenTTD.')
parser.add_argument('-it', '--infile-tif', dest='infile_tif', metavar='FILE',
help='Specify a TIF file to use to generate the heightmap, rather than fetching SRTM data from '
'the internet.')
parser.add_argument('-b', '--bounds', type=float, nargs=4, metavar=('BOTTOM', 'LEFT', 'TOP', 'RIGHT'),
help='Specify the bounds (coordinates of bottom-left and top-right points) for which to '
'download the SRTM data.')
parser.add_argument('-ot', '--outfile-tif', dest='outfile_tif', metavar='FILE',
help='Save the generated TIF file (before processing) so that it can be re-used.')
parser.add_argument('-sd', '--scale-data', type=float, dest='scale_data', metavar='FACTOR',
help='Factor by which to scale the data prior to converting to a PNG. Values lower than one '
'will downsample the data; values greater than one will upsample the data. Bilinear '
'interpolation is used.')
parser.add_argument('-e', '--epsg', type=int, metavar='CODE', default=PSEUDO_MERCATOR,
help='The EPSG code of the projection to use for the image. Default is 3857, the '
'"Pseudo-Mercator" projection commonly used by web applications such as Google Maps '
'and OpenStreetMap.')
parser.add_argument('-l', '--lakes', type=int, nargs='?', const=80, metavar='SIZE',
help='Detect lakes (as contiguous regions of a minimum size with the exact same elevation) '
'and set their elevation to zero (so they are rendered as water in OpenTTD). If provided, '
'the integer argument determines the minimum number of pixels such an area must contain '
'in order to be considered a lake (default is 80).')
parser.add_argument('-ru', '--raise-undersea', type=int, nargs='?', const=1, metavar='ELEVATION',
dest='raise_undersea', help='Raise pixels with an elevation below zero (ie, land that is below '
'sea level) to the given level.')
parser.add_argument('-rl', '--raise-low', type=float, nargs='?', const=0.0, metavar='ELEVATION', dest='raise_low',
help='Raise pixels with low elevation so that they have a non-zero value in the resulting '
'greyscale image. If a numerical argument is provided, only pixels with elevations '
'above that value will be raised (default is 0).')
parser.add_argument('-ns', '--no-sea', dest='no_sea', action='store_true',
help='Increase (or decrease) all elevations so that the lowest elevation is just above sea '
'level. This is helpful when your real world data contains land that is below sea level '
'but the map is entirely inland, so that there is no actual sea.')
parser.add_argument('-mb', '--max-brightness', type=int, default=255, dest='max_brightness',
help='Set the maximum brightness in the greyscale PNG (ie, the brightness of the highest point '
'in the data). Should be between 1 and 255; a lower value will lead to a flatter map in '
'OpenTTD.')
parser.add_argument('-ip', '--infile-png', dest='infile_png', metavar='FILE',
help='Load FILE to perform cropping and/or scaling, rather than generating a new PNG file from '
'SRTM data.')
parser.add_argument('-c', '--crop', nargs=2, metavar=('WIDTHxHEIGHT', 'MODE'),
help='Crop the resulting image to WIDTH x HEIGHT. MODE determines which region of the image to '
'crop to and must be one of nw, n, ne, e, c, w, sw, s, se. Note that you may prefer to '
'do the cropping and scaling in your favourite image editor.')
parser.add_argument('-si', '--scale-image', dest='scale_image', metavar='WIDTHxHEIGHT',
help='Scale the resulting image to WIDTH x HEIGHT. Note that you may prefer to do the cropping '
'and scaling in your favourite image editor.')
parser.add_argument('-cc', '--clear-cache', dest='clear_cache', action='store_true',
help='Clear cached SRTM data.')
parser.add_argument('outfile', help='The file to which the greyscale PNG image will be written.')
return parser
if __name__ == '__main__':
parser = get_arg_parser()
ns = parser.parse_args()
check_namespace(ns)
parse_namespace(ns)