GDAL和OPENCV 对tif影像进行平滑模糊

GDAL和OPENCV 对tif影像进行平滑模糊

import os
from osgeo import  ogr
from osgeo import gdal
from osgeo import osr
import numpy as np
import cv2 
import matplotlib.pyplot as plt
import matplotlib.colors as mpc
import argparse

def geo2pixel(geot:list,geo:list):
    px=(geot[5]*geo[0]-geot[2]*geo[1]-geot[0]*geot[5]+geot[2]*geot[3])/(geot[1]*geot[5]-geot[2]*geot[4])
    py=(geot[4]*geo[0]-geot[1]*geo[1]+geot[1]*geot[3]-geot[4]*geot[0])/(geot[2]*geot[4]-geot[1]*geot[5])
    return [int(px),int(py)]

def pixel2geo(geot:list,pixel:list):
    x_geo=geot[0]+pixel[0]*geot[1]+pixel[1]*geot[2]
    y_geo=geot[3]+pixel[0]*geot[4]+pixel[1]*geot[5]
    return [x_geo,y_geo]

def edgedetect(bblock):
    avg=np.nanmean(bblock)
    std=np.nanstd(bblock)
    minv=avg-2.5*std
    maxv=avg+2.5*std
    bblock=(bblock-minv)/(maxv-minv)*255
    bblock=np.where(bblock<0,0,bblock)
    bblock=np.where(bblock>255,255,bblock)
    bblock=np.uint8(bblock)
    canny=cv2.Canny(bblock,20, 100)
    lines=cv2.HoughLinesP(canny,1,np.pi/180,200,minLineLength=100,maxLineGap=400)
    try:
        lines=lines[:,0,:]
        for x1,y1,x2,y2 in lines:
            dx=x2-x1
            dy=y2-y1
            angle=np.arctan(dy/dx)/np.pi*180
            if angle>=10 and angle<=20:
                return True
    except:
        return False
    return False

def stdStretch(array, colorList,outPath):
    # array 影像
    # colorList 色带参数
    avg = np.nanmean(array)
    std = np.nanstd(array)
    # 默认按2.5倍标准差进行拉伸
    v_max = avg+2.5*std
    v_min = avg-2.5*std
    norm = plt.Normalize(v_min, v_max)
    colors = colorList
    cm = mpc.LinearSegmentedColormap.from_list('colormap', colors, 512)
    plt.imshow(array, cm, norm)
    plt.axis('off')
    plt.subplots_adjust(bottom=0, top=1, left=0, right=1, hspace=0, wspace=0)
    plt.colorbar()
    plt.savefig(outPath, transparent=True, dpi=100)
    plt.close()


