我正在尝试找到将dask数组保存到geotiff的最佳/最快方法。
我一直在研究如何使rasterio / GDAL线程安全,但我没有提出任何可以从 python 轻松访问的内容。
如果我使用da.store(..., lock=False)
,那么该文件可能会损坏。
我尝试为每次写入打开和关闭文件,但仍然使用相同的文件描述符(如果我错了,请纠正我)并且无论如何都不是一个好的解决方案。
有没有人有另一种方法可以让每个dask worker(线程)安全地写入 rasterio 库创建的 geotiff 文件?目前的工作解决方案是将dask.store
保留为默认值lock-True
。
我猜任何其他解决方案都会涉及到线程锁定,但我认为无论如何都有这个解决方案。
我的工作示例代码如下:
import dask.array as da
import numpy as np
import rasterio
from rasterio.windows import Window
class RIOFile(object):
"""Rasterio wrapper to allow da.store to do window saving."""
def __init__(self, *args, **kwargs):
"""Initialize the object."""
self.args = args
self.kwargs = kwargs
self.rfile = None
def __setitem__(self, key, item):
"""Put the data chunk in the image."""
if len(key) == 3:
indexes = list(range(
key[0].start + 1,
key[0].stop + 1,
key[0].step or 1
))
y = key[1]
x = key[2]
else:
indexes = 1
y = key[0]
x = key[1]
chy_off = y.start
chy = y.stop - y.start
chx_off = x.start
chx = x.stop - x.start
# band indexes
self.rfile.write(item, window=Window(chx_off, chy_off, chx, chy),
indexes=indexes)
def __enter__(self):
"""Enter method."""
self.rfile = rasterio.open(*self.args, **self.kwargs)
return self
def __exit__(self, exc_type, exc_value, traceback):
"""Exit method."""
self.rfile.close()
rows = cols = 21696
a = da.ones((4, rows, cols), dtype=np.float64, chunks=(1, 4096, 4096) )
a = a * np.array([255., 255., 255., 255.])[:, None, None]
a = a.astype(np.uint8)
with RIOFile('test.tif', 'w', driver='GTiff', width=cols, height=rows, count=4, dtype=np.uint8) as r_file:
da.store(a, r_file, lock=True)
更改lock=False
可能已损坏的文件。
我也能够通过增加GDAL的内部缓存大小来获得成功,尽管不能保证输出。