-
Notifications
You must be signed in to change notification settings - Fork 3
/
resample_mask.py
150 lines (132 loc) · 5.32 KB
/
resample_mask.py
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
"""Resample a mask with custom rules."""
import argparse
import logging
import math
import os
import shutil
import sys
from osgeo import gdal
from osgeo import osr
from ecoshard import taskgraph
import ecoshard.geoprocessing
import numpy
gdal.SetCacheMax(2**27)
logging.basicConfig(
level=logging.DEBUG,
format=(
'%(asctime)s (%(relativeCreated)d) %(levelname)s %(name)s'
' [%(funcName)s:%(lineno)d] %(message)s'),
stream=sys.stdout)
LOGGER = logging.getLogger(__name__)
logging.getLogger('taskgraph').setLevel(logging.WARN)
def threshold_mask(prop_array, threshold):
result = numpy.zeros(prop_array.shape, dtype=bool)
result[prop_array >= threshold] = 1
return result
def main():
"""Entry point."""
parser = argparse.ArgumentParser(description='Resample mask')
parser.add_argument(
'mask_raster_path',
help='path to mask raster')
parser.add_argument(
'target_cell_size',
type=float, help='new cell size in raster units')
parser.add_argument(
'pixel_proportion_to_mask',
type=float, help='what percent of pixels need to be masked to turn on')
parser.add_argument(
'target_raster_path', type=str,
help='path to resampled target raster')
parser.add_argument(
'--reproject_file', nargs=6, type=str, help=(
'6 arguments: {path to projection wkt} '
'{pixel size in projected units} '
'bounding box {xmin}, {ymin} {xmax} {ymax}'))
args = parser.parse_args()
WORKSPACE_DIR = 'tmp_resample_workspace'
os.makedirs(WORKSPACE_DIR, exist_ok=True)
task_graph = taskgraph.TaskGraph(WORKSPACE_DIR, -1)
mask_info = ecoshard.geoprocessing.get_raster_info(args.mask_raster_path)
mask_basename = os.path.basename(
os.path.splitext(args.mask_raster_path)[0])
# calculate new raster size
n_cols, n_rows = mask_info['raster_size']
change_ratio = args.target_cell_size / mask_info['pixel_size'][0]
if change_ratio > 1:
raise ValueError(
f'the selection of a new target cell size of '
f'{args.target_cell_size} will result in making the cells '
f'{change_ratio} times bigger.')
LOGGER.info(f'raster size is scaled by {change_ratio}')
LOGGER.info(
f"original raster is {mask_info['pixel_size']}, "
f'new raster is {n_cols/change_ratio}, {n_rows/change_ratio}')
ones_array = numpy.ones(
(int(math.ceil(change_ratio)), int(math.ceil(change_ratio))),
dtype=numpy.float32)
ones_array[:] /= ones_array.size
kernel_raster_path = os.path.join(
WORKSPACE_DIR, f'tmp_kernel_{ones_array.shape}.tif')
task_graph.add_task(
func=ecoshard.geoprocessing.numpy_array_to_raster,
args=(
ones_array, None, (1, -1), (0, 0), osr.SRS_WKT_WGS84_LAT_LONG,
kernel_raster_path),
target_path_list=[kernel_raster_path])
sampled_raster_path = os.path.join(
WORKSPACE_DIR, f'tmp_sampled_{mask_basename}.tif')
task_graph.add_task(
func=ecoshard.geoprocessing.convolve_2d,
args=(
(args.mask_raster_path, 1),
(kernel_raster_path, 1), sampled_raster_path),
kwargs={'mask_nodata': False},
target_path_list=[sampled_raster_path],
task_name='convolve_2d')
threshold_raster_path = os.path.join(
WORKSPACE_DIR,
f'tmp_threshold_{mask_basename}_{args.pixel_proportion_to_mask}.tif')
task_graph.add_task(
func=ecoshard.geoprocessing.raster_calculator,
args=(
[(sampled_raster_path, 1), (args.pixel_proportion_to_mask, 'raw')],
threshold_mask, threshold_raster_path,
gdal.GDT_Byte, None),
target_path_list=[threshold_raster_path],
task_name='threshold')
task_graph.add_task(
func=ecoshard.geoprocessing.warp_raster,
args=(
threshold_raster_path,
(args.target_cell_size, -args.target_cell_size),
args.target_raster_path, 'near'),
target_path_list=[args.target_raster_path],
task_name=f'warp to {args.target_cell_size}')
task_graph.join()
if args.reproject_file:
LOGGER.info(f'reproject to {args.reproject_file}')
unprojected_path = os.path.join(
WORKSPACE_DIR, os.path.basename(args.target_raster_path))
shutil.copy(args.target_raster_path, unprojected_path)
with open(args.reproject_file[0], 'r') as reproject_file:
target_projection_wkt = reproject_file.read()
# target pixel size is specified by the user
target_pixel_size = [
float(args.reproject_file[1]), -float(args.reproject_file[1])]
target_bounding_box = [float(x) for x in args.reproject_file[2:]]
LOGGER.debug(
f'warp raster to new projection with bounding box '
f'{target_bounding_box}')
task_graph.add_task(
func=ecoshard.geoprocessing.warp_raster,
args=(
unprojected_path, target_pixel_size, args.target_raster_path,
'near'),
kwargs={
'target_projection_wkt': target_projection_wkt,
'target_bb': target_bounding_box},
target_path_list=[args.target_raster_path],
task_name=f'reproject {args.target_raster_path}')
if __name__ == '__main__':
main()