Permalink
Cannot retrieve contributors at this time
import os | |
import random | |
from functools import partial | |
from itertools import product | |
from collections import namedtuple | |
from collections.abc import Container | |
from gbdxtools.rda.io import to_geotiff | |
from gbdxtools.rda.util import RatPolyTransform, AffineTransform, pad_safe_positive, pad_safe_negative, RDA_TO_DTYPE, get_proj | |
from gbdxtools.images.mixins import PlotMixin, BandMethodsTemplate, Deprecations | |
from shapely import ops, wkt | |
from shapely.geometry import box, shape, mapping, asShape | |
from shapely.geometry.base import BaseGeometry | |
import pyproj | |
import dask | |
from dask import optimization | |
from dask.delayed import delayed | |
import dask.array as da | |
import numpy as np | |
from affine import Affine | |
threads = int(os.environ.get('GBDX_THREADS', 8)) | |
threaded_get = partial(dask.threaded.get, num_workers=threads) | |
class DaskMeta(namedtuple("DaskMeta", ["dask", "name", "chunks", "dtype", "shape"])): | |
__slots__ = () | |
@classmethod | |
def from_darray(cls, darr, new=tuple.__new__, len=len): | |
dsk, _ = optimization.cull(darr.dask, darr.__dask_keys__()) | |
itr = [dsk, darr.name, darr.chunks, darr.dtype, darr.shape] | |
return cls._make(itr) | |
@property | |
def values(self): | |
return self._asdict().values() | |
class DaskImage(da.Array): | |
""" | |
A DaskImage is a 2 or 3 dimension dask array that contains implements the `__daskmeta__` interface. | |
""" | |
def __new__(cls, dm, **kwargs): | |
if isinstance(dm, da.Array): | |
dm = DaskMeta.from_darray(dm) | |
elif isinstance(dm, dict): | |
dm = DaskMeta(**dm) | |
elif isinstance(dm, DaskMeta): | |
pass | |
elif dm.__class__.__name__ in ("Op", "GraphMeta", "TmsMeta", "TemplateMeta"): | |
itr = [dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape] | |
dm = DaskMeta._make(itr) | |
else: | |
raise ValueError("{} must be initialized with a DaskMeta, a dask array, or a dict with DaskMeta fields".format(cls.__name__)) | |
self = da.Array.__new__(cls, dm.dask, dm.name, dm.chunks, dtype=dm.dtype, shape=dm.shape) | |
if "__geo_transform__" in kwargs: | |
self.__geo_transform__ = kwargs["__geo_transform__"] | |
if "__geo_interface__" in kwargs: | |
self.__geo_interface__ = kwargs["__geo_interface__"] | |
return self | |
@property | |
def __daskmeta__(self): | |
return DaskMeta(self) | |
def read(self, bands=None, **kwargs): | |
"""Reads data from a dask array and returns the computed ndarray matching the given bands | |
Args: | |
bands (list): band indices to read from the image. Returns bands in the order specified in the list of bands. | |
Returns: | |
ndarray: a numpy array of image data | |
""" | |
arr = self | |
if bands is not None: | |
arr = self[bands, ...] | |
return arr.compute(scheduler=threaded_get) | |
def randwindow(self, window_shape): | |
"""Get a random window of a given shape from within an image | |
Args: | |
window_shape (tuple): The desired shape of the returned image as (height, width) in pixels. | |
Returns: | |
image: a new image object of the specified shape and same type | |
""" | |
row = random.randrange(window_shape[0], self.shape[1]) | |
col = random.randrange(window_shape[1], self.shape[2]) | |
return self[:, row-window_shape[0]:row, col-window_shape[1]:col] | |
def iterwindows(self, count=64, window_shape=(256, 256)): | |
""" Iterate over random windows of an image | |
Args: | |
count (int): the number of the windows to generate. Defaults to 64, if `None` will continue to iterate over random windows until stopped. | |
window_shape (tuple): The desired shape of each image as (height, width) in pixels. | |
Yields: | |
image: an image of the given shape and same type. | |
""" | |
if count is None: | |
while True: | |
yield self.randwindow(window_shape) | |
else: | |
for i in range(count): | |
yield self.randwindow(window_shape) | |
def window_at(self, geom, window_shape): | |
"""Return a subsetted window of a given size, centered on a geometry object | |
Useful for generating training sets from vector training data | |
Will throw a ValueError if the window is not within the image bounds | |
Args: | |
geom (shapely,geometry): Geometry to center the image on | |
window_shape (tuple): The desired shape of the image as (height, width) in pixels. | |
Returns: | |
image: image object of same type | |
""" | |
# Centroids of the input geometry may not be centered on the object. | |
# For a covering image we use the bounds instead. | |
# This is also a workaround for issue 387. | |
y_size, x_size = window_shape[0], window_shape[1] | |
bounds = box(*geom.bounds) | |
px = ops.transform(self.__geo_transform__.rev, bounds).centroid | |
miny, maxy = int(px.y - y_size/2), int(px.y + y_size/2) | |
minx, maxx = int(px.x - x_size/2), int(px.x + x_size/2) | |
_, y_max, x_max = self.shape | |
if minx < 0 or miny < 0 or maxx > x_max or maxy > y_max: | |
raise ValueError("Input geometry resulted in a window outside of the image") | |
return self[:, miny:maxy, minx:maxx] | |
def window_cover(self, window_shape, pad=True): | |
""" Iterate over a grid of windows of a specified shape covering an image. | |
The image is divided into a grid of tiles of size window_shape. Each iteration returns | |
the next window. | |
Args: | |
window_shape (tuple): The desired shape of each image as (height, | |
width) in pixels. | |
pad: (bool): Whether or not to pad edge cells. If False, cells that do not | |
have the desired shape will not be returned. Defaults to True. | |
Yields: | |
image: image object of same type. | |
""" | |
size_y, size_x = window_shape[0], window_shape[1] | |
_ndepth, _nheight, _nwidth = self.shape | |
nheight, _m = divmod(_nheight, size_y) | |
nwidth, _n = divmod(_nwidth, size_x) | |
img = self | |
if pad is True: | |
new_height, new_width = _nheight, _nwidth | |
if _m != 0: | |
new_height = (nheight + 1) * size_y | |
if _n != 0: | |
new_width = (nwidth + 1) * size_x | |
if (new_height, new_width) != (_nheight, _nwidth): | |
bounds = box(0, 0, new_width, new_height) | |
geom = ops.transform(self.__geo_transform__.fwd, bounds) | |
img = self[geom] | |
row_lims = range(0, img.shape[1], size_y) | |
col_lims = range(0, img.shape[2], size_x) | |
for maxy, maxx in product(row_lims, col_lims): | |
reg = img[:, maxy:(maxy + size_y), maxx:(maxx + size_x)] | |
if pad is False: | |
if reg.shape[1:] == window_shape: | |
yield reg | |
else: | |
yield reg | |
class GeoDaskImage(DaskImage, Container, PlotMixin, BandMethodsTemplate, Deprecations): | |
_default_proj = "EPSG:4326" | |
def map_blocks(self, *args, **kwargs): | |
''' Queue a deferred function to run on each block of image | |
This is identical to Dask's map_block functinos, but returns a GeoDaskImage to preserve | |
the geospatial information. | |
Args: see dask.Array.map_blocks | |
Returns: | |
GeoDaskImage: a dask array with the function queued up to run when the image is read | |
''' | |
darr = super(GeoDaskImage, self).map_blocks(*args, **kwargs) | |
return GeoDaskImage(darr, __geo_interface__ = self.__geo_interface__, | |
__geo_transform__ = self.__geo_transform__) | |
def rechunk(self, *args, **kwargs): | |
darr = super(GeoDaskImage, self).rechunk(*args, **kwargs) | |
return GeoDaskImage(darr, __geo_interface__ = self.__geo_interface__, | |
__geo_transform__ = self.__geo_transform__) | |
def asShape(self): | |
return asShape(self) | |
@property | |
def affine(self): | |
""" The geo transform of the image | |
Returns: | |
dict: The image's affine transform | |
""" | |
# TODO add check for Ratpoly or whatevs | |
return self.__geo_transform__._affine | |
@property | |
def bounds(self): | |
""" Access the spatial bounding box of the image | |
Returns: | |
list: list of bounds in image projected coordinates (minx, miny, maxx, maxy) | |
""" | |
return shape(self).bounds | |
@property | |
def proj(self): | |
""" The projection of the image """ | |
# keep EPSG code strings upper case | |
return self.__geo_transform__.proj.replace('epsg', 'EPSG') | |
def aoi(self, **kwargs): | |
""" Subsets the Image by the given bounds | |
Args: | |
bbox (list): optional. A bounding box array [minx, miny, maxx, maxy] | |
wkt (str): optional. A WKT geometry string | |
geojson (str): optional. A GeoJSON geometry dictionary | |
Returns: | |
image: an image instance of the same type | |
""" | |
g = self._parse_geoms(**kwargs) | |
if g is None: | |
return self | |
else: | |
return self[g] | |
def pxbounds(self, geom, clip=False): | |
""" Returns the bounds of a geometry object in pixel coordinates | |
Args: | |
geom: Shapely geometry object or GeoJSON as Python dictionary or WKT string | |
clip (bool): Clip the bounds to the min/max extent of the image | |
Returns: | |
list: bounds in pixels [min x, min y, max x, max y] clipped to image bounds | |
""" | |
try: | |
if isinstance(geom, dict): | |
if 'geometry' in geom: | |
geom = shape(geom['geometry']) | |
else: | |
geom = shape(geom) | |
elif isinstance(geom, BaseGeometry): | |
geom = shape(geom) | |
else: | |
geom = wkt.loads(geom) | |
except: | |
raise TypeError ("Invalid geometry object") | |
# if geometry doesn't overlap the image, return an error | |
if geom.disjoint(shape(self)): | |
raise ValueError("Geometry outside of image bounds") | |
# clip to pixels within the image | |
(xmin, ymin, xmax, ymax) = ops.transform(self.__geo_transform__.rev, geom).bounds | |
_nbands, ysize, xsize = self.shape | |
if clip: | |
xmin = max(xmin, 0) | |
ymin = max(ymin, 0) | |
xmax = min(xmax, xsize) | |
ymax = min(ymax, ysize) | |
return (xmin, ymin, xmax, ymax) | |
def geotiff(self, **kwargs): | |
""" Creates a geotiff on the filesystem | |
Args: | |
path (str): optional, path to write the geotiff file to, default is ./output.tif | |
proj (str): optional, EPSG string of projection to reproject to | |
spec (str): optional, if set to 'rgb', write out color-balanced 8-bit RGB tif | |
bands (list): optional, list of bands to export. If spec='rgb' will default to RGB bands, | |
otherwise will export all bands | |
Returns: | |
str: path the geotiff was written to """ | |
if 'proj' not in kwargs: | |
kwargs['proj'] = self.proj | |
return to_geotiff(self, **kwargs) | |
def _parse_geoms(self, **kwargs): | |
""" Finds supported geometry types, parses them and returns the bbox """ | |
bbox = kwargs.get('bbox', None) | |
wkt_geom = kwargs.get('wkt', None) | |
geojson = kwargs.get('geojson', None) | |
if bbox is not None: | |
g = box(*bbox) | |
elif wkt_geom is not None: | |
g = wkt.loads(wkt_geom) | |
elif geojson is not None: | |
g = shape(geojson) | |
else: | |
return None | |
if self.proj is None: | |
return g | |
else: | |
return self._reproject(g, from_proj=kwargs.get('from_proj', 'EPSG:4326')) | |
def _reproject(self, geometry, from_proj=None, to_proj=None): | |
if from_proj is None: | |
from_proj = self._default_proj | |
if to_proj is None: | |
to_proj = self.proj if self.proj is not None else "EPSG:4326" | |
tfm = pyproj.Transformer.from_crs(get_proj(from_proj), get_proj(to_proj), always_xy=True) | |
return ops.transform(tfm.transform, geometry) | |
def _slice_padded(self, _bounds): | |
pads = (max(-_bounds[0], 0), max(-_bounds[1], 0), | |
max(_bounds[2]-self.shape[2], 0), max(_bounds[3]-self.shape[1], 0)) | |
bounds = (max(_bounds[0], 0), | |
max(_bounds[1], 0), | |
max(min(_bounds[2], self.shape[2]), 0), | |
max(min(_bounds[3], self.shape[1]), 0)) | |
result = self[:, bounds[1]:bounds[3], bounds[0]:bounds[2]] | |
if pads[0] > 0: | |
dims = (result.shape[0], result.shape[1], pads[0]) | |
result = da.concatenate([da.zeros(dims, chunks=dims, dtype=result.dtype), | |
result], axis=2) | |
if pads[2] > 0: | |
dims = (result.shape[0], result.shape[1], pads[2]) | |
result = da.concatenate([result, | |
da.zeros(dims, chunks=dims, dtype=result.dtype)], axis=2) | |
if pads[1] > 0: | |
dims = (result.shape[0], pads[1], result.shape[2]) | |
result = da.concatenate([da.zeros(dims, chunks=dims, dtype=result.dtype), | |
result], axis=1) | |
if pads[3] > 0: | |
dims = (result.shape[0], pads[3], result.shape[2]) | |
result = da.concatenate([result, | |
da.zeros(dims, chunks=dims, dtype=result.dtype)], axis=1) | |
return (result, _bounds[0], _bounds[1]) | |
def __contains__(self, g): | |
geometry = ops.transform(self.__geo_transform__.rev, g) | |
img_bounds = box(0, 0, *self.shape[2:0:-1]) | |
return img_bounds.contains(geometry) | |
def __getitem__(self, geometry): | |
if isinstance(geometry, BaseGeometry) or getattr(geometry, "__geo_interface__", None) is not None: | |
g = shape(geometry) | |
if g.disjoint(shape(self)): | |
raise ValueError("AOI does not intersect image: {} not in {}".format(g.bounds, self.bounds)) | |
bounds = ops.transform(self.__geo_transform__.rev, g).bounds | |
result, xmin, ymin = self._slice_padded(bounds) | |
else: | |
if len(geometry) == 1: | |
assert geometry[0] == Ellipsis | |
return self | |
elif len(geometry) == 2: | |
arg0, arg1 = geometry | |
if isinstance(arg1, slice): | |
assert arg0 == Ellipsis | |
return self[:, :, arg1.start:arg1.stop] | |
elif arg1 == Ellipsis: | |
return self[arg0, :, :] | |
elif len(geometry) == 3: | |
try: | |
nbands, ysize, xsize = self.shape | |
except: | |
ysize, xsize = self.shape | |
band_idx, y_idx, x_idx = geometry | |
if y_idx == Ellipsis: | |
y_idx = slice(0, ysize) | |
if x_idx == Ellipsis: | |
x_idx = slice(0, xsize) | |
if not(isinstance(y_idx, slice) and isinstance(x_idx, slice)): | |
di = DaskImage(self) | |
return di.__getitem__(geometry) | |
xmin, ymin, xmax, ymax = x_idx.start, y_idx.start, x_idx.stop, y_idx.stop | |
xmin = 0 if xmin is None else xmin | |
ymin = 0 if ymin is None else ymin | |
xmax = xsize if xmax is None else xmax | |
ymax = ysize if ymax is None else ymax | |
if ymin > ysize and xmin > xsize: | |
raise IndexError("Index completely out of image bounds") | |
g = ops.transform(self.__geo_transform__.fwd, box(xmin, ymin, xmax, ymax)) | |
result = super(GeoDaskImage, self).__getitem__(geometry) | |
else: | |
return super(GeoDaskImage, self).__getitem__(geometry) | |
gi = mapping(g) | |
gt = self.__geo_transform__ + (xmin, ymin) | |
image = super(GeoDaskImage, self.__class__).__new__(self.__class__, result, __geo_interface__ = gi, __geo_transform__ = gt) | |
return image |