# -*- coding: utf-8 -*-
"""
@author lzugis
@date 2017-06-02
@brief 利用shp裁剪影像
"""
 
from osgeo import gdal, gdalnumeric, ogr
from PIL import Image, ImageDraw
import os
import operator
 
gdal.UseExceptions()
 
# This function will convert the rasterized clipper shapefile
# to a mask for use within GDAL.
def imageToArray(i):
 """
 Converts a Python Imaging Library array to a
 gdalnumeric image.
 """
 a=gdalnumeric.fromstring(i.tobytes(),'b')
 a.shape=i.im.size[1], i.im.size[0]
 return a
 
def arrayToImage(a):
 """
 Converts a gdalnumeric array to a
 Python Imaging Library Image.
 """
 i=Image.frombytes('L',(a.shape[1],a.shape[0]),
   (a.astype('b')).tobytes())
 return i
 
def world2Pixel(geoMatrix, x, y):
 """
 Uses a gdal geomatrix (gdal.GetGeoTransform()) to calculate
 the pixel location of a geospatial coordinate
 """
 ulX = geoMatrix[0]
 ulY = geoMatrix[3]
 xDist = geoMatrix[1]
 pixel = int((x - ulX) / xDist)
 line = int((ulY - y) / xDist)
 return (pixel, line)
 
#
# EDIT: this is basically an overloaded
# version of the gdal_array.OpenArray passing in xoff, yoff explicitly
# so we can pass these params off to CopyDatasetInfo
#
def OpenArray( array, prototype_ds = None, xoff=0, yoff=0 ):
 ds = gdal.Open( gdalnumeric.GetArrayFilename(array) )
 
 if ds is not None and prototype_ds is not None:
  if type(prototype_ds).__name__ == 'str':
   prototype_ds = gdal.Open( prototype_ds )
  if prototype_ds is not None:
   gdalnumeric.CopyDatasetInfo( prototype_ds, ds, xoff=xoff, yoff=yoff )
 return ds
 
def histogram(a, bins=range(0,256)):
 """
 Histogram function for multi-dimensional array.
 a = array
 bins = range of numbers to match
 """
 fa = a.flat
 n = gdalnumeric.searchsorted(gdalnumeric.sort(fa), bins)
 n = gdalnumeric.concatenate([n, [len(fa)]])
 hist = n[1:]-n[:-1]
 return hist
 
def stretch(a):
 """
 Performs a histogram stretch on a gdalnumeric array image.
 """
 hist = histogram(a)
 im = arrayToImage(a)
 lut = []
 for b in range(0, len(hist), 256):
  # step size
  step = reduce(operator.add, hist[b:b+256]) / 255
  # create equalization lookup table
  n = 0
  for i in range(256):
   lut.append(n / step)
   n = n + hist[i+b]
  im = im.point(lut)
 return imageToArray(im)
 
