从dask数组存储

时间:2018-03-03 23:38:27

标签: python gdal dask geotiff rasterio

我正在尝试找到将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的内部缓存大小来获得成功,尽管不能保证输出。

0 个答案:

没有答案