First draft of a classification_to_shapefile routine

classification_to_shapefile.py

####################################################################################################################################
##Test

import matplotlib 
# matplotlib.use('Agg') # non interactive
# %matplotlib qt
# from   matplotlib import pyplot as plt
import numpy as np

plot_matrix =  np.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   	                   [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   	                   [ 0, 1, 1, 1, 1, 0, 2, 2, 2, 2, 0, 1, 1, 1, 0, 2, 0, 0, 0],
   	                   [ 0, 1, 0, 0, 0, 0, 0, 2, 0, 2, 0, 1, 0, 1, 0, 2, 0, 0, 0],
   	                   [ 0, 1, 0, 1, 1, 0, 0, 2, 0, 2, 0, 1, 1, 1, 0, 2, 0, 0, 0],
   	                   [ 0, 1, 0, 0, 1, 0, 0, 2, 0, 2, 0, 1, 0, 1, 0, 2, 0, 0, 0],
   	                   [ 0, 1, 1, 1, 1, 0, 2, 2, 2, 2, 0, 1, 0, 1, 0, 2, 2, 2, 0],
   	                   [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   	                   [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   	                   [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

# class 0 represent No Data, so we have 3-1 classes
classes = np.unique(plot_matrix)[1:]


# Some Fake GeoTIFF georereference
img_extent = (-180, 180, -90, 90)
xmin,xmax,ymin,ymax = img_extent
nrows,ncols = plot_matrix.shape
xres=(xmax-xmin)/float(ncols)
yres=(ymax-ymin)/float(nrows)
geotransform = (xmin,xres,0,ymax,0, -yres)

out_path = '/Users/damo_ma/Desktop/test_gdal/'


##########################################################################
# create 1 bands in raster file in Memory

from osgeo import gdal
from osgeo import ogr
from osgeo import osr
# from osgeo import gdalnumeric

import os 

# create a Raster Layer in memory
driver = gdal.GetDriverByName( 'MEM' )
gdal_datasource = driver.Create( '',ncols, nrows, 1,gdal.GDT_Int16)  
gdal_datasource.SetGeoTransform(geotransform)
srs = osr.SpatialReference()
srs.ImportFromEPSG(4326) #WGS04 lat long
gdal_datasource.SetProjection( srs.ExportToWkt() )

raster_band = gdal_datasource.GetRasterBand(1)
raster_band.SetNoDataValue(plot_matrix.astype(int)[0,0])
raster_band.WriteArray(plot_matrix.astype(int))  # it start from band 1!!

# # Create GDAL dataset for array, and set georeferencing.
# gdal_datasource = gdalnumeric.OpenArray( plot_matrix.astype(int) )
# gdal_datasource.SetGeoTransform( geotransform )
# #create the spatial reference, WGS04
# srs = osr.SpatialReference()
# srs.ImportFromEPSG(4326)
# gdal_datasource.SetProjection( srs.ExportToWkt() )

##########################################################################
# create a Vector Layer in memory
drv = ogr.GetDriverByName( 'Memory' )
ogr_datasource = drv.CreateDataSource( 'out' )

input_layer     = ogr_datasource.CreateLayer('polygonized', srs, ogr.wkbPolygon )
layerDefinition = input_layer.GetLayerDefn()

field_defn = ogr.FieldDefn(b'class', ogr.OFTInteger) 
input_layer.CreateField(field_defn) 

print input_layer.GetFeatureCount()

gdal.Polygonize( raster_band, raster_band.GetMaskBand(), input_layer,  0)
print input_layer.GetFeatureCount()

########################################################################################################################################
# create 1 bands in raster file

driver = ogr.GetDriverByName("ESRI Shapefile")

for i in range(layerDefinition.GetFieldCount()):
    print 'custom Field Name : %s ' % layerDefinition.GetFieldDefn(i).GetName()

# select the field name to use for merge the polygon

field_name = layerDefinition.GetFieldDefn(0).GetName()

# it should be just one field!
for i in classes:
   input_layer.SetAttributeFilter("%s = %s" % (field_name,i))
   print 'field %s == %s : %s features ' % (field_name,i,input_layer.GetFeatureCount())

out_shp = 'testone.shp'

# Remove output shapefile if it already exists
if os.path.exists(out_path + out_shp):
    driver.DeleteDataSource(out_path + out_shp)
out_datasource = driver.CreateDataSource( out_path + out_shp )

print out_path + out_shp

# create a new layer with wkbMultiPolygon
multi_layer = out_datasource.CreateLayer("merged", input_layer.GetSpatialRef(), ogr.wkbMultiPolygon) # add a Layer, same Spatial Ref as input, MultiPolygon
# Add the fields we're interested in
field_field_name = ogr.FieldDefn(field_name, ogr.OFTInteger) # add a Field named Class
multi_layer.CreateField(field_defn) 


for i in classes:
   input_layer.SetAttributeFilter("%s = %s" % (field_name,i))
   print 'field %s == %s : %s features ' % (field_name,i,input_layer.GetFeatureCount())

   # # create a new layer with wkbMultiPolygon
   # multi_layer = out_datasource.CreateLayer("class %s"  % i, input_layer.GetSpatialRef(), ogr.wkbMultiPolygon) # add a Layer
   # field_name = ogr.FieldDefn("class", ogr.OFTInteger) # add a Field
   # multi_layer.CreateField(field_name)

   multi_feature = ogr.Feature(multi_layer.GetLayerDefn()) # generate a feature
   multipolygon  = ogr.Geometry(ogr.wkbMultiPolygon) # generate a polygon based on layer unique Geometry definition
   for feature in input_layer:
      multipolygon.AddGeometry(feature.geometry()) #aggregate all the input geometry sharing the class value i

   multi_feature.SetGeometry(multipolygon) # add the merged geoemtry to the current feature
   multi_feature.SetField(field_name, i) # set the field of the current feature
   multi_layer.CreateFeature(multi_feature) # add the current feature to the layer
   multi_feature.Destroy() # desrtroy the current feature

gdal_datasource = None
ogr_datasource  = None
out_datasource  = None