sanya-dbms/processing/tools/Raster/ClipRasterByGrid.py

# -*- coding: utf-8 -*-

__author__ = 'wanger'
__description__ = '按照掩膜范围裁剪栅格数据'
__date__ = '2024-11-25'
__copyright__ = '(C) 2024 by siwei'
__revision__ = '1.0'

import os
import subprocess

import math

from qgis._core import QgsRasterFileWriter, QgsProcessingParameterFolderDestination
from qgis.core import (
    QgsVectorLayer,
    QgsRasterLayer,
    QgsFeature,
    QgsGeometry
)
import processing
from PyQt5.QtCore import QFileInfo
from qgis.analysis import QgsRasterCalculator, QgsRasterCalculatorEntry

from qgis.PyQt.QtCore import QCoreApplication
from qgis._core import QgsSpatialIndex, QgsField, QgsRectangle, QgsProcessingParameterVectorDestination, \
    QgsVectorFileWriter, QgsWkbTypes, QgsProcessingParameterNumber, QgsProcessingParameterBoolean, QgsGeometry, \
    QgsFeature
from qgis.core import (QgsProcessing,
                       QgsVectorLayer,
                       QgsFeatureSink,
                       QgsProcessingException,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFile,
                       QgsProcessingParameterFeatureSource,
                       QgsProcessingParameterFeatureSink)
from qgis import processing


