Skip to content

Instantly share code, notes, and snippets.

@v0lat1le
Created November 25, 2015 22:30
Show Gist options
  • Save v0lat1le/beb7d9a94dc5b4f6de7e to your computer and use it in GitHub Desktop.
Save v0lat1le/beb7d9a94dc5b4f6de7e to your computer and use it in GitHub Desktop.
xray DataStore for geotiffs
class GeoTiffBandArrayWrapper(xray.core.utils.NdimSizeLenMixin):
def __init__(self, raster, band):
self._band = raster.GetRasterBand(band)
self.dtype = {gdalconst.GDT_Int16: numpy.int16}[self._band.DataType]
self.shape = (raster.RasterYSize, raster.RasterXSize)
def __array__(self, dtype=None):
return self._band.ReadAsArray()
def __getitem__(self, key):
idxs = []
slicer = []
for width, idx in zip(self.shape, key):
if isinstance(idx, slice):
idxs.append(slice(idx.start or 0, idx.stop or width, idx.step or 1))
slicer.append(slice(0, idxs[-1].stop-idxs[-1].start, idx.step))
else:
idxs.append(slice(idx, idx+1, 1))
slicer.append(0)
data = self._band.ReadAsArray(idxs[1].start, idxs[0].start, idxs[1].stop-idxs[1].start, idxs[0].stop-idxs[0].start)
return data[slicer]
def __repr__(self):
return '%s(array=%r)' % (type(self).__name__, None)
class DataCubeV1DataStore(xray.backends.common.AbstractDataStore):
def __init__(self, filepath):
super(DataCubeV1DataStore, self).__init__()
self._raster = gdal.Open(filepath, gdalconst.GA_ReadOnly)
def get_attrs(self):
return {}
def _read_data(self, band):
print "read all"
return self._raster.GetRasterBand(band).ReadAsArray()
def get_variables(self):
transform = self._raster.GetGeoTransform()
shape = (self._raster.RasterXSize, self._raster.RasterYSize)
lon_idx = numpy.linspace(transform[0], transform[0]+transform[1]*shape[0], num=shape[0])
lat_idx = numpy.linspace(transform[3], transform[3]+transform[5]*shape[1], num=shape[1])
name = "BLUEskvnskldvn"
dask = {(name, 0, 0): (self._read_data, BLUE)}
blue = da.Array(name, dask, ((shape[1],),(shape[0],)), numpy.int16)
print blue
return {
"longitude": xray.Variable("longitude", lon_idx),
"latitude": xray.Variable("latitude", lat_idx),
"BLUE": xray.Variable(["latitude", "longitude"], blue)
}
def close(self):
super(DataCubeV1DataStore, self).close()
self._raster = None
def reader(band, x, y, w, h, raster):
data = band.ReadAsArray(x,y,w,h)
return data
def chunkedBand(dataset_path, band, name):
raster = gdal.Open(dataset_path, gdalconst.GA_ReadOnly)
band = raster.GetRasterBand(band)
block = [4000,250]#band.GetBlockSize()
#name = name
dtype = numpy.int16
dsk = {}
chunks = ((block[1], )*(raster.RasterYSize//block[1]), (block[0],)*(raster.RasterXSize//block[0]))
for j in xrange(raster.RasterYSize//block[1]):
for i in xrange(raster.RasterXSize//block[0]):
dsk[(name,j,i)] = (reader, band, i*block[0], j*block[1], block[0], block[1], raster)
return da.Array(dsk, name, chunks, dtype)
red = chunkedBand(dataset_path, RED, "RED_jbfjv")
nir = chunkedBand(dataset_path, NEAR_IR, "NIR_sdjnsl")
ndvi = (nir-red)/(nir+red)
ndvi.compute()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment