sanya-dbms/processing/algs/gdal/ExportRasterByMask.py

"""
***************************************************************************
    ExportRasterByMask.py
    ---------------------
    Date                 : September 2013
    Copyright            : (C) 2013 by Alexander Bruy
    Email                : alexander bruy at gmail dot com
***************************************************************************
*                                                                         *
*   This program is free software; you can redistribute it and/or modify  *
*   it under the terms of the GNU General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
***************************************************************************
"""

__author__ = 'Alexander Bruy'
__date__ = 'September 2013'
__copyright__ = '(C) 2013, Alexander Bruy'

import os

from qgis.PyQt.QtGui import QIcon

from qgis.core import (QgsRasterFileWriter,
                       QgsProcessing,
                       QgsProcessingException,
                       QgsProcessingParameterDefinition,
                       QgsProcessingParameterFeatureSource,
                       QgsProcessingParameterRasterLayer,
                       QgsProcessingParameterCrs,
                       QgsProcessingParameterEnum,
                       QgsProcessingParameterExtent,
                       QgsProcessingParameterString,
                       QgsProcessingParameterNumber,
                       QgsProcessingParameterBoolean,
                       QgsProcessingParameterRasterDestination)
from processing.algs.gdal.GdalAlgorithm import GdalAlgorithm
from processing.algs.gdal.GdalUtils import GdalUtils
from processing.tools.StringUtils import (get_download_dir, dict_to_str_with_equals, sourceToDBConfig, getUUID, get_temp_dir)

pluginPath = os.path.split(os.path.split(os.path.dirname(__file__))[0])[0]


class ExportRasterByMask(GdalAlgorithm):
    INPUT = 'INPUT'
    MASK = 'MASK'
    SOURCE_CRS = 'SOURCE_CRS'
    TARGET_CRS = 'TARGET_CRS'
    EXTENT = 'TARGET_EXTENT'
    NODATA = 'NODATA'
    ALPHA_BAND = 'ALPHA_BAND'
    CROP_TO_CUTLINE = 'CROP_TO_CUTLINE'
    KEEP_RESOLUTION = 'KEEP_RESOLUTION'
    SET_RESOLUTION = 'SET_RESOLUTION'
    X_RESOLUTION = 'X_RESOLUTION'
    Y_RESOLUTION = 'Y_RESOLUTION'
    OPTIONS = 'OPTIONS'
    DATA_TYPE = 'DATA_TYPE'
    MULTITHREADING = 'MULTITHREADING'
    EXTRA = 'EXTRA'
    OUTPUT = 'OUTPUT'

    def __init__(self):
        super().__init__()

    def initAlgorithm(self, config=None):

        self.TYPES = [self.tr('Use Input Layer Data Type'), 'Byte', 'Int16', 'UInt16', 'UInt32', 'Int32', 'Float32',
                      'Float64', 'CInt16', 'CInt32', 'CFloat32', 'CFloat64', 'Int8']

        self.addParameter(QgsProcessingParameterRasterLayer(self.INPUT,
                                                            self.tr('数据源')))
        self.addParameter(QgsProcessingParameterFeatureSource(self.MASK,
                                                              self.tr('导出的数据范围'),
                                                              [QgsProcessing.TypeVectorPolygon]))
        self.addParameter(QgsProcessingParameterNumber(self.NODATA,
                                                       self.tr('无效值'),
                                                       type=QgsProcessingParameterNumber.Double,
                                                       defaultValue=None,
                                                       optional=True))
        # self.addParameter(QgsProcessingParameterBoolean(self.CROP_TO_CUTLINE,
        #                                                 self.tr('按范围进行导出'),
        #                                                 defaultValue=True))

        options_param = QgsProcessingParameterString(self.OPTIONS,
                                                     self.tr('高级设置'),
                                                     defaultValue='',
                                                     optional=True)
        options_param.setFlags(options_param.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
        options_param.setMetadata({
            'widget_wrapper': {
                'class': 'processing.algs.gdal.ui.RasterOptionsWidget.RasterOptionsWidgetWrapper'}})
        self.addParameter(options_param)

        dataType_param = QgsProcessingParameterEnum(self.DATA_TYPE,
                                                    self.tr('输出栅格类型'),
                                                    self.TYPES,
                                                    allowMultiple=False,
                                                    defaultValue=0)
        dataType_param.setFlags(dataType_param.flags() | QgsProcessingParameterDefinition.FlagAdvanced)
        self.addParameter(dataType_param)

        self.addParameter(QgsProcessingParameterRasterDestination(self.OUTPUT,
                                                                  self.tr('导出位置')))

    def name(self):
        return 'rasterexport'

    def displayName(self):
        return self.tr('栅格数据分发')

    def icon(self):
        return QIcon(os.path.join(pluginPath, 'images', 'dbms', 'export-img.png'))

    def group(self):
        return self.tr('数据分发')

    def groupId(self):
        return 'dataexport'

    def commandName(self):
        return 'gdalwarp'

    def getConsoleCommands(self, parameters, context, feedback, executing=True):
        inLayer = self.parameterAsRasterLayer(parameters, self.INPUT, context)
        if inLayer is None:
            raise QgsProcessingException(self.invalidRasterError(parameters, self.INPUT))

        maskLayer, maskLayerName = self.getOgrCompatibleSource(self.MASK, parameters, context, feedback, executing)

        # sourceCrs = self.parameterAsCrs(parameters, self.SOURCE_CRS, context)
        # targetCrs = self.parameterAsCrs(parameters, self.TARGET_CRS, context)

        # bbox = self.parameterAsExtent(parameters, self.EXTENT, context)
        # bboxCrs = self.parameterAsExtentCrs(parameters, self.EXTENT, context)

        if self.NODATA in parameters and parameters[self.NODATA] is not None:
            nodata = self.parameterAsDouble(parameters, self.NODATA, context)
        else:
            nodata = None
        options = self.parameterAsString(parameters, self.OPTIONS, context)
        out = self.parameterAsOutputLayer(parameters, self.OUTPUT, context)
        self.setOutputValue(self.OUTPUT, out)

        arguments = ['-overwrite']

        data_type = self.parameterAsEnum(parameters, self.DATA_TYPE, context)
        if data_type:
            if self.TYPES[data_type] == 'Int8' and GdalUtils.version() < 3070000:
                raise QgsProcessingException(self.tr('Int8 data type requires GDAL version 3.7 or later'))

            arguments.append('-ot ' + self.TYPES[data_type])

        output_format = QgsRasterFileWriter.driverForExtension(os.path.splitext(out)[1])
        if not output_format:
            raise QgsProcessingException(self.tr('Output format is invalid'))

        arguments.append('-of')
        arguments.append(output_format)

        arguments.append('-cutline')
        arguments.append(maskLayer)
        arguments.append('-cl')
        arguments.append(maskLayerName)

        # if self.parameterAsBoolean(parameters, self.CROP_TO_CUTLINE, context):
        #     arguments.append('-crop_to_cutline')

        if nodata is not None:
            arguments.append(f'-dstnodata {nodata}')

        if options:
            arguments.extend(GdalUtils.parseCreationOptions(options))
        source = inLayer.source()
        if source.__contains__("dbname"):
            uid = getUUID() + ".tif"
            tempdir = get_temp_dir()
            tempfile = f'{tempdir}/{uid}'
            arguments.append(tempfile)
            arguments.append(out)
            ####### 开始处理gdal导出postgis栅格数据 ########
            real = source.replace("\"\"\"", "").replace(" (rast)", "")
            re = sourceToDBConfig(real)
            st = dict_to_str_with_equals(re)
            source = f"PG:{st}"
            print(source)
            cmd = f'gdal_translate -of GTiff "{source}" {tempfile}'
            return [cmd, self.commandName(), GdalUtils.escapeAndJoin(arguments)]
        else:
            arguments.append(source)
            arguments.append(out)
            return [self.commandName(), GdalUtils.escapeAndJoin(arguments)]