class ClipRasterByMaskProcessingAlgorithm(QgsProcessingAlgorithm):
    INPUT_SHP = 'INPUT_SHP'
    INPUT_RASTER = 'INPUT_RASTER'
    OUTPUT = 'OUTPUT'

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def createInstance(self):
        return ClipRasterByMaskProcessingAlgorithm()

    def name(self):
        return 'cliprasterbymask'

    def displayName(self):
        return self.tr('按照掩膜范围裁剪栅格数据')

    def group(self):
        return self.tr('栅格裁剪')

    def groupId(self):
        return 'rasterclip'

    def shortHelpString(self):
        return self.tr("遍历输入shp数据要素裁剪输入栅格，保存到输出文件夹。")

    def initAlgorithm(self, config=None):
        self.addParameter(QgsProcessingParameterFile(
            self.INPUT_SHP,
            '网格数据',
            extension='shp'
        ))
        self.addParameter(QgsProcessingParameterFile(
            self.INPUT_RASTER,
            '栅格数据',
            extension='tif'
        ))
        # 通过 setExtensions 方法设置多个扩展名
        # self.parameter(self.INPUT_RASTER).setExtensions(['tif', 'img'])  # 支持 .shp, .csv, .txt
        self.addParameter(QgsProcessingParameterFolderDestination(self.OUTPUT, '输出文件夹'))

    # 执行
    def processAlgorithm(self, parameters, context, feedback):
        input_shp = self.parameterAsString(parameters, self.INPUT_SHP, context)
        input_raster = self.parameterAsOutputLayer(parameters, self.INPUT_RASTER, context)
        output_raster = self.parameterAsString(parameters, self.OUTPUT, context)
        # 输入Shapefile路径
        # input_shp = "E:\\projects\\遥感技术部需求\\裁剪栅格\\data\\grid_extent.shp"
        # 输入栅格数据路径
        # input_raster = "D:\\gisdata\\tif\\unnamed_卫图_Level_12.tif"
        # 输出裁剪后的栅格路径
        # output_raster = "E:\\projects\\遥感技术部需求\\裁剪栅格\\output\\"
        # 加载输入Shapefile和栅格数据
        vector_layer = QgsVectorLayer(input_shp, "vector_layer", "ogr")
        raster_layer = QgsRasterLayer(input_raster, "raster_layer")
        fields = vector_layer.fields()
        # 获取分辨率（像素大小）
        pixel_width = raster_layer.rasterUnitsPerPixelX()
        pixel_height = raster_layer.rasterUnitsPerPixelY()
        # 遍历输入图层的所有要素，并为每个要素保存一个单独的Shapefile文件
        total = 0
        for idx, feature in enumerate(vector_layer.getFeatures()):
            # 获取当前要素的几何数据
            geometry = feature.geometry()
            key = feature["新图号"]
            extent = geometry.boundingBox()

            # 创建一个新的 Shapefile 文件名，可以使用索引或 ID 来命名
            # output_shp = os.path.join(output_raster, f"{key}.shp")
            #
            # # 创建输出shapefile
            # writer = QgsVectorFileWriter(output_shp, 'UTF-8', fields, QgsWkbTypes.MultiPolygon, vector_layer.crs(),
            #                              'ESRI Shapefile')
            # writer.addFeature(feature)
            #
            # # 完成写入并关闭输出文件
            # del writer

            # 使用QGIS的processing模块进行栅格裁剪
            # params = {
            #     'INPUT': raster_layer,  # 输入栅格
            #     'MASK': output_shp,  # 用于裁剪的矢量图层
            #     'OUTPUT': f"{output_raster}\\{key}.tif",  # 输出裁剪后的栅格文件
            #     'CROP_TO_CUTLINE': True
            # }
            # 调用 clipRasterByExtent 进行裁剪
            # vector_layer_cur = QgsVectorLayer(output_shp, 'Shapefile Layer')
            # params = {
            #     'INPUT': input_raster,
            #     'EXTENT': vector_layer_cur.extent(),
            #     'PROJWIN': vector_layer_cur.extent(),
            #     'OUTPUT': f"{output_raster}\\{key}.tif"
            # }
            # extent = vector_layer_cur.extent()
            # 计算裁剪范围的整数倍，确保裁剪范围对齐
            minX = extent.xMinimum()
            maxX = extent.xMaximum()
            minY = extent.yMinimum()
            maxY = extent.yMaximum()

            # 通过对齐最大值，确保裁剪范围完整
            maxX_aligned = maxX + pixel_width
            maxY_aligned = maxY + pixel_width
            minX_aligned = minX - pixel_width
            minY_aligned = minY - pixel_width
            # maxX_aligned = math.ceil(maxX / pixel_width) * pixel_width + pixel_width
            # maxY_aligned = math.ceil(maxY / pixel_height) * pixel_height + pixel_width
            # minX_aligned = math.floor(minX / pixel_width) * pixel_width - pixel_width
            # minY_aligned = math.floor(minY / pixel_height) * pixel_height - pixel_width
            # extent_cur = f'{minX},{extent.yMaximum()},{extent.xMaximum()},{minY}'  # 裁剪范围：minX,maxY,maxX,minY
            # extent_cur = f'524101.737 3310956.7264999966 530229.7547000019 3306243.2199'  # 裁剪范围：minX,maxY,maxX,minY
            # print(extent_cur)

            # 设置 gdal_translate 命令和参数
            gdal_translate_cmd = [
                'gdal_translate',  # GDAL 工具名称
                '-projwin', str(minX_aligned), str(maxY_aligned), str(maxX_aligned), str(minY_aligned),
                # 裁剪范围（minX, maxY, maxX, minY）
                '-of', 'GTiff',  # 输出格式为 GeoTIFF
                input_raster,  # 输入栅格文件路径
                f"{output_raster}\\{key}.tif"  # 输出文件路径
            ]
            print(gdal_translate_cmd)

            # 运行命令
            try:
                subprocess.run(gdal_translate_cmd, check=True)
                print("gdal_translate 执行成功！")
                total = total + 1
            except subprocess.CalledProcessError as e:
                print(f"gdal_translate 执行失败: {e}")

            # 3. 调用 `gdal:translate` 进行裁剪
            # parameters = {
            #     'INPUT': input_raster,  # 输入栅格
            #     'OUTPUT': f"{output_raster}\\{key}.tif",  # 输出路径
            #     'PROJWIN': extent_cur,  # 裁剪范围 (左下右上)
            #     # 'TARGET_CRS': 'EPSG:4544'  # 可选：指定目标坐标系统（如果需要）
            # }
            # print(parameters)

            # 4. 运行工具
            # processing.run('gdal:translate', parameters)

            # 执行裁剪操作
            # processing.run('gdal:cliprasterbyextent', params)
            # processing.run('gdal:cliprasterbymasklayer', params)
            # print(f"Cropped raster saved to {output_raster}")

        return {
            "状态": f"处理成功，共裁剪出{total}块栅格数据。"
        }