conversion.py 4.21 KB
Newer Older
1
2
3
# -*- coding: utf-8 -*-
__author__ = "Daniel Scheffler"

Daniel Scheffler's avatar
Daniel Scheffler committed
4
from six import PY3
5
6
7
8
9
10
11
12
13
14
15
16

from shapely.wkb import loads
import numpy as np
try:
    import gdal
    import ogr
    import osr
except ImportError:
    from osgeo import gdal
    from osgeo import ogr
    from osgeo import osr

17
18
from ...io.raster.gdal                  import get_GDAL_ds_inmem
from ...processing.progress_mon         import printProgress, is_timed_out
Daniel Scheffler's avatar
Daniel Scheffler committed
19
from ...compatibility.python.exceptions import TimeoutError as TimeoutError_comp
20
21


Daniel Scheffler's avatar
Daniel Scheffler committed
22
23
def raster2polygon(array_or_GeoArray, gt=None, prj=None, DN2extract=1, exact=True, maxfeatCount=None,
                   timeout=None, progress=True, q=False):
24
25
26
27
28
    """Calculates a footprint polygon for the given array or GeoArray.

    :param array_or_GeoArray:
    :param gt:
    :param prj:
29
    :param DN2extract:        <int, float> pixel value to create polygons for
30
    :param exact:
31
32
    :param maxfeatCount:      <int> the maximum expected number of polygons. If more polygons are found, every further
                              processing is cancelled and a RunTimeError is raised.
33
34
35
    :param timeout:           breaks the process after a given time in seconds
    :param progress:          show progress bars (default: True)
    :param q:                 quiet mode (default: False)
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    :return:
    """
    from ... import GeoArray
    geoArr = array_or_GeoArray if isinstance(array_or_GeoArray, GeoArray) else GeoArray(array_or_GeoArray, gt, prj)

    src_ds   = get_GDAL_ds_inmem(geoArr.mask_nodata.astype(np.uint8), geoArr.gt, geoArr.prj)
    src_band = src_ds.GetRasterBand(1)

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

    srs = osr.SpatialReference()
    srs.ImportFromWkt(geoArr.prj)

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

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

56
57
    # set callback
    progressBar      = lambda percent01, message, user_data: \
Daniel Scheffler's avatar
Daniel Scheffler committed
58
59
        printProgress(percent01 * 100, prefix='Polygonize progress    ', suffix='Complete', barLength=50, timeout=timeout)
    timeout_callback = lambda percent01, message, user_data: is_timed_out(timeout)
60
61
62
63
64
65

    callback = progressBar if progress and not q else timeout_callback if timeout else None

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

67
68
69
    errMsg = gdal.GetLastErrorMsg()

    if errMsg == 'User terminated':
Daniel Scheffler's avatar
Daniel Scheffler committed
70
        raise TimeoutError('raster2polygon timed out!') if PY3 else TimeoutError_comp('raster2polygon timed out!')
71
72

    if result is None:
73
        raise Exception(errMsg)
74
75

    # extract polygon
Daniel Scheffler's avatar
Daniel Scheffler committed
76
77
78
79
80
    mem_layer.SetAttributeFilter('DN = %s' %DN2extract)

    from geopandas import GeoDataFrame
    featCount = mem_layer.GetFeatureCount()

81
82
    if not featCount:
        raise RuntimeError('No features with DN=%s found in the input image.' %DN2extract)
Daniel Scheffler's avatar
Daniel Scheffler committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
    if maxfeatCount and featCount > maxfeatCount:
        raise RuntimeError('Found %s features with DN=%s but maximum feature count was set to %s.'
                           %(featCount, DN2extract, maxfeatCount))

    #tmp = np.full((featCount,2), DN, geoArr.dtype)
    #tmp[:,0] = range(featCount)
    #GDF = GeoDataFrame(tmp, columns=['idx','DN'])

    #def get_shplyPoly(GDF_row):
    #    if not is_timed_out(3):
    #        element   = mem_layer.GetNextFeature()
    #        shplyPoly = loads(element.GetGeometryRef().ExportToWkb()).buffer(0)
    #        element   = None
    #        return shplyPoly
    #    else:
    #        raise TimeoutError

    #GDF['geometry'] = GDF.apply(get_shplyPoly, axis=1)

    GDF = GeoDataFrame(columns=['geometry', 'DN'])
    for i in range(featCount):
        if not is_timed_out(timeout):
            element    = mem_layer.GetNextFeature()
            GDF.loc[i] = [loads(element.GetGeometryRef().ExportToWkb()).buffer(0), DN2extract]
            element    = None
        else:
            raise TimeoutError('raster2polygon timed out!') if PY3 else TimeoutError_comp('raster2polygon timed out!')

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

    mem_ds = mem_layer = None
114

Daniel Scheffler's avatar
Daniel Scheffler committed
115
    shplyPoly = GDF.loc[1,'geometry']
116
117
118
    return shplyPoly