if __name__=='__main__':
    # 获取输入参数
    parser=argparse.ArgumentParser()
    parser.add_argument('tifinput',type=str,help='The path of the tif file to process!')
    parser.add_argument('burst',type=str,help='SHP file to define the processed area!')
    parser.add_argument('mask',type=str,help='SHP file to clip the result tif!')
    parser.add_argument('tifoutput',type=str,help='The path to store the result!')
    parser.add_argument('pngout',type=str,help='The path to store the rendered result!')
    args=parser.parse_args()
    tifinput=args.tifinput
    burstSHP=args.burst
    maskSHP=args.mask
    tifoutput=args.tifoutput
    pngout=args.pngout

    if not(os.path.exists(tifinput)):
        print("{} doesn't exist!".format(tifinput))
        exit()
    if not(os.path.exists(burstSHP)):
        print("{} doesn't exist!".format(burstSHP))
        exit()
    if not(os.path.exists(maskSHP)):
        print("{} doesn't exist!".format(maskSHP))
        exit()

    # 最大平滑窗口 最小平滑窗口 缓冲区半径
    rmax=25
    rmin=1
    rbuffer=3000

    # 读取栅格数据
    gdal.AllRegister()
    tifin=gdal.Open(tifinput)
    bandin=gdal.Dataset.GetRasterBand(tifin,1)
    arrin=gdal.Band.ReadAsArray(bandin)
    # 将nodata设置为np.nan
    arrin=np.where(arrin<-9999,np.nan,arrin)
    arrout=arrin[:,:]
    gt=gdal.Dataset.GetGeoTransform(tifin)
    # sr=gdal.Dataset.GetSpatialRef(tifin)
    sr=gdal.Dataset.GetProjection(tifin)
    # sr=gdal.Dataset.GetProjectionRef(tifin)
    xSize=bandin.XSize
    ySize=bandin.YSize


    ogr.RegisterAll()
    ds=ogr.Open(burstSHP)
    lay=ogr.DataSource.GetLayerByIndex(ds,0)
    # ogr.Layer.SetAttributeFilter(lay,"code in ('e3','e4','e6','e7','m5','m6','m8','w7')")

    # 加载 研究区域shp

    while True:
        fea=ogr.Layer.GetNextFeature(lay)
        if not(fea):
            break
        fgeo=ogr.Feature.geometry(fea)
        fbuffer=ogr.Geometry.Buffer(fgeo,rbuffer,0)
        fbound=ogr.Geometry.GetEnvelope(fgeo)
        x_geo_min=fbound[0]
        x_geo_max=fbound[1]
        y_geo_min=fbound[2]
        y_geo_max=fbound[3]
        x_pix_min,y_pix_min=geo2pixel(gt,[x_geo_min,y_geo_max])
        x_pix_max,y_pix_max=geo2pixel(gt,[x_geo_max,y_geo_min])
        bblock=arrin[y_pix_min:(y_pix_max+1),x_pix_min:(x_pix_max+1)]
        flag=edgedetect(bblock)
        if not(flag):
            continue
        # print(ogr.Feature.GetField(fea,'code'))
        for row in range(y_pix_min,y_pix_max+1):
            for col in range(x_pix_min,x_pix_max+1):
                if row<0 or row>=ySize or col<0 or col>=xSize:
                    # 像素不在影像覆盖范围内不处理
                    continue
                pvalue=arrin[row,col]
                if np.isnan(pvalue):
                    # nodata像素不处理
                    continue
                x_geo,y_geo=pixel2geo(gt,[col,row])
                pgeo=ogr.CreateGeometryFromWkt('POINT({} {})'.format(x_geo,y_geo),osr.SpatialReference(sr))
                if not(ogr.Geometry.Within(pgeo,fbuffer)):
                    # 该像素不在缓冲区内部不处理
                    continue
                
                # 计算平滑窗口大小
                dist=ogr.Geometry.Distance(pgeo,fgeo)
                r=int(dist/rbuffer*(rmin-rmax)+rmax)
                pblock=arrin[(row-r):(row+r+1),(col-r):(col+r+1)]
                arrout[row,col]=np.nanmean(pblock)

    # 渲染为PNG输出为Preview
    colorList=['#0000FF', '#3B7FFF', '#00FFFF','#B4FF91', '#FFFF00', '#FF9100', '#FF0000']
    tifname=os.path.basename(tifinput)
    pngname=tifname.replace('.tif','.png')
    pngfile=os.path.join(pngout,pngname)
    stdStretch(arrout,colorList,pngfile)
    
    # 输出结果
    tiffile=os.path.join(tifoutput,tifname)
    driver=gdal.GetDriverByName('GTiff')
    tifout=gdal.Driver.Create(driver,tiffile,xSize,ySize,1,gdal.GDT_Float32,['COMPRESS=LZW'])
    bandout=gdal.Dataset.GetRasterBand(tifout,1)
    gdal.Band.WriteArray(bandout,arrout)
    gdal.Dataset.SetGeoTransform(tifout,gt)
    # gdal.Dataset.SetSpatialRef(tifout,sr)
    gdal.Dataset.SetProjection(tifout,sr)
    print('{} has been processed!'.format(tifname))