drawing polygon

geosentinel/rasterWkt.py

@@ -34,32 +34,31 @@ def rasterWkt(wkt, inputfile, outputfile):
 	print (outputfile)
 	intput = gdal.Open(inputfile)
 	rows, cols, geotransform = intput.RasterYSize, intput.RasterXSize, intput.GetGeoTransform()
+	sr = osr.SpatialReference()
+	sr.ImportFromWkt(intput.GetProjection())
 	# Read the input bands as numpy arrays.
 
 
 	driverTiff = gdal.GetDriverByName('GTiff')
 	output = driverTiff.Create(outputfile, cols, rows, 3, gdal.GDT_Byte)
-	print ("BANDS")
-	print (output.RasterCount)
+	polygonRaster = driverTiff.Create("polygonTmp.tiff", cols, rows, 1, gdal.GDT_Byte)
 
-	for b in range(1, output.RasterCount+1):
-		np_input = intput.GetRasterBand(b).ReadAsArray(0,0,cols,rows)
-		output.GetRasterBand(b).SetNoDataValue(-99)
-		output.GetRasterBand(b).WriteArray(np_input)
-		output.GetRasterBand(b).FlushCache()
 
+	#Set up polygon raster
+	polygonRaster.GetRasterBand(b).SetNoDataValue(-99)
+	polygonRaster.GetRasterBand(b).WriteArray(np.zeros(cols*rows))
+	polygonRaster.GetRasterBand(b).FlushCache()
+	polygonRaster.SetGeoTransform(geotransform)
 
-	output.SetGeoTransform(geotransform)
 
-	sr = osr.SpatialReference()
-	sr.ImportFromWkt(intput.GetProjection())
-	output.SetProjection(sr.ExportToWkt())
 
 
 	# shapefile
 	driver = ogr.GetDriverByName("ESRI Shapefile")
 	data_source = driver.CreateDataSource("myShape.shp")
 
+
+	#Burn Polygon
 	srs = osr.SpatialReference()
 	srs.ImportFromEPSG(4326)
 	layer = data_source.CreateLayer("wkt", srs, ogr.wkbPolygon)
@@ -70,10 +69,26 @@ def rasterWkt(wkt, inputfile, outputfile):
 	feature.SetGeometry(geometry)
 	layer.CreateFeature(feature)
 
-	gdal.RasterizeLayer(output, [1,2,3], layer, burn_values=[255,1,1])
+	gdal.RasterizeLayer(polygonRaster, [1], layer, burn_values=[255])
+
+	#Draw output image
+	np_polygon = polygonRaster.GetRasterBand(1).ReadAsArray(0,0,cols,rows)
+	for b in range(1, output.RasterCount+1):
+		np_input = intput.GetRasterBand(b).ReadAsArray(0,0,cols,rows) | polygonRaster
+		output.GetRasterBand(b).SetNoDataValue(-99)
+		output.GetRasterBand(b).WriteArray(np_input)
+		output.GetRasterBand(b).FlushCache()
+
+	output.SetGeoTransform(geotransform)
+	output.SetProjection(sr.ExportToWkt())
+
+
+	#Close files
 	feature = None
 	data_source = None
 	output = None
+	polygonRaster = None
+
 
 def main(argv):
 	rasterWkt(argv[1], argv[2], argv[3])