Python osgeo.ogr.GetDriverByName() Examples

def writeFile(filename,geotransform,geoprojection,data):
	(x,y) = data.shape
	format = "GTiff"
	noDataValue = -9999
	driver = gdal.GetDriverByName(format)
	# you can change the dataformat but be sure to be able to store negative values including -9999
	dst_datatype = gdal.GDT_Float32

	#print(data)

	dst_ds = driver.Create(filename,y,x,1,dst_datatype)
	dst_ds.GetRasterBand(1).WriteArray(data)
	dst_ds.GetRasterBand(1).SetNoDataValue( noDataValue )
	dst_ds.SetGeoTransform(geotransform)
	dst_ds.SetProjection(geoprojection)
	return 1 

def writeFile(filename,geotransform,geoprojection,data):
    (x,y) = data.shape
    format = "GTiff"
    noDataValue = -9999
    driver = gdal.GetDriverByName(format)
    # you can change the dataformat but be sure to be able to store negative values including -9999
    dst_datatype = gdal.GDT_Float32

    #print(data)

    dst_ds = driver.Create(filename,y,x,1,dst_datatype)
    dst_ds.GetRasterBand(1).WriteArray(data)
    dst_ds.GetRasterBand(1).SetNoDataValue( noDataValue )
    dst_ds.SetGeoTransform(geotransform)
    dst_ds.SetProjection(geoprojection)
    return 1 

def shp2OsmosisPolygon(daShapefile,txtFn):        
    driver = ogr.GetDriverByName('ESRI Shapefile') 
    infile = driver.Open(daShapefile) 
    layer = infile.GetLayer()
    f= open(txtFn,"w") 
    f.write("osmosis polygon\nfirst_area\n")
    for area in layer: 
        area_shape = area.GetGeometryRef() 
        area_polygon = area_shape.GetGeometryRef(0) 
        no_of_polygon_vertices = area_polygon.GetPointCount()         
        for vertex in range(no_of_polygon_vertices): 
            lon, lat, z = area_polygon.GetPoint(vertex)  #获取经纬度坐标
            print(lon,lat,z)
            f.write("%s  %s\n"%(lon,lat))
    f.write("END\nEND")        
            
    f.close() 

def writeRaster(arr, out_path, template):
    no_data = 0
    # First of all, gather some information from the template file
    data = gdal.Open(template)
    [cols, rows] = arr.shape
    trans = data.GetGeoTransform()
    proj = data.GetProjection()
    # nodatav = 0 #data.GetNoDataValue()
    # Create the file, using the information from the template file
    outdriver = gdal.GetDriverByName("GTiff")
    # http://www.gdal.org/gdal_8h.html
    # GDT_Byte = 1, GDT_UInt16 = 2, GDT_UInt32 = 4, GDT_Int32 = 5, GDT_Float32 = 6,
    outdata   = outdriver.Create(str(out_path), rows, cols, 1, gdal.GDT_UInt32)
    # Write the array to the file, which is the original array in this example
    outdata.GetRasterBand(1).WriteArray(arr)
    # Set a no data value if required
    outdata.GetRasterBand(1).SetNoDataValue(no_data)
    # Georeference the image
    outdata.SetGeoTransform(trans)
    # Write projection information
    outdata.SetProjection(proj)
    return arr


# raster to vector 

def polygonize(img,shp_path):
    # mapping between gdal type and ogr field type
    type_mapping = {gdal.GDT_Byte: ogr.OFTInteger,
                    gdal.GDT_UInt16: ogr.OFTInteger,
                    gdal.GDT_Int16: ogr.OFTInteger,
                    gdal.GDT_UInt32: ogr.OFTInteger,
                    gdal.GDT_Int32: ogr.OFTInteger,
                    gdal.GDT_Float32: ogr.OFTReal,
                    gdal.GDT_Float64: ogr.OFTReal,
                    gdal.GDT_CInt16: ogr.OFTInteger,
                    gdal.GDT_CInt32: ogr.OFTInteger,
                    gdal.GDT_CFloat32: ogr.OFTReal,
                    gdal.GDT_CFloat64: ogr.OFTReal}

    ds = gdal.Open(img)
    prj = ds.GetProjection()
    srcband = ds.GetRasterBand(1)
    dst_layername = "Shape"
    drv = ogr.GetDriverByName("ESRI Shapefile")
    dst_ds = drv.CreateDataSource(shp_path)
    srs = osr.SpatialReference(wkt=prj)

    dst_layer = dst_ds.CreateLayer(dst_layername, srs=srs)
    raster_field = ogr.FieldDefn('id', type_mapping[srcband.DataType])
    dst_layer.CreateField(raster_field)
    gdal.Polygonize(srcband, srcband, dst_layer, 0, [], callback=None)
    del img, ds, srcband, dst_ds, dst_layer


