conversion.py 2.47 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# -*- coding: utf-8 -*-
__author__ = "Daniel Scheffler"



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
19
from ...io.raster.gdal                  import get_GDAL_ds_inmem
from ...processing.progress_mon         import printProgress, is_timed_out
from ...compatibility.python.exceptions import TimeoutError
20
21


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

    :param array_or_GeoArray:
    :param gt:
    :param prj:
    :param exact:
29
30
31
    :param timeout:           breaks the process after a given time in seconds
    :param progress:          show progress bars (default: True)
    :param q:                 quiet mode (default: False)
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
    :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)

52
53
54
55
56
57
58
59
60
61
62
63
64
65
    # set callback
    progressBar      = lambda percent01, message, user_data: \
        printProgress(percent01 * 100, prefix='Polygonize progress    ', suffix='Complete', barLength=50, timeout=3)
    timeout_callback = lambda percent01, message, user_data: is_timed_out(3)

    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)
    errMsg = gdal.GetLastErrorMsg()

    if errMsg == 'User terminated':
        raise TimeoutError('raster2polygon timed out!')
66
67

    if result is None:
68
        raise Exception(errMsg)
69
70
71
72
73
74
75
76
77

    # extract polygon
    mem_layer.SetAttributeFilter('DN = 1')
    element   = mem_layer.GetNextFeature() # asserts that
    shplyPoly = loads(element.GetGeometryRef().ExportToWkb())

    return shplyPoly