python gdal批量计算NDVI

最近尝试用python解决一些图像处理的问题，前面已经实现了图像的批量裁剪和镶嵌，现在来实现植被指数NDVI的批量计算。

import numpy as np
import os
from osgeo import gdal

def NDVI_batch(input_path,output_path):    #定义批量计算NDVI的函数
    for filename in os.listdir(input_path):
        # 打开需要裁剪的文件夹,将所有文件以列表的形式列出
        if filename.endswith('.tif'):
            # 判断文件是否为.tif结尾
            dataset = gdal.Open(os.path.join(input_path, filename), gdal.GA_ReadOnly)     #打开影像，只读模式
            width = dataset.RasterXSize   #获取数据宽度
            height = dataset.RasterYSize  #获取数据高度

            driver = gdal.GetDriverByName('GTiff')   # 选择输出格式
            output_dataset = driver.Create(os.path.join(output_path, "NDVI_" + filename), width, height, 1, gdal.GDT_Float32)   #创建输出文件

            band_red = dataset.GetRasterBand(3).ReadAsArray(0, 0, width, height).astype(np.float64)   #获取影像波段信息
            band_nir = dataset.GetRasterBand(4).ReadAsArray(0, 0, width, height).astype(np.float64)
            data_ndvi = (band_nir - band_red) / (band_nir + band_red)

            output_band = output_dataset.GetRasterBand(1)
            output_band.WriteArray(data_ndvi)   #写入数据
            output_band.FlushCache()
            output_dataset.SetGeoTransform(dataset.GetGeoTransform()) #设置地理变换
            output_dataset.SetProjection(dataset.GetProjection())     #设置投影信息

            dataset = None      #关闭文件
            output_dataset = None
            print(filename, "finished!")

以上是核心代码块，加入输入输出即可执行