# convert images in a selected folder to shapefiles 

def get_idx_as_shp(self, path, gt, wkt):
        '''
        Exports a Shapefile containing the locations of the extracted
        endmembers. Assumes the coordinates are in decimal degrees.
        '''
        coords = pixel_to_xy(self.get_idx(), gt=gt, wkt=wkt, dd=True)

        driver = ogr.GetDriverByName('ESRI Shapefile')
        ds = driver.CreateDataSource(path)
        srs = osr.SpatialReference()
        srs.ImportFromEPSG(4326)

        layer = ds.CreateLayer(path.split('.')[0], srs, ogr.wkbPoint)
        for pair in coords:
            feature = ogr.Feature(layer.GetLayerDefn())

            # Create the point from the Well Known Text
            point = ogr.CreateGeometryFromWkt('POINT(%f %f)' %  pair)
            feature.SetGeometry(point) # Set the feature geometry
            layer.CreateFeature(feature) # Create the feature in the layer
            feature.Destroy() # Destroy the feature to free resources

        # Destroy the data source to free resources
        ds.Destroy() 

def CreateMultiBandGeoTiff(OutPath, Array):
    '''
    Author: Jake Shermeyer
    Array has shape:
        Channels, Y, X?
    '''
    driver = gdal.GetDriverByName('GTiff')
    DataSet = driver.Create(OutPath, Array.shape[2], Array.shape[1],
                            Array.shape[0], gdal.GDT_Byte,
                            ['COMPRESS=LZW'])
    for i, image in enumerate(Array, 1):
        DataSet.GetRasterBand(i).WriteArray(image)
    del DataSet

    return OutPath


############################################################################### 

def save_state_gauges(out_fn, bbox=None):
    """Save stream gauge data to a geojson file."""
    url = 'http://gis.srh.noaa.gov/arcgis/services/ahps_gauges/' + \
          'MapServer/WFSServer'
    parms = {
        'version': '1.1.0',
        'typeNames': 'ahps_gauges:Observed_River_Stages',
        'srsName': 'urn:ogc:def:crs:EPSG:6.9:4326',
    }
    if bbox:
        parms['bbox'] = bbox
    try:
        request = 'WFS:{0}?{1}'.format(url, urllib.urlencode(parms))
    except:
        request = 'WFS:{0}?{1}'.format(url, urllib.parse.urlencode(parms))
    wfs_ds = ogr.Open(request)
    if wfs_ds is None:
        raise RuntimeError('Could not open WFS.')
    wfs_lyr = wfs_ds.GetLayer(0)

    driver = ogr.GetDriverByName('GeoJSON')
    if os.path.exists(out_fn):
        driver.DeleteDataSource(out_fn)
    json_ds = driver.CreateDataSource(out_fn)
    json_ds.CopyLayer(wfs_lyr, '') 

def create_shapedir(self):
        drv = ogr.GetDriverByName("ESRI Shapefile")
        shp = drv.CreateDataSource(self.path)
        lyr = shp.CreateLayer("nodes", None, ogr.wkbPoint)
        feature = ogr.Feature(lyr.GetLayerDefn())
        feature.SetGeometry(None)
        lyr.CreateFeature(feature)
        feature.Destroy() 

def upsample_like_that(self, ext_img, method='bilinear', no_data_value=None):
        """Use gdal.ReprojectImage to upsample the object so that it looks
        like the geoio image object passed in as the ext_img argument.  Method
        can be those listed in GeoImage.upsample.
        """

        if ext_img.meta.resolution[0]>self.meta.resolution[0]:
            raise ValueError('The requested resolution is not at or higher '
                             'than the base object.  Use downsample or '
                             'resample methods.')

        # Set up destination image in memory
        drv = gdal.GetDriverByName('MEM')
        dst = drv.Create('',
                         ext_img.shape[2],
                         ext_img.shape[1],
                         ext_img.shape[0],
                         ext_img.meta.gdal_dtype)
        dst.SetGeoTransform(ext_img.meta.geo_transform)
        dst.SetProjection(ext_img.meta.projection_string)
        if no_data_value is not None:
            blist = [x + 1 for x in range(ext_img.shape[0])]
            [dst.GetRasterBand(x).SetNoDataValue(no_data_value) for x in blist]

        # Run resample, pull data, and del the temp gdal object.
        data = self._upsample_from_gdalobj(self.get_gdal_obj(),dst,method)
        del dst

        return data 

