2

我正在尝试找到将 dask 数组保存到 geotiff 的最佳/最快方法。

我一直在研究如何使 rasterio/GDAL 线程安全，但我没有想出任何可以从python轻松访问的东西。

如果我使用da.store(..., lock=False)，那么文件可能会损坏。我尝试为每次写入打开和关闭文件，但这仍然使用相同的文件描述符（如果我错了，请纠正我）并且无论如何都不是一个很好的解决方案。

有没有人有另一种方法来做到这一点，以便每个 dask 工作人员（线程）可以安全地写入由rasterio库创建的geotiff文件？当前的工作解决方案是保留其默认值.dask.storelock-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 的内部缓存大小来获得成功的输出，尽管不能保证。

4

0 回答 0