conversion.py 5.32 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"

Daniel Scheffler's avatar
Daniel Scheffler committed
26
from six import PY3
27
28

from shapely.wkb import loads
29
from osgeo import gdal, osr, ogr
30

31
32
from ...io.raster.gdal import get_GDAL_ds_inmem
from ...processing.progress_mon import ProgressBar, Timer
Daniel Scheffler's avatar
Daniel Scheffler committed
33
from ...compatibility.python.exceptions import TimeoutError as TimeoutError_comp
34
from ..raster.reproject import warp_ndarray
35
36


37
def raster2polygon(array, gt, prj, DN2extract=1, exact=True, maxfeatCount=None,
Daniel Scheffler's avatar
Daniel Scheffler committed
38
                   timeout=None, progress=True, q=False):
Daniel Scheffler's avatar
Daniel Scheffler committed
39
    """Calculate a footprint polygon for the given array.
40

41
    :param array:             2D numpy array
42
43
    :param gt:
    :param prj:
44
    :param DN2extract:        <int, float> pixel value to create polygons for
45
    :param exact:
46
47
    :param maxfeatCount:      <int> the maximum expected number of polygons. If more polygons are found, every further
                              processing is cancelled and a RunTimeError is raised.
48
49
50
    :param timeout:           breaks the process after a given time in seconds
    :param progress:          show progress bars (default: True)
    :param q:                 quiet mode (default: False)
51
52
    :return:
    """
53
54
    assert array.ndim == 2, "Only 2D arrays are supported. Got a %sD array." % array.ndim

55
56
    # downsample input array in case is has more than 1e8 pixels to prevent crash
    if not exact and array.size > 1e8:  # 10000 x 10000 px
57
        # downsample with nearest neighbour
58
59
60
61
62
        zoom_factor = (8000 * 8000 / array.size) ** 0.5
        array, gt, prj = warp_ndarray(array, gt, prj,
                                      out_gsd=(gt[1] / zoom_factor, gt[5] / zoom_factor),
                                      rspAlg='near',
                                      q=True)
63

64
    src_ds = get_GDAL_ds_inmem(array, gt, prj)
65
66
67
68
    src_band = src_ds.GetRasterBand(1)

    # Create a memory OGR datasource to put results in.
    mem_drv = ogr.GetDriverByName('Memory')
69
    mem_ds = mem_drv.CreateDataSource('out')
70
71

    srs = osr.SpatialReference()
72
    srs.ImportFromWkt(prj)
73
74
75
76
77
78

    mem_layer = mem_ds.CreateLayer('poly', srs, ogr.wkbPolygon)

    fd = ogr.FieldDefn('DN', ogr.OFTInteger)
    mem_layer.CreateField(fd)

79
    # set callback
Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
80
81
    callback = ProgressBar(prefix='Polygonize progress    ', suffix='Complete', barLength=50, timeout=timeout,
                           use_as_callback=True) \
82
        if progress and not q else Timer(timeout, use_as_callback=True) if timeout else None
83
84

    # run the algorithm
Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
85
    status = gdal.Polygonize(src_band, src_band.GetMaskBand(), mem_layer, 0, ["8CONNECTED=8"] if exact else [],
86
                             callback=callback)
Daniel Scheffler's avatar
Daniel Scheffler committed
87

Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
88
89
90
91
92
    # handle exit status other than 0 (fail)
    if status != 0:
        errMsg = gdal.GetLastErrorMsg()
        if errMsg == 'User terminated':
            raise TimeoutError('raster2polygon timed out!') if PY3 else TimeoutError_comp('raster2polygon timed out!')
93
        raise Exception(errMsg)
94
95

    # extract polygon
96
    mem_layer.SetAttributeFilter('DN = %s' % DN2extract)
Daniel Scheffler's avatar
Daniel Scheffler committed
97
98
99
100

    from geopandas import GeoDataFrame
    featCount = mem_layer.GetFeatureCount()

101
    if not featCount:
102
        raise RuntimeError('No features with DN=%s found in the input image.' % DN2extract)
Daniel Scheffler's avatar
Daniel Scheffler committed
103
104
    if maxfeatCount and featCount > maxfeatCount:
        raise RuntimeError('Found %s features with DN=%s but maximum feature count was set to %s.'
105
                           % (featCount, DN2extract, maxfeatCount))
Daniel Scheffler's avatar
Daniel Scheffler committed
106

107
108
109
    # tmp = np.full((featCount,2), DN, geoArr.dtype)
    # tmp[:,0] = range(featCount)
    # GDF = GeoDataFrame(tmp, columns=['idx','DN'])
Daniel Scheffler's avatar
Daniel Scheffler committed
110

111
    # def get_shplyPoly(GDF_row):
Daniel Scheffler's avatar
Daniel Scheffler committed
112
113
114
115
116
117
118
119
    #    if not is_timed_out(3):
    #        element   = mem_layer.GetNextFeature()
    #        shplyPoly = loads(element.GetGeometryRef().ExportToWkb()).buffer(0)
    #        element   = None
    #        return shplyPoly
    #    else:
    #        raise TimeoutError

120
    # GDF['geometry'] = GDF.apply(get_shplyPoly, axis=1)
Daniel Scheffler's avatar
Daniel Scheffler committed
121
122

    GDF = GeoDataFrame(columns=['geometry', 'DN'])
123
    timer = Timer(timeout)
Daniel Scheffler's avatar
Daniel Scheffler committed
124
    for i in range(featCount):
125
        if not timer.timed_out:
126
            element = mem_layer.GetNextFeature()
Daniel Scheffler's avatar
Daniel Scheffler committed
127
            GDF.loc[i] = [loads(element.GetGeometryRef().ExportToWkb()).buffer(0), DN2extract]
128
            del element
Daniel Scheffler's avatar
Daniel Scheffler committed
129
130
131
132
133
        else:
            raise TimeoutError('raster2polygon timed out!') if PY3 else TimeoutError_comp('raster2polygon timed out!')

    GDF = GDF.dissolve(by='DN')

134
    del mem_ds, mem_layer
135

136
    shplyPoly = GDF.loc[1, 'geometry']
137
    return shplyPoly