def create_shp_datasource(self):
        self.shp_layers = {}
        if not os.path.exists(self.dir_path):
            os.makedirs(self.dir_path)
        driver = ogr.GetDriverByName('ESRI Shapefile')
        return driver.CreateDataSource(self.dir_path) 

def write_shp(path, gjson):
    proj = osr.SpatialReference(osr.SRS_WKT_WGS84)
    g = ogr.CreateGeometryFromJson(json.dumps(gjson))
    vdriver = ogr.GetDriverByName('ESRI Shapefile')
    ds = vdriver.CreateDataSource(path)
    layer = ds.CreateLayer('', proj, g.GetGeometryType())
    featdef = layer.GetLayerDefn()
    feature = ogr.Feature(featdef)
    feature.SetGeometry(g)
    layer.CreateFeature(feature)
    feature.Destroy()
    ds.Destroy() 

def parsed(args):
    kml_file = args.geo

    def kml2geojson(kml_file):
        drv = ogr.GetDriverByName("KML")
        kml_ds = drv.Open(kml_file)
        for kml_lyr in kml_ds:
            for feat in kml_lyr:
                outfile = feat.ExportToJson()
                geom2 = str(outfile).replace(", 0.0", "")
                with open(args.loc + "./kmlout.geojson", "w") as csvfile:
                    writer = csv.writer(csvfile)
                    writer.writerow([geom2])

    kml2geojson(args.geo)
    raw = open(args.loc + "./kmlout.geojson")
    for line in raw:
        fields = line.strip().split(":")[3]
        f2 = fields.strip().split("}")[0]
        filenames = (
            p1
            + f2
            + p2
            + str(args.start)
            + p3
            + str(args.end)
            + p4
            + p5
            + str(args.cloud)
            + p6
        )
    with open(args.loc + "./aoi.json", "w") as outfile:
        outfile.write(filenames)
        outfile.close() 

def layers_to_feature_dataset(ds_name, gdb_fn, dataset_name):
    """Copy layers to a feature dataset in a file geodatabase."""

    # Open the input datasource.
    in_ds = ogr.Open(ds_name)
    if in_ds is None:
        raise RuntimeError('Could not open datasource')

    # Open the geodatabase or create it if it doesn't exist.
    gdb_driver = ogr.GetDriverByName('FileGDB')
    if os.path.exists(gdb_fn):
        gdb_ds = gdb_driver.Open(gdb_fn, 1)
    else:
        gdb_ds = gdb_driver.CreateDataSource(gdb_fn)
    if gdb_ds is None:
        raise RuntimeError('Could not open file geodatabase')

    # Create an option list so the feature classes will be
    # saved in a feature dataset.
    options = ['FEATURE_DATASET=' + dataset_name]

    # Loop through the layers in the input datasource and copy
    # each one into the geodatabase.
    for i in range(in_ds.GetLayerCount()):
        lyr = in_ds.GetLayer(i)
        lyr_name = lyr.GetName()
        print('Copying ' + lyr_name + '...')
        gdb_ds.CopyLayer(lyr, lyr_name, options) 

def setUp(self):

        def createlayer(driver):
            lyr = shp.CreateLayer("edges", None, ogr.wkbLineString)
            namedef = ogr.FieldDefn("Name", ogr.OFTString)
            namedef.SetWidth(32)
            lyr.CreateField(namedef)
            return lyr

        drv = ogr.GetDriverByName("ESRI Shapefile")

        testdir = os.path.join(tempfile.gettempdir(), 'shpdir')
        shppath = os.path.join(tempfile.gettempdir(), 'tmpshp.shp')

        self.deletetmp(drv, testdir, shppath)
        os.mkdir(testdir)

        shp = drv.CreateDataSource(shppath)
        lyr = createlayer(shp)
        self.names = ['a', 'b', 'c', 'c']  # edgenames
        self.paths = ([(1.0, 1.0), (2.0, 2.0)],
                      [(2.0, 2.0), (3.0, 3.0)],
                      [(0.9, 0.9), (4.0, 0.9), (4.0, 2.0)])

        self.simplified_names = ['a', 'b', 'c']  # edgenames
        self.simplified_paths = ([(1.0, 1.0), (2.0, 2.0)],
                                 [(2.0, 2.0), (3.0, 3.0)],
                                 [(0.9, 0.9), (4.0, 2.0)])

 
        for path, name in zip(self.paths, self.names):
            feat = ogr.Feature(lyr.GetLayerDefn())
            g = ogr.Geometry(ogr.wkbLineString)
            for p in path:
                g.AddPoint_2D(*p)
            feat.SetGeometry(g)
            feat.SetField("Name", name)
            lyr.CreateFeature(feat)
        self.shppath = shppath
        self.testdir = testdir
        self.drv = drv 

