conversion.py 1.57 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
# -*- 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

from ...io.raster.gdal import get_GDAL_ds_inmem


def raster2polygon(array_or_GeoArray, gt=None, prj=None, exact=True):
    """Calculates a footprint polygon for the given array or GeoArray.

    :param array_or_GeoArray:
    :param gt:
    :param prj:
    :param exact:
    :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)

    # run the algorithm.
    result = gdal.Polygonize(src_band, src_band.GetMaskBand(), mem_layer, 0, ["8CONNECTED=8"] if exact else [])

    if result is None:
        raise Exception(gdal.GetLastErrorMsg())

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

    return shplyPoly