conversion.py 4.49 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

from shapely.wkb import loads
7
8
import numpy as np
from skimage.transform import rescale
9

10
11
12
13
14
15
16
17
18
try:
    import gdal
    import ogr
    import osr
except ImportError:
    from osgeo import gdal
    from osgeo import ogr
    from osgeo import osr

19
20
from ...io.raster.gdal import get_GDAL_ds_inmem
from ...processing.progress_mon import ProgressBar, Timer
Daniel Scheffler's avatar
Daniel Scheffler committed
21
from ...compatibility.python.exceptions import TimeoutError as TimeoutError_comp
22
23


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

28
    :param array:             2D numpy array
29
30
    :param gt:
    :param prj:
31
    :param DN2extract:        <int, float> pixel value to create polygons for
32
    :param exact:
33
34
    :param maxfeatCount:      <int> the maximum expected number of polygons. If more polygons are found, every further
                              processing is cancelled and a RunTimeError is raised.
35
36
37
    :param timeout:           breaks the process after a given time in seconds
    :param progress:          show progress bars (default: True)
    :param q:                 quiet mode (default: False)
38
39
40
    :return:
    """

41
42
    assert array.ndim == 2, "Only 2D arrays are supported. Got a %sD array." % array.ndim

43
44
    # 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
45
46
47
48
49
50
51
52
53
54
        zoom_factor = 0.5

        # downsample to half size, nearest neighbour
        array = rescale(array, zoom_factor, order=0, preserve_range=True, mode='edge').astype(np.bool)

        # update pixel sizes within gt
        gt = list(gt)
        gt[1] /= zoom_factor
        gt[5] /= zoom_factor

55
    src_ds = get_GDAL_ds_inmem(array, gt, prj)
56
57
58
59
    src_band = src_ds.GetRasterBand(1)

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

    srs = osr.SpatialReference()
63
    srs.ImportFromWkt(prj)
64
65
66
67
68
69

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

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

70
    # set callback
Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
71
72
    callback = ProgressBar(prefix='Polygonize progress    ', suffix='Complete', barLength=50, timeout=timeout,
                           use_as_callback=True) \
73
        if progress and not q else Timer(timeout, use_as_callback=True) if timeout else None
74
75

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

Daniel Scheffler's avatar
Bugfix    
Daniel Scheffler committed
79
80
81
82
83
    # 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!')
84
        raise Exception(errMsg)
85
86

    # extract polygon
87
    mem_layer.SetAttributeFilter('DN = %s' % DN2extract)
Daniel Scheffler's avatar
Daniel Scheffler committed
88
89
90
91

    from geopandas import GeoDataFrame
    featCount = mem_layer.GetFeatureCount()

92
    if not featCount:
93
        raise RuntimeError('No features with DN=%s found in the input image.' % DN2extract)
Daniel Scheffler's avatar
Daniel Scheffler committed
94
95
    if maxfeatCount and featCount > maxfeatCount:
        raise RuntimeError('Found %s features with DN=%s but maximum feature count was set to %s.'
96
                           % (featCount, DN2extract, maxfeatCount))
Daniel Scheffler's avatar
Daniel Scheffler committed
97

98
99
100
    # 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
101

102
    # def get_shplyPoly(GDF_row):
Daniel Scheffler's avatar
Daniel Scheffler committed
103
104
105
106
107
108
109
110
    #    if not is_timed_out(3):
    #        element   = mem_layer.GetNextFeature()
    #        shplyPoly = loads(element.GetGeometryRef().ExportToWkb()).buffer(0)
    #        element   = None
    #        return shplyPoly
    #    else:
    #        raise TimeoutError

111
    # GDF['geometry'] = GDF.apply(get_shplyPoly, axis=1)
Daniel Scheffler's avatar
Daniel Scheffler committed
112
113

    GDF = GeoDataFrame(columns=['geometry', 'DN'])
114
    timer = Timer(timeout)
Daniel Scheffler's avatar
Daniel Scheffler committed
115
    for i in range(featCount):
116
        if not timer.timed_out:
117
            element = mem_layer.GetNextFeature()
Daniel Scheffler's avatar
Daniel Scheffler committed
118
            GDF.loc[i] = [loads(element.GetGeometryRef().ExportToWkb()).buffer(0), DN2extract]
119
            del element
Daniel Scheffler's avatar
Daniel Scheffler committed
120
121
122
123
124
        else:
            raise TimeoutError('raster2polygon timed out!') if PY3 else TimeoutError_comp('raster2polygon timed out!')

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

125
    del mem_ds, mem_layer
126

127
    shplyPoly = GDF.loc[1, 'geometry']
128
    return shplyPoly