Python 遥感影像分块读取并计算NDVI（或其他操作）

# -*- coding: utf-8 -*-
# @Time    : 2022/8/25 13:30
# @Author  : Zhang Min
# @FileName: tif按步长分块.py
import glob
import math
import time

import numpy as np
from osgeo import gdal


def readTif(tifFile):
    dataset = gdal.Open(tifFile)
    width = dataset.RasterXSize
    height = dataset.RasterYSize
    geotrans = dataset.GetGeoTransform()
    proj = dataset.GetProjection()
    data = dataset.ReadAsArray(0, 0, width, height)

    del dataset
    return data, geotrans, proj, width, height


def writeTif(filename, geotrans, proj, data):
    # 判断栅格数据的数据类型
    if 'int8' in data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    # 判读数组维数
    if len(data.shape) == 3:
        im_bands, im_height, im_width = data.shape
    else:
        im_bands, (im_height, im_width) = 1, data.shape

    # 创建文件
    driver = gdal.GetDriverByName("GTiff")  # 数据类型必须有，因为要计算需要多大内存空间
    dataset = driver.Create(filename, im_width, im_height, im_bands, datatype)
    dataset.SetGeoTransform(geotrans)  # 写入仿射变换参数
    dataset.SetProjection(proj)  # 写入投影
    if im_bands == 1:
        dataset.GetRasterBand(1).WriteArray(data)  # 写入数组数据
    else:
        for i in range(im_bands):
            dataset.GetRasterBand(i + 1).WriteArray(data[i])
    del dataset


def cal_nvdi(data):
    ndvi_molecule = data[3, :, :] - data[2, :, :]
    ndvi_denominator = data[3, :, :] + data[2, :, :] + 0.00001

    ndvi = ndvi_molecule / ndvi_denominator

    return ndvi


if __name__ == '__main__':
    # 时序影像的文件夹位置
    tif_list = glob.glob(r'D:\工作日志\NDVI分块\耕地\*.tif')

    # 保存文件夹
    dataNewArea = r"D:\工作日志\NDVI分块\耕地分块结果"
    # 步长
    step = 1024

    for i in range(len(tif_list)):
        start_time = time.time()

        dataset = gdal.Open(tif_list[i])
        width = dataset.RasterXSize  # 列
        height = dataset.RasterYSize  # 行
        geotrans = dataset.GetGeoTransform()
        proj = dataset.GetProjection()
        # data = dataset.ReadAsArray(0, 0, width, height)
        # print(width, height)  # 列 6787 行：7431
        # print(data.shape)  # C H（行:7431） W（列:6787）

        # ReadAsArray的起始点
        startX = 0
        startY = 0

        # 列 行 按照step 分成多少块，向上取整
        wcNum = math.ceil(width / step)  # 列块
        hrNum = math.ceil(height / step)  # 行块

        for i in range(0, wcNum):  # 逐列块
            for j in range(0, hrNum):  # 逐行块

                # 处于最后一列最后一行，行列都向内缩进forwardX，forwardY个像素
                if (i == wcNum - 1) and (j == hrNum - 1):
                    differentialX = step - (width - startX)  # 381 缺失381列刚好凑满1024
                    forwardX = startX - differentialX  # 向前挪动到6144-381

                    differentialY = step - (height - startY)  # 761
                    forwardY = startY - differentialY  # 向前挪动到6407
                    new_data = dataset.ReadAsArray(forwardX, forwardY, step, step).astype(np.float32)

                    # NDVI计算
                    ndvi = cal_nvdi(new_data)

                    # 投影坐标计算
                    px = geotrans[0] + forwardX * geotrans[1] + forwardY * geotrans[2]
                    py = geotrans[3] + forwardX * geotrans[4] + forwardY * geotrans[5]

                # 处于最后一列,列向内缩进forwardX个像素
                elif (i == wcNum - 1):
                    differentialX = step - (width - startX)  # 381 缺失381列刚好凑满1024
                    forwardX = startX - differentialX  # 向前挪动到6144-381
                    new_data = dataset.ReadAsArray(forwardX, startY, step, step).astype(np.float32)

                    ndvi = cal_nvdi(new_data)

                    px = geotrans[0] + forwardX * geotrans[1] + startY * geotrans[2]
                    py = geotrans[3] + forwardX * geotrans[4] + startY * geotrans[5]

                # 处于最后一行，行向内缩进forwardY个像素
                elif (j == hrNum - 1):
                    differentialY = step - (height - startY)  # 761 缺失761列刚好凑满1024
                    forwardY = startY - differentialY  # 向前挪动到7431-761
                    new_data = dataset.ReadAsArray(startX, forwardY, step, step).astype(np.float32)

                    ndvi = cal_nvdi(new_data)

                    px = geotrans[0] + startX * geotrans[1] + forwardY * geotrans[2]
                    py = geotrans[3] + startX * geotrans[4] + forwardY * geotrans[5]

                # 非以上情况
                else:
                    new_data = dataset.ReadAsArray(startX, startY, step, step).astype(np.float32)
                    ndvi = cal_nvdi(new_data)
                    px = geotrans[0] + startX * geotrans[1] + startY * geotrans[2]
                    py = geotrans[3] + startX * geotrans[4] + startY * geotrans[5]

                startY = startY + step

                newGeotrans = (px, geotrans[1], geotrans[2], py, geotrans[4], geotrans[5])
                
                # 写入
                writeTif(dataNewArea + '//ndvi_' + str(i) + "_" + str(j) + ".tif", geotrans=newGeotrans, proj=proj,
                         data=ndvi)

            startX = startX + step
            startY = 0

        end_time = time.time()
        print(end_time - start_time)