conversion.py 5.27 KB
Newer Older
1
# -*- coding: utf-8 -*-
2
3
4

# py_tools_ds
#
Daniel Scheffler's avatar
Daniel Scheffler committed
5
# Copyright (C) 2016-2021  Daniel Scheffler (GFZ Potsdam, daniel.scheffler@gfz-potsdam.de)
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
#
# 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
26
__author__ = "Daniel Scheffler"

from shapely.wkb import loads
27
from osgeo import gdal, osr, ogr
28

29
30
from ...io.raster.gdal import get_GDAL_ds_inmem
from ...processing.progress_mon import ProgressBar, Timer
31
from ..raster.reproject import warp_ndarray
32
33


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

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

52
53
    # 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
54
        # downsample with nearest neighbour
55
56
57
58
59
        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)
60

61
    src_ds = get_GDAL_ds_inmem(array, gt, prj)
62
63
64
65
    src_band = src_ds.GetRasterBand(1)

    # Create a memory OGR datasource to put results in.
    mem_drv = ogr.GetDriverByName('Memory')
66
    mem_ds = mem_drv.CreateDataSource('out')
67
68

    srs = osr.SpatialReference()
69
    srs.ImportFromWkt(prj)
70
71
72
73
74
75

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

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

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

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

91
92
93
94
95
96
        # handle exit status other than 0 (fail)
        if status != 0:
            errMsg = gdal.GetLastErrorMsg()
            if errMsg == 'User terminated':
                raise TimeoutError('raster2polygon timed out!')
            raise Exception(errMsg)
97
98

    # extract polygon
99
    mem_layer.SetAttributeFilter('DN = %s' % DN2extract)
Daniel Scheffler's avatar
Daniel Scheffler committed
100
101
102
103

    from geopandas import GeoDataFrame
    featCount = mem_layer.GetFeatureCount()

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

110
111
112
    # 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
113

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

123
    # GDF['geometry'] = GDF.apply(get_shplyPoly, axis=1)
Daniel Scheffler's avatar
Daniel Scheffler committed
124
125

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

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

137
    del mem_ds, mem_layer
138

139
    shplyPoly = GDF.loc[1, 'geometry']
140
    return shplyPoly