from osgeo import gdal, ogr, osr
import os
import datetime
import numpy as np

path = "Z_NAFP20210727.tif"


if __name__ == '__main__':
    start_time = datetime.datetime.now()

    inraster = gdal.Open(path)  # 讀取路徑中的柵格數(shù)據(jù)
    inband = inraster.GetRasterBand(1)  # 這個(gè)波段就是最后想要轉(zhuǎn)為矢量的波段，如果是單波段數(shù)據(jù)的話那就都是1
    prj = osr.SpatialReference()
    prj.ImportFromWkt(inraster.GetProjection())  # 讀取柵格數(shù)據(jù)的投影信息，用來為后面生成的矢量做準(zhǔn)備

    outshp = path[:-4] + ".shp"  # 給后面生成的矢量準(zhǔn)備一個(gè)輸出文件名，這里就是把原柵格的文件名后綴名改成shp了
    drv = ogr.GetDriverByName("ESRI Shapefile")
    if os.path.exists(outshp):  # 若文件已經(jīng)存在，則刪除它繼續(xù)重新做一遍
        drv.DeleteDataSource(outshp)
    Polygon = drv.CreateDataSource(outshp)  # 創(chuàng)建一個(gè)目標(biāo)文件
    Poly_layer = Polygon.CreateLayer(path[:-4], srs=prj, geom_type=ogr.wkbMultiPolygon)  # 對(duì)shp文件創(chuàng)建一個(gè)圖層，定義為多個(gè)面類
    newField = ogr.FieldDefn('value', ogr.OFTReal)  # 給目標(biāo)shp文件添加一個(gè)字段，用來存儲(chǔ)原始柵格的pixel value,浮點(diǎn)型，
    Poly_layer.CreateField(newField)

    gdal.Polygonize(inband, None, Poly_layer, 0)  # 核心函數(shù)，執(zhí)行的就是柵格轉(zhuǎn)矢量操作
    # gdal.FPolygonize(inband, None, Poly_layer, 0)  # 只轉(zhuǎn)矩形，不合并
    Polygon.SyncToDisk()
    Polygon = None
    end_time = datetime.datetime.now()
    print("Succeeded at", end_time)
    print("Elapsed Time:", end_time - start_time)  # 輸出程序運(yùn)行所需時(shí)間