221 lines
6.6 KiB
Python
221 lines
6.6 KiB
Python
from osgeo import ogr, gdal
|
||
import os
|
||
import argparse
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
def get_filename_without_ext(path):
|
||
base_name = os.path.basename(path)
|
||
if '.' not in base_name or base_name.startswith('.'):
|
||
return base_name
|
||
return base_name.rsplit('.', 1)[0]
|
||
|
||
def read_tif(path):
|
||
dataset = gdal.Open(path) # 打开TIF文件
|
||
if dataset is None:
|
||
print("无法打开文件")
|
||
return None, None, None
|
||
|
||
cols = dataset.RasterXSize # 图像宽度
|
||
rows = dataset.RasterYSize # 图像高度
|
||
bands = dataset.RasterCount
|
||
im_proj = dataset.GetProjection() # 获取投影信息
|
||
im_Geotrans = dataset.GetGeoTransform() # 获取仿射变换信息
|
||
im_data = dataset.ReadAsArray(0, 0, cols, rows) # 读取栅格数据为NumPy数组
|
||
print("行数：", rows)
|
||
print("列数：", cols)
|
||
print("波段：", bands)
|
||
del dataset # 关闭数据集
|
||
return im_proj, im_Geotrans, im_data
|
||
|
||
def write_envi(im_data, im_geotrans, im_proj, output_path):
|
||
"""
|
||
将数组数据写入ENVI格式文件
|
||
:param im_data: 输入的numpy数组（2D或3D）
|
||
:param im_geotrans: 仿射变换参数（6元组）
|
||
:param im_proj: 投影信息（WKT字符串）
|
||
:param output_path: 输出文件路径（无需扩展名，会自动生成.dat和.hdr）
|
||
"""
|
||
im_bands = 1
|
||
im_height, im_width = im_data.shape
|
||
# 创建ENVI格式驱动
|
||
driver = gdal.GetDriverByName("ENVI")
|
||
dataset = driver.Create(output_path, im_width, im_height, 1, gdal.GDT_Byte)
|
||
|
||
if dataset is not None:
|
||
dataset.SetGeoTransform(im_geotrans) # 设置地理变换参数
|
||
dataset.SetProjection(im_proj) # 设置投影
|
||
|
||
dataset.GetRasterBand(1).WriteArray(im_data)
|
||
|
||
dataset.FlushCache() # 确保数据写入磁盘
|
||
dataset = None # 关闭文件
|
||
|
||
|
||
def Strech_linear(im_data):
|
||
im_data_dB=10*np.log10(im_data)
|
||
immask=np.isfinite(im_data_dB)
|
||
infmask = np.isinf(im_data_dB)
|
||
imvail_data=im_data[immask]
|
||
im_data_dB=0
|
||
|
||
minvalue=np.nanmin(imvail_data)
|
||
maxvalue=np.nanmax(imvail_data)
|
||
|
||
infmask = np.isinf(im_data_dB)
|
||
im_data[infmask] = minvalue-100
|
||
im_data = (im_data - minvalue) / (maxvalue - minvalue) * 254+1
|
||
im_data=np.clip(im_data,0,255)
|
||
return im_data.astype(np.uint8)
|
||
|
||
def Strech_linear1(im_data):
|
||
im_data_dB = 10 * np.log10(im_data)
|
||
immask = np.isfinite(im_data_dB)
|
||
infmask = np.isinf(im_data_dB)
|
||
imvail_data = im_data[immask]
|
||
im_data_dB=0
|
||
|
||
minvalue=np.percentile(imvail_data,1)
|
||
maxvalue = np.percentile(imvail_data, 99)
|
||
|
||
|
||
im_data[infmask] = minvalue - 100
|
||
im_data = (im_data - minvalue) / (maxvalue - minvalue) * 254 + 1
|
||
im_data = np.clip(im_data, 0, 255)
|
||
|
||
return im_data.astype(np.uint8)
|
||
|
||
|
||
def Strech_linear2(im_data):
|
||
im_data_dB = 10 * np.log10(im_data)
|
||
immask = np.isfinite(im_data_dB)
|
||
infmask = np.isinf(im_data_dB)
|
||
imvail_data = im_data[immask]
|
||
im_data_dB = 0
|
||
|
||
minvalue = np.percentile(imvail_data, 2)
|
||
maxvalue = np.percentile(imvail_data, 98)
|
||
|
||
im_data[infmask] = minvalue - 100
|
||
im_data = (im_data - minvalue) / (maxvalue - minvalue) * 254 + 1
|
||
im_data = np.clip(im_data, 0, 255)
|
||
|
||
return im_data.astype(np.uint8)
|
||
|
||
def Strech_linear5(im_data):
|
||
im_data_dB = 10 * np.log10(im_data)
|
||
immask = np.isfinite(im_data_dB)
|
||
infmask = np.isinf(im_data_dB)
|
||
imvail_data = im_data[immask]
|
||
im_data_dB = 0
|
||
|
||
minvalue = np.percentile(imvail_data, 5)
|
||
maxvalue = np.percentile(imvail_data, 95)
|
||
|
||
im_data[infmask] = minvalue - 100
|
||
im_data = (im_data - minvalue) / (maxvalue - minvalue) * 254 + 1
|
||
im_data = np.clip(im_data, 0, 255)
|
||
|
||
return im_data.astype(np.uint8)
|
||
|
||
def Strech_SquareRoot(im_data):
|
||
im_data=np.sqrt(im_data)
|
||
immask = np.isfinite(im_data)
|
||
imvail_data = im_data[immask]
|
||
|
||
minvalue=np.nanmin(imvail_data)
|
||
maxvalue=np.nanmax(imvail_data)
|
||
im_data = (im_data - minvalue) / (maxvalue - minvalue) * 254 + 1
|
||
im_data = np.clip(im_data, 0, 255)
|
||
return im_data.astype(np.uint8)
|
||
|
||
def DataStrech(im_data,strechmethod):
|
||
# [,"Linear1","Linear2","Linear5","SquareRoot"]
|
||
if strechmethod == "Linear" :
|
||
return Strech_linear(im_data)
|
||
elif strechmethod == "Linear1":
|
||
return Strech_linear1(im_data)
|
||
elif strechmethod == "Linear2":
|
||
return Strech_linear2(im_data)
|
||
elif strechmethod == "Linear5":
|
||
return Strech_linear5(im_data)
|
||
elif strechmethod == "SquareRoot":
|
||
return Strech_SquareRoot(im_data)
|
||
else:
|
||
return im_data.astype(np.uint8)
|
||
|
||
|
||
|
||
|
||
|
||
def stretchProcess(infilepath,outfilepath,strechmethod):
|
||
im_proj, im_Geotrans, im_data=read_tif(infilepath)
|
||
envifilepath=get_filename_without_ext(outfilepath)+".bin"
|
||
envifilepath=os.path.join(os.path.dirname(outfilepath),envifilepath)
|
||
im_data = DataStrech(im_data,strechmethod)
|
||
im_data = im_data.astype(np.uint8)
|
||
write_envi(im_data,im_Geotrans,im_proj,envifilepath)
|
||
Image.fromarray(im_data).save(outfilepath,compress_level=0)
|
||
print("图像拉伸处理结束")
|
||
|
||
def getParams():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('-i','--infile',type=str,default=r"F:\天仪SAR卫星数据集\舰船数据\bc2-sp-org-vv-20250205t032055-021998-000036-0055ee-01.tiff", help='输入shapefile文件')
|
||
parser.add_argument('-o', '--outfile',type=str,default=r"F:\天仪SAR卫星数据集\舰船数据\bc2-sp-org-vv-20250205t032055-021998-000036-0055ee-01.png", help='输出geojson文件')
|
||
group = parser.add_mutually_exclusive_group()
|
||
group.add_argument(
|
||
'--Linear',
|
||
action='store_const',
|
||
const='Linear',
|
||
dest='method',
|
||
help='线性拉伸'
|
||
)
|
||
group.add_argument(
|
||
'--Linear1prec',
|
||
action='store_const',
|
||
const='Linear1',
|
||
dest='method',
|
||
help='1%线性拉伸'
|
||
)
|
||
group.add_argument(
|
||
'--Linear2prec',
|
||
action='store_const',
|
||
const='Linear2',
|
||
dest='method',
|
||
help='2%线性拉伸'
|
||
)
|
||
group.add_argument(
|
||
'--Linear5prec',
|
||
action='store_const',
|
||
const='Linear5',
|
||
dest='method',
|
||
help='5%线性拉伸'
|
||
)
|
||
group.add_argument(
|
||
'--SquareRoot',
|
||
action='store_const',
|
||
const='SquareRoot',
|
||
dest='method',
|
||
help='平方根拉伸'
|
||
)
|
||
parser.set_defaults(method='SquareRoot')
|
||
|
||
args = parser.parse_args()
|
||
return args
|
||
|
||
if __name__ == '__main__':
|
||
parser = getParams()
|
||
intiffPath=parser.infile
|
||
outbinPath=parser.outfile
|
||
methodstr=parser.method
|
||
print('infile=',intiffPath)
|
||
print('outfile=',outbinPath)
|
||
print('method=',methodstr)
|
||
stretchProcess(intiffPath, outbinPath, methodstr)
|
||
|
||
|
||
|
||
|
||
|
||
|