conversion.py 3.89 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:
Daniel Scheffler's avatar
Daniel Scheffler committed
29
    :param DN2extract:
30
    :param exact:
Daniel Scheffler's avatar
Daniel Scheffler committed
31
    :param maxfeatCount:
32
33
34
    :param timeout:           breaks the process after a given time in seconds
    :param progress:          show progress bars (default: True)
    :param q:                 quiet mode (default: False)
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
    :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)

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

    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
65

66
67
68
    errMsg = gdal.GetLastErrorMsg()

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

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

    # extract polygon
Daniel Scheffler's avatar
Daniel Scheffler committed
75
76
77
78
79
80
81
82
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
    mem_layer.SetAttributeFilter('DN = %s' %DN2extract)

    from geopandas import GeoDataFrame
    featCount = mem_layer.GetFeatureCount()

    if maxfeatCount and featCount > maxfeatCount:
        raise RuntimeError('Found %s features with DN=%s but maximum feature count was set to %s.'
                           %(featCount, DN2extract, maxfeatCount))

    #if featCount>1:
    #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
112

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