gdal.py 4.31 KB
Newer Older
1
# -*- coding: utf-8 -*-
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# py_tools_ds
#
# Copyright (C) 2019  Daniel Scheffler (GFZ Potsdam, daniel.scheffler@gfz-potsdam.de)
#
# This software was developed within the context of the GeoMultiSens project funded
# by the German Federal Ministry of Education and Research
# (project grant code: 01 IS 14 010 A-C).
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with this program.  If not, see <http://www.gnu.org/licenses/>.

24
25
__author__ = "Daniel Scheffler"

26
27
28
import time
import os

29
import numpy as np
30
from pandas import DataFrame
31
from osgeo import gdal_array
32

33
34
35
try:
    from osgeo import gdal
except ImportError:
Daniel Scheffler's avatar
Daniel Scheffler committed
36
37
    import gdal

38

39
from ...dtypes.conversion import dTypeDic_NumPy2GDALcompatible
40
41


42
def get_GDAL_ds_inmem(array, gt=None, prj=None, nodata=None):
43
44
    """Convert a numpy array into a GDAL dataset.
    NOTE: Possibly the data type has to be automatically changed in order ensure GDAL compatibility!
45

46
    :param array:   <numpy.ndarray> in the shape (rows, columns, bands)
47
48
    :param gt:
    :param prj:
Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
49
    :param nodata:  <int> nodata value to be set (GDAL seems to have issues with non-int nodata values.)
50
51
    :return:
    """
52
    # FIXME does not respect different nodata values for each band
53

54
55
    if len(array.shape) == 3:
        array = np.rollaxis(array, 2)  # rows,cols,bands => bands,rows,cols
56

57
58
59
60
    # convert data type to GDAL compatible data type
    if gdal_array.NumericTypeCodeToGDALTypeCode(array.dtype) is None:
        array = array.astype(dTypeDic_NumPy2GDALcompatible[str(np.dtype(array.dtype))])

61
    ds = gdal_array.OpenArray(array)  # uses interleave='band' by default
62

63
64
    if ds is None:
        raise Exception(gdal.GetLastErrorMsg())
65
66
67
68
69
    if gt:
        ds.SetGeoTransform(gt)
    if prj:
        ds.SetProjection(prj)

Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
70
    if nodata is not None:
71
        for i in range(ds.RasterCount):
72
            band = ds.GetRasterBand(i + 1)
Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
73
            try:
74
                band.SetNoDataValue(nodata)
Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
75
            except TypeError:
76
                raise TypeError(type(nodata), 'TypeError while trying to set NoDataValue to %s. ' % nodata)
77
78
79
            del band

    ds.FlushCache()  # Write to disk.
80
81
82
83
84
    return ds


def get_GDAL_driverList():
    count = gdal.GetDriverCount()
85
    df = DataFrame(np.full((count, 5), np.nan), columns=['drvCode', 'drvLongName', 'ext1', 'ext2', 'ext3'])
86
87
88
89
90
    for i in range(count):
        drv = gdal.GetDriver(i)
        if drv.GetMetadataItem(gdal.DCAP_RASTER):
            meta = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
            extensions = meta.split() if meta else []
Daniel Scheffler's avatar
Daniel Scheffler committed
91
92
93
94
95
            df.loc[i] = [drv.GetDescription(),
                         drv.GetMetadataItem(gdal.DMD_LONGNAME),
                         extensions[0] if len(extensions) > 0 else np.nan,
                         extensions[1] if len(extensions) > 1 else np.nan,
                         extensions[2] if len(extensions) > 2 else np.nan]
96
    df = df.dropna(how='all')
97
    return df
98
99
100
101
102
103
104
105


def wait_if_used(path_file, lockfile, timeout=100, try_kill=0):
    globs = globals()
    same_gdalRefs = [k for k, v in globs.items() if
                     isinstance(globs[k], gdal.Dataset) and globs[k].GetDescription() == path_file]
    t0 = time.time()

106
107
    def update_same_gdalRefs(sRs):
        return [sR for sR in sRs if sR in globals() and globals()[sR] is not None]
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

    while same_gdalRefs != [] or os.path.exists(lockfile):
        if os.path.exists(lockfile):
            continue

        if time.time() - t0 > timeout:
            if try_kill:
                for sR in same_gdalRefs:
                    globals()[sR] = None
                    print('had to kill %s' % sR)
            else:
                if os.path.exists(lockfile):
                    os.remove(lockfile)

                raise TimeoutError('The file %s is permanently used by another variable.' % path_file)

        same_gdalRefs = update_same_gdalRefs(same_gdalRefs)