def TestLand(lon, lat):
    latlon = osr.SpatialReference()
    latlon.ImportFromEPSG(4326)

    # create a point

    pt = ogr.Geometry(ogr.wkbPoint)
    pt.SetPoint_2D(0, lon, lat)

    # read shapefile
    shapefile = "land_polygons_osm/simplified_land_polygons.shp"
    driver = ogr.GetDriverByName("ESRI Shapefile")
    dataSource = driver.Open(shapefile, 0)
    layer = dataSource.GetLayer()
    targetProj = layer.GetSpatialRef()
    land = False

    # conversion to shapefile projection
    transform = osr.CoordinateTransformation(latlon, targetProj)
    pt.Transform(transform)

    # search point in layers
    for feature in layer:
        geom = feature.GetGeometryRef()
        if geom.Contains(pt):
            land = True
            break

    return land


##################################### Lecture de fichier de parametres "Mot_clé=Valeur" 

def saveToShape(self, array, srs, outShapeFile):
        # Parse a delimited text file of volcano data and create a shapefile
        # use a dictionary reader so we can access by field name
        # set up the shapefile driver
        outDriver = ogr.GetDriverByName('ESRI Shapefile')

        # create the data source
        if os.path.exists(outShapeFile):
            outDriver.DeleteDataSource(outShapeFile)
        # Remove output shapefile if it already exists

        # options = ['SPATIALITE=YES'])
        ds = outDriver.CreateDataSource(outShapeFile)

        # create the spatial reference, WGS84

        lyrout = ds.CreateLayer('randomSubset', srs)
        fields = [
            array[1].GetFieldDefnRef(i).GetName() for i in range(
                array[1].GetFieldCount())]

        for f in fields:
            field_name = ogr.FieldDefn(f, ogr.OFTString)
            field_name.SetWidth(24)
            lyrout.CreateField(field_name)

        for k in array:
            lyrout.CreateFeature(k)

        # Save and close the data source
        ds = None 

def delete_shapedir(self):
        drv = ogr.GetDriverByName("ESRI Shapefile")
        if os.path.exists(self.path):
            drv.DeleteDataSource(self.path) 

def delete_shapedir(self):
        drv = ogr.GetDriverByName("ESRI Shapefile")
        if os.path.exists(self.path):
            drv.DeleteDataSource(self.path) 

def featuresToFile(features, dst_drv, dst_name, dst_srs, layer_name=None,
                   geomtype=None, overwrite=True):
    if not features:  # features is empty list
        print("No Features Created")
        return

    drv = ogr.GetDriverByName(dst_drv)
    if drv is None:
        print("Driver not available ({})".format(dst_drv))
        return

    dsrc = drv.CreateDataSource(dst_name)
    if dsrc is None:
        print("DataSource creation failed")
        return

    if not geomtype:
        f0 = features[0]
        geomref = features[0].GetGeometryRef()
        if geomref is not None:
            geomtype = geomref.GetGeometryType()
        else:
            return

    layer = dsrc.CreateLayer(layer_name, srs=dst_srs, geom_type=geomtype)

    # Create the fields for the new file
    for i in range(features[0].GetFieldCount()):
        fieldDef = features[0].GetFieldDefnRef(i)
        if "List" in ogr.GetFieldTypeName(fieldDef.GetType()):
            t = ogr.GetFieldTypeName(fieldDef.GetType())[:-4]
            if t == "String":
                fieldDef = ogr.FieldDefn(fieldDef.GetName(), ogr.OFTString)
            elif t == "Integer":
                fieldDef = ogr.FieldDefn(fieldDef.GetName(), ogr.OFTInteger)

        layer.CreateField(fieldDef)

    # print layer_name
    for feature in features:
        layer.CreateFeature(feature) 

