68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
|
from osgeo import gdal
|
|||
|
|
# from tqdm import tqdm
|
|||
|
|
|
|||
|
|
def ReadTif(tif_path):
|
|||
|
|
dataset = gdal.Open(tif_path)
|
|||
|
|
width = dataset.RasterXSize
|
|||
|
|
height = dataset.RasterYSize
|
|||
|
|
|
|||
|
|
geotrans = list(dataset.GetGeoTransform()) # 仿射矩阵
|
|||
|
|
proj = dataset.GetProjection() # 地图投影信息
|
|||
|
|
data = dataset.ReadAsArray(0, 0, width, height) # 将数据写成数组，对应栅格矩阵
|
|||
|
|
|
|||
|
|
del dataset # 关闭对象，文件dataset
|
|||
|
|
return proj, geotrans, data, width, height
|
|||
|
|
|
|||
|
|
def writeTif(fileroute, im_proj, im_geotrans, im_data):
|
|||
|
|
# 判断栅格数据的数据类型
|
|||
|
|
"""
|
|||
|
|
GDAL中的GDALDataType是一个枚举型，其中的值为：
|
|||
|
|
GDT_Unknown : 未知数据类型
|
|||
|
|
GDT_Byte : 8bit正整型 (C++中对应unsigned char)
|
|||
|
|
GDT_UInt16 : 16bit正整型 (C++中对应 unsigned short)
|
|||
|
|
GDT_Int16 : 16bit整型 (C++中对应 short 或 short int)
|
|||
|
|
GDT_UInt32 : 32bit 正整型 (C++中对应unsigned long)
|
|||
|
|
GDT_Int32 : 32bit整型 (C++中对应int 或 long 或 long int)
|
|||
|
|
GDT_Float32 : 32bit 浮点型 (C++中对应float)
|
|||
|
|
GDT_Float64 : 64bit 浮点型 (C++中对应double)
|
|||
|
|
GDT_CInt16 : 16bit复整型 (?)
|
|||
|
|
GDT_CInt32 : 32bit复整型 (?)
|
|||
|
|
GDT_CFloat32 : 32bit复浮点型 (?)
|
|||
|
|
GDT_CFloat64 : 64bit复浮点型 (?)
|
|||
|
|
"""
|
|||
|
|
if 'int8' in im_data.dtype.name:
|
|||
|
|
datatype = gdal.GDT_Byte
|
|||
|
|
elif 'int16' in im_data.dtype.name:
|
|||
|
|
datatype = gdal.GDT_UInt16
|
|||
|
|
else:
|
|||
|
|
datatype = gdal.GDT_Float32
|
|||
|
|
|
|||
|
|
# 判读数组维数
|
|||
|
|
if len(im_data.shape) == 3:
|
|||
|
|
im_bands, im_height, im_width = im_data.shape
|
|||
|
|
else:
|
|||
|
|
im_bands, (im_height, im_width) = 1, im_data.shape
|
|||
|
|
im_data = im_data.reshape(im_bands,im_height, im_width)
|
|||
|
|
# 创建文件
|
|||
|
|
driver = gdal.GetDriverByName("GTiff") # 数据类型必须有，因为要计算需要多大内存空间
|
|||
|
|
dataset = driver.Create(fileroute, im_width,im_height, im_bands, datatype)
|
|||
|
|
|
|||
|
|
dataset.SetGeoTransform(im_geotrans) # 写入仿射变换参数
|
|||
|
|
dataset.SetProjection(im_proj) # 写入投影
|
|||
|
|
|
|||
|
|
if im_bands == 1:
|
|||
|
|
dataset.GetRasterBand(1).WriteArray(im_data[0]) # 写入数组数据
|
|||
|
|
else:
|
|||
|
|
#for i in range(im_bands):
|
|||
|
|
for i in tqdm(range(im_bands)):
|
|||
|
|
dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
|
|||
|
|
|
|||
|
|
del dataset
|
|||
|
|
|
|||
|
|
def np2gdal(hwc): # hwc2chw
|
|||
|
|
chw = hwc.swapaxes(2, 0).swapaxes(1, 2) # h,w,c to c,h,w
|
|||
|
|
return chw
|
|||
|
|
|
|||
|
|
def gdal2np(chw): #chw2hwc
|
|||
|
|
hwc = chw.swapaxes(1, 0).swapaxes(1, 2) # h,w,c
|
|||
|
|
return hwc
|