def main( shapefile_path, raster_path ):
 # Load the source data as a gdalnumeric array
 srcArray = gdalnumeric.LoadFile(raster_path)
 
 # Also load as a gdal image to get geotransform
 # (world file) info
 srcImage = gdal.Open(raster_path)
 geoTrans = srcImage.GetGeoTransform()
 
 # Create an OGR layer from a boundary shapefile
 shapef = ogr.Open(shapefile_path)
 lyr = shapef.GetLayer( os.path.split( os.path.splitext( shapefile_path )[0] )[1] )
 poly = lyr.GetNextFeature()
 
 # Convert the layer extent to image pixel coordinates
 minX, maxX, minY, maxY = lyr.GetExtent()
 ulX, ulY = world2Pixel(geoTrans, minX, maxY)
 lrX, lrY = world2Pixel(geoTrans, maxX, minY)
 
 # Calculate the pixel size of the new image
 pxWidth = int(lrX - ulX)
 pxHeight = int(lrY - ulY)
 
 clip = srcArray[:, ulY:lrY, ulX:lrX]
 
 #
 # EDIT: create pixel offset to pass to new image Projection info
 #
 xoffset = ulX
 yoffset = ulY
 print "Xoffset, Yoffset = ( %f, %f )" % ( xoffset, yoffset )
 
 # Create a new geomatrix for the image
 geoTrans = list(geoTrans)
 geoTrans[0] = minX
 geoTrans[3] = maxY
 
 # Map points to pixels for drawing the
 # boundary on a blank 8-bit,
 # black and white, mask image.
 points = []
 pixels = []
 geom = poly.GetGeometryRef()
 pts = geom.GetGeometryRef(0)
 for p in range(pts.GetPointCount()):
  points.append((pts.GetX(p), pts.GetY(p)))
 for p in points:
  pixels.append(world2Pixel(geoTrans, p[0], p[1]))
 rasterPoly = Image.new("L", (pxWidth, pxHeight), 1)
 rasterize = ImageDraw.Draw(rasterPoly)
 rasterize.polygon(pixels, 0)
 mask = imageToArray(rasterPoly)
 
 # Clip the image using the mask
 clip = gdalnumeric.choose(mask,   (clip, 0)).astype(gdalnumeric.uint8)
 
 # This image has 3 bands so we stretch each one to make them
 # visually brighter
 for i in range(3):
  clip[i,:,:] = stretch(clip[i,:,:])
 
 # Save new tiff
 #
 # EDIT: instead of SaveArray, let's break all the
 # SaveArray steps out more explicity so
 # we can overwrite the offset of the destination
 # raster
 #
 ### the old way using SaveArray
 #
 # gdalnumeric.SaveArray(clip, "OUTPUT.tif", format="GTiff", prototype=raster_path)
 #
 ###
 #
 gtiffDriver = gdal.GetDriverByName( 'GTiff' )
 if gtiffDriver is None:
  raise ValueError("Can't find GeoTiff Driver")
 gtiffDriver.CreateCopy( "beijing.tif",
  OpenArray( clip, prototype_ds=raster_path, xoff=xoffset, yoff=yoffset )
 )
 
 # Save as an 8-bit jpeg for an easy, quick preview
 clip = clip.astype(gdalnumeric.uint8)
 gdalnumeric.SaveArray(clip, "beijing.jpg", format="JPEG")
 
 gdal.ErrorReset()
 
 
if __name__ == '__main__':
 #shapefile_path, raster_path
 shapefile_path = 'beijing.shp'
 raster_path = 'world.tif'
 main( shapefile_path, raster_path )

import ogr,sys,os
import numpy as np

os.chdir(r'E:\')

#设置driver,并打开矢量文件
driver = ogr.GetDriverByName('ESRI Shapefile')
ds = driver.Open('sites.shp', 0)
if ds is None:
  print("Could not open", 'sites.shp')
  sys.exit(1)
#获取图册
layer = ds.GetLayer()

#要素数量
numFeatures = layer.GetFeatureCount()
print("Feature count: "+str(numFeatures))

#获取范围
extent = layer.GetExtent()
print("Extent:", extent)
print("UL:", extent[0],extent[3])
print("LR:", extent[1],extent[2])

#获取要素
feature = layer.GetNextFeature()
ids = []
xs = []
ys = []
covers = []
#循环每个要素属性
while feature:
  #获取字段“id”的属性
  id = feature.GetField('id')
  #获取空间属性
  geometry = feature.GetGeometryRef()
  x = geometry.GetX()
  y = geometry.GetY()
  cover = feature.GetField('cover')
  ids.append(id)
  xs.append(x)
  ys.append(y)
  covers.append(cover)
  feature = layer.GetNextFeature()

data = [ids, xs, ys, covers]
data = np.array(data)
data = data.transpose()

#写出致txt
np.savetxt('myfile.txt',data, fmt='%s %s %s %s')
np.savetxt('myfile.csv',data, fmt='%s %s %s %s')

#释放文件空间
layer.ResetReading()
feature.Destroy()
ds.Destroy()