def saveToShape(array, srs, outShapeFile):
    # Parse a delimited text file of volcano data and create a shapefile
    # use a dictionary reader so we can access by field name
    # set up the shapefile driver
    import ogr
    outDriver = ogr.GetDriverByName('ESRI Shapefile')

    # create the data source
    if os.path.exists(outShapeFile):
        outDriver.DeleteDataSource(outShapeFile)
    # Remove output shapefile if it already exists

    # options = ['SPATIALITE=YES'])
    ds = outDriver.CreateDataSource(outShapeFile)

    # create the spatial reference, WGS84

    lyrout = ds.CreateLayer('randomSubset', srs, ogr.wkbPoint)
    fields = [array[1].GetFieldDefnRef(i).GetName()
              for i in range(array[1].GetFieldCount())]

    for i, f in enumerate(fields):
        field_name = ogr.FieldDefn(f, array[1].GetFieldDefnRef(i).GetType())
        field_name.SetWidth(array[1].GetFieldDefnRef(i).GetWidth())
        lyrout.CreateField(field_name)

    for k in array:
        lyrout.CreateFeature(k)

    # Save and close the data source
    ds = None 

def __init__(self, iface, tableName, user, apiKey, owner=None, sql=None, geoJSON=None,
                 filterByExtent=False, spatiaLite=None, readonly=False, multiuser=False, isSQL=False):
        self.iface = iface
        self.user = user
        self._apiKey = apiKey
        self.layerType = 'ogr'
        self.owner = owner
        self.isSQL = isSQL
        self.multiuser = multiuser
        # SQLite available?
        driverName = "SQLite"
        sqLiteDrv = ogr.GetDriverByName(driverName)
        self.database_path = QgisCartoDB.CartoDBPlugin.PLUGIN_DIR + '/db/database.sqlite'
        self.datasource = sqLiteDrv.Open(self.database_path, True)
        self.layerName = tableName
        self.cartoTable = tableName
        self.readonly = False
        self._deletedFeatures = []
        self.sql = sql

        self.forceReadOnly = False or readonly

        if sql is None:
            sql = 'SELECT * FROM ' + self._schema() + self.cartoTable
            if filterByExtent:
                extent = self.iface.mapCanvas().extent()
                sql = sql + " WHERE ST_Intersects(ST_GeometryFromText('{}', 4326), the_geom)".format(extent.asWktPolygon())
        else:
            self.forceReadOnly = True

        self._loadData(sql, geoJSON, spatiaLite) 

def output_ogr(self, output):
    driver = ogr.GetDriverByName( 'ESRI Shapefile' )
    if os.path.exists( output['file_name'] ):
      driver.DeleteDataSource( output['file_name'] )
    source = driver.CreateDataSource( output['file_name'] )
    layer = source.CreateLayer( self.layer_dfn.GetName(),
                                geom_type = self.layer_dfn.GetGeomType(),
                                srs = self.layer.GetSpatialRef() )

    for field in self.fields:
      fd = ogr.FieldDefn( str(field['name']), field['type'] )
      fd.SetWidth( field['width'] )
      if 'precision' in field:
        fd.SetPrecision( field['precision'] )
      layer.CreateField( fd )

    for geometry in self.geometries:
      if geometry.geom is not None:
        feature = ogr.Feature( feature_def = layer.GetLayerDefn() )
        for index, field in enumerate(self.fields):
          if field['name'] in geometry.properties:
            feature.SetField( index, geometry.properties[field['name']].encode('utf-8') )
          else:
            feature.SetField( index, '' )
        feature.SetGeometryDirectly(
          ogr.CreateGeometryFromWkb(
            shapely.wkb.dumps(
              geometry.geom
            )
          )
        )
        layer.CreateFeature( feature )
        feature.Destroy()

    source.Destroy() 

