-
Notifications
You must be signed in to change notification settings - Fork 23
/
worker.py
172 lines (138 loc) · 5.07 KB
/
worker.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
#!/usr/bin/env python
from __future__ import division
import click
import numpy as np
import rasterio
import riomucho
from rio_pansharpen.methods import Brovey
from rasterio.transform import guard_transform
from . utils import (
_pad_window, _upsample, _calc_windows, _check_crs,
_create_apply_mask, _half_window, _rescale)
def pansharpen(vis, vis_transform, pan, pan_transform,
pan_dtype, r_crs, dst_crs, weight,
method="Brovey", src_nodata=0):
"""Pansharpen a lower-resolution visual band
Parameters
=========
vis: ndarray, 3D with shape == (3, vh, vw)
Visual band array with RGB bands
vis_transform: Affine
affine transform defining the georeferencing of the vis array
pan: ndarray, 2D with shape == (ph, pw)
Panchromatic band array
pan_transform: Affine
affine transform defining the georeferencing of the pan array
method: string
Algorithm for pansharpening; default Brovey
Returns:
======
pansharp: ndarray, 3D with shape == (3, ph, pw)
pansharpened visual band
affine transform is identical to `pan_transform`
"""
rgb = _upsample(_create_apply_mask(vis), pan.shape, vis_transform, r_crs,
pan_transform, dst_crs)
# Main Pansharpening Processing
if method == "Brovey":
pansharp, _ = Brovey(rgb, pan, weight, pan_dtype)
# TODO: add other methods
return pansharp
def _pansharpen_worker(open_files, pan_window, _, g_args):
"""rio mucho worker for pansharpening. It reads input
files and performing pansharpening on each window.
Parameters
------------
open_files: list of rasterio open files
pan_window: tuples
g_args: dictionary
Returns
---------
out: None
Output is written to dst_path
"""
pan = open_files[0].read(1, window=pan_window).astype(np.float32)
pan_dtype = open_files[0].meta['dtype']
# Get the rgb window that covers the pan window
if g_args.get("half_window"):
rgb_window = _half_window(pan_window)
else:
padding = 2
pan_bounds = open_files[0].window_bounds(pan_window)
rgb_base_window = open_files[1].window(*pan_bounds)
rgb_window = _pad_window(rgb_base_window, padding)
# Determine affines for those windows
pan_affine = open_files[0].window_transform(pan_window)
rgb_affine = open_files[1].window_transform(rgb_window)
rgb = riomucho.utils.array_stack(
[src.read(window=rgb_window, boundless=True).astype(np.float32)
for src in open_files[1:]])
if g_args["verb"]:
click.echo('pan shape: %s, rgb shape %s' % (pan.shape, rgb.shape))
pansharpened = pansharpen(
rgb, rgb_affine, pan, pan_affine, pan_dtype,
g_args["r_crs"], g_args["dst_crs"],
g_args["weight"], method="Brovey")
pan_rescale = _rescale(
pansharpened, g_args["src_nodata"], g_args["dst_dtype"],
out_alpha=g_args.get("out_alpha", True))
return pan_rescale
def calculate_landsat_pansharpen(src_paths, dst_path, dst_dtype,
weight, verbosity, jobs, half_window,
customwindow, out_alpha, creation_opts):
"""Parameters
------------
src_paths: list of string (pan_path, r_path, g_path, b_path)
dst_path: string
dst_dtype: 'uint16', 'uint8'.
weight: float
jobs: integer
half_window: boolean
customwindow: integer
out_alpha: boolean
output an alpha band?
creation_opts: dict
creation options to update the write profile
Returns
---------
out: None
Output is written to dst_path
"""
with rasterio.open(src_paths[0]) as pan_src:
windows = _calc_windows(pan_src, customwindow)
profile = pan_src.profile
if profile['count'] > 1:
raise RuntimeError(
"Pan band must be 1 band - is {}".format(profile['count']))
dst_dtype = np.__dict__[dst_dtype]
profile.update(
transform=guard_transform(pan_src.transform),
dtype=dst_dtype,
count=3,
photometric='rgb')
if out_alpha:
profile['count'] = 4
if creation_opts:
profile.update(**creation_opts)
with rasterio.open(src_paths[1]) as r_src:
r_meta = r_src.meta
if profile['width'] <= r_meta['width'] or \
profile['height'] <= r_meta['height']:
raise RuntimeError(
"Pan band must be larger than RGB bands")
_check_crs([r_meta, profile])
g_args = {
"verb": verbosity,
"half_window": half_window,
"dst_dtype": dst_dtype,
"out_alpha": out_alpha,
"weight": weight,
"dst_aff": guard_transform(profile['transform']),
"dst_crs": profile['crs'],
"r_aff": guard_transform(r_meta['transform']),
"r_crs": r_meta['crs'],
"src_nodata": 0}
with riomucho.RioMucho(src_paths, dst_path, _pansharpen_worker,
windows=windows, global_args=g_args,
options=profile, mode='manual_read') as rm:
rm.run(jobs)