geolib.py 文件源码

def shp2array(shp_fn, r_ds=None, res=None, extent=None, t_srs=None):
    """Rasterize input shapefile to match existing raster Dataset (or specified res/extent/t_srs)
    """
    shp_ds = ogr.Open(shp_fn)
    lyr = shp_ds.GetLayer()
    #This returns xmin, ymin, xmax, ymax
    shp_extent = lyr_extent(lyr)
    shp_srs = lyr.GetSpatialRef()
    # dst_dt = gdal.GDT_Byte
    ndv = 0
    if r_ds is not None:
        r_extent = ds_extent(r_ds)
        res = get_res(r_ds, square=True)[0] 
        if extent is None:
            extent = r_extent
        r_srs = get_ds_srs(r_ds)
        r_geom = ds_geom(r_ds)
        # dst_ns = r_ds.RasterXSize
        # dst_nl = r_ds.RasterYSize
        #Convert raster extent to shp_srs
        cT = osr.CoordinateTransformation(r_srs, shp_srs)
        r_geom_reproj = geom_dup(r_geom)
        r_geom_reproj.Transform(cT)
        r_geom_reproj.AssignSpatialReference(t_srs)
        lyr.SetSpatialFilter(r_geom_reproj)
        #lyr.SetSpatialFilter(ogr.CreateGeometryFromWkt(wkt))
    else:
        #TODO: clean this up
        if res is None:
            sys.exit("Must specify input res")
        if extent is None:
            print("Using input shp extent")
            extent = shp_extent
    if t_srs is None:
        t_srs = r_srs
    if not shp_srs.IsSame(t_srs):
        print("Input shp srs: %s" % shp_srs.ExportToProj4())
        print("Specified output srs: %s" % t_srs.ExportToProj4())
        out_ds = lyr_proj(lyr, t_srs)
        outlyr = out_ds.GetLayer()
    else:
        outlyr = lyr
    #outlyr.SetSpatialFilter(r_geom)
    m_ds = mem_ds(res, extent, srs=t_srs, dtype=gdal.GDT_Byte)
    b = m_ds.GetRasterBand(1)
    b.SetNoDataValue(ndv)
    gdal.RasterizeLayer(m_ds, [1], outlyr, burn_values=[1])
    a = b.ReadAsArray()
    a = ~(a.astype('Bool'))
    return a