def write_geometry(geometry, out_path, srs_id=4326):
    """
    Saves the geometry in an ogr.Geometry object to a shapefile.

    Parameters
    ----------
    geometry
        An ogr.Geometry object
    out_path
        The location to save the output shapefile
    srs_id
        The projection of the output shapefile. Can be an EPSG number or a WKT string.

    Notes
    -----
    The shapefile consists of one layer named 'geometry'.


    """
    # TODO: Fix this needing an extra filepath on the end
    driver = ogr.GetDriverByName("ESRI Shapefile")
    data_source = driver.CreateDataSource(out_path)
    srs = osr.SpatialReference()
    if type(srs_id) is int:
        srs.ImportFromEPSG(srs_id)
    if type(srs_id) is str:
        srs.ImportFromWkt(srs_id)
    layer = data_source.CreateLayer(
        "geometry",
        srs,
        geom_type=geometry.GetGeometryType())
    feature_def = layer.GetLayerDefn()
    feature = ogr.Feature(feature_def)
    feature.SetGeometry(geometry)
    layer.CreateFeature(feature)
    data_source.FlushCache()
    data_source = None 

def make_class_dict(path):
    # Set up dict to save Xs and Ys
    driver = ogr.GetDriverByName('ESRI Shapefile')
    data_source = driver.Open(path, 0)
    if data_source is None:
        report_and_exit("File read failed: %s", path)

    layer = data_source.GetLayer(0)
    class_labels = []
    data = []

    for feature in layer:
        try:
            var1 = float(feature.GetField('NFDI_mag')) 
   	    var2 = float(feature.GetField('NFDI_rmse'))
	    var3 = float(feature.GetField('NFDI_sin')) 
	    var4 = float(feature.GetField('NFDI_cos')) 
	    var5 = float(feature.GetField('Gv_mag')) 
	    var6 = float(feature.GetField('Shade_mag'))
	    var7 = float(feature.GetField('NPV_mag')) 
	    var8 = float(feature.GetField('Soil_mag')) 
	    label = feature.GetField('class')
        except:
            continue 
        class_labels.append(label)
        data.append([var1, var2, var3, var4, var5, var6, var7, var8])
#        data.append([var1, var3, var4, var5, var6, var7, var8])

    return class_labels, data 

def make_class_dict(path):
    # Set up dict to save Xs and Ys
    driver = ogr.GetDriverByName('ESRI Shapefile')
#    data_source = gdal.OpenEx(vector_data_path, gdal.OF_VECTOR)
    data_source = driver.Open(path, 0)
    if data_source is None:
        report_and_exit("File read failed: %s", path)

    layer = data_source.GetLayer(0)
    class_labels = []
    data = []

    for feature in layer:
        try:
            var1 = float(feature.GetField('NFDI_mag')) 
#   	    var2 = float(feature.GetField('NFDI_rmse'))
	    var3 = float(feature.GetField('NFDI_sin')) 
	    var4 = float(feature.GetField('NFDI_cos')) 
	    var5 = float(feature.GetField('Gv_mag')) 
	    var6 = float(feature.GetField('Shade_mag'))
	    var7 = float(feature.GetField('NPV_mag')) 
	    var8 = float(feature.GetField('Soil_mag')) 
	    label = feature.GetField('class')
        except:
            continue 
        class_labels.append(label)
        #data.append([var1, var2, var3, var4, var5, var6, var7, var8])
        data.append([var1, var3, var4, var5, var6, var7, var8])

    return class_labels, data 

def write_geotiff(fname, data, geo_transform, projection, data_type=gdal.GDT_Byte):

    """
    Create a GeoTIFF file with the given data.
    :param fname: Path to a directory with shapefiles
    :param data: Number of rows of the result
    :param geo_transform: Returned value of 
	gdal.Dataset.GetGeoTransform (coefficients for transforming between 
	pixel/line (P,L) raster space, and projection coordinates (Xp,Yp) space.
    :param projection: Projection definition string (Returned by 
	gdal.Dataset.GetProjectionRef)
    """

    driver = gdal.GetDriverByName('GTiff')
    rows, cols = data.shape
    dataset = driver.Create(fname, cols, rows, 1, data_type)
    dataset.SetGeoTransform(geo_transform)
    dataset.SetProjection(projection)
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)

    metadata = {
        'TIFFTAG_COPYRIGHT': 'CC BY 4.0',
        'TIFFTAG_DOCUMENTNAME': 'classification',
        'TIFFTAG_IMAGEDESCRIPTION': 'Supervised classification.',
        'TIFFTAG_SOFTWARE': 'Python, GDAL, scikit-learn'
    }
    dataset.SetMetadata(metadata)

    dataset = None  # Close the file
    return 

def create_mask_from_vector(vector_data_path, cols, rows, geo_transform, 
			    projection, target_value=1,
                            output_fname='', dataset_format='MEM'):

    """
    Rasterize the given vector (wrapper for gdal.RasterizeLayer). 
    Return a gdal.Dataset.
    :param vector_data_path: Path to a shapefile
    :param cols: Number of columns of the result
    :param rows: Number of rows of the result
    :param geo_transform: Returned value of gdal.Dataset.GetGeoTransform 
	(coefficients for transforming between pixel/line (P,L) raster space,
	 and projection coordinates (Xp,Yp) space.
    :param projection: Projection definition string (Returned by 
	gdal.Dataset.GetProjectionRef)
    :param target_value: Pixel value for the pixels. Must be a valid 
	gdal.GDT_UInt16 value.
    :param output_fname: If the dataset_format is GeoTIFF, this is the output 
	file name
    :param dataset_format: The gdal.Dataset driver name. [default: MEM]
    """

    driver = ogr.GetDriverByName('ESRI Shapefile')
    data_source = driver.Open(vector_data_path, 0)
    if data_source is None:
        report_and_exit("File read failed: %s", vector_data_path)
    layer = data_source.GetLayer(0)
    driver = gdal.GetDriverByName(dataset_format)
    target_ds = driver.Create(output_fname, cols, rows, 1, gdal.GDT_UInt16)
    target_ds.SetGeoTransform(geo_transform)
    target_ds.SetProjection(projection)
    gdal.RasterizeLayer(target_ds, [1], layer, burn_values=[target_value])
    return target_ds 

def polygonize(img,shp_path):
    # mapping between gdal type and ogr field type
    type_mapping = {gdal.GDT_Byte: ogr.OFTInteger,
                    gdal.GDT_UInt16: ogr.OFTInteger,
                    gdal.GDT_Int16: ogr.OFTInteger,
                    gdal.GDT_UInt32: ogr.OFTInteger,
                    gdal.GDT_Int32: ogr.OFTInteger,
                    gdal.GDT_Float32: ogr.OFTReal,
                    gdal.GDT_Float64: ogr.OFTReal,
                    gdal.GDT_CInt16: ogr.OFTInteger,
                    gdal.GDT_CInt32: ogr.OFTInteger,
                    gdal.GDT_CFloat32: ogr.OFTReal,
                    gdal.GDT_CFloat64: ogr.OFTReal}

    ds = gdal.Open(img)
    prj = ds.GetProjection()
    srcband = ds.GetRasterBand(1)

    dst_layername = "Shape"
    drv = ogr.GetDriverByName("ESRI Shapefile")
    dst_ds = drv.CreateDataSource(shp_path)
    srs = osr.SpatialReference(wkt=prj)

    dst_layer = dst_ds.CreateLayer(dst_layername, srs=srs)
    # raster_field = ogr.FieldDefn('id', type_mapping[srcband.DataType])
    raster_field = ogr.FieldDefn('id', type_mapping[gdal.GDT_Int32])
    dst_layer.CreateField(raster_field)
    gdal.Polygonize(srcband, srcband, dst_layer, 0, [], callback=None)
    del img, ds, srcband, dst_ds, dst_layer


# extract sinks from dem 

def array2raster(rasterfn,newRasterfn,array,driver_name = "ENVI", noDataValue = -9999):
    """Takes an array and writes to a GDAL compatible raster. It needs another raster to map the dimensions.

    Args:
        FileName (str): The filename (with path and extension) of a raster that has the same dimensions as the raster to be written.
        newRasterfn (str): The filename (with path and extension) of the new raster.
        array (np.array): The array to be written
        driver_name (str): The type of raster to write. Default is ENVI since that is the LSDTOpoTools format
        noDataValue (float): The no data value

    Return:
        np.array: A numpy array with the data from the raster.

    Author: SMM
    """

    raster = gdal.Open(rasterfn)
    geotransform = raster.GetGeoTransform()
    originX = geotransform[0]
    originY = geotransform[3]
    pixelWidth = geotransform[1]
    pixelHeight = geotransform[5]
    cols = raster.RasterXSize
    rows = raster.RasterYSize

    driver = gdal.GetDriverByName(driver_name)
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((originX, pixelWidth, 0, originY, 0, pixelHeight))
    outRaster.GetRasterBand(1).SetNoDataValue( noDataValue )
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromWkt(raster.GetProjectionRef())
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()
#==============================================================================