Blog SIG & Territoires

from qgis.core import (
    QgsProcessingAlgorithm,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterNumber,
    QgsProcessingParameterEnum,
    QgsProcessingParameterRasterDestination,
    QgsProcessingException
)
import os
import subprocess
from osgeo import gdal, osr

class ESRGANSuperResolution(QgsProcessingAlgorithm):

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer("INPUT_RASTER", "Raster d’entrée")
        )
        self.addParameter(
            QgsProcessingParameterNumber(
                "SCALE",
                "Facteur d’agrandissement",
                type=QgsProcessingParameterNumber.Integer,
                defaultValue=4
            )
        )
        self.addParameter(
            QgsProcessingParameterEnum(
                "MODEL",
                "Modèle ESRGAN",
                options=["realesrgan-x4plus", "realesrgan-x4plus-anime", "realesrnet-x4plus"],
                defaultValue=0
            )
        )
        self.addParameter(
            QgsProcessingParameterRasterDestination("OUTPUT_RASTER", "Raster de sortie")
        )

    def processAlgorithm(self, parameters, context, feedback):
        input_layer = self.parameterAsRasterLayer(parameters, "INPUT_RASTER", context)
        if input_layer is None:
            raise QgsProcessingException("Raster d’entrée invalide.")
        
        input_raster = input_layer.source()
        scale = self.parameterAsInt(parameters, "SCALE", context)
        model_index = self.parameterAsInt(parameters, "MODEL", context)
        output_raster = self.parameterAsOutputLayer(parameters, "OUTPUT_RASTER", context)

        exe = r"C:/SR/realesrgan-ncnn-vulkan-20220424-windows/realesrgan-ncnn-vulkan.exe"
        if not os.path.exists(exe):
            raise QgsProcessingException("Executable Real-ESRGAN introuvable : " + exe)

        # Fichiers temporaires
        tmp_png = os.path.splitext(input_raster)[0] + "_tmp.png"
        sr_png = os.path.splitext(input_raster)[0] + "_sr.png"

        model_list = ["realesrgan-x4plus", "realesrgan-x4plus-anime", "realesrnet-x4plus"]
        model_name = model_list[model_index]

        # -----------------------------
        # Conversion vers PNG
        feedback.pushInfo("Conversion en PNG…")
        subprocess.run(["gdal_translate", "-of", "PNG", input_raster, tmp_png], check=True)

        # -----------------------------
        # Super-résolution ESRGAN
        feedback.pushInfo("Exécution ESRGAN…")
        subprocess.run([exe, "-i", tmp_png, "-o", sr_png, "-n", model_name, "-s", str(scale)], check=True)

        # -----------------------------
        # Calcul de la nouvelle géotransformation
        src = gdal.Open(input_raster)
        gt = src.GetGeoTransform()
        proj = src.GetProjection()
        # Ajustement du pixel pour conserver l'emprise
        new_gt = (gt[0], gt[1]/scale, gt[2], gt[3], gt[4], gt[5]/scale)

        # Création du worldfile pour sr_png
        wld_path = sr_png.replace(".png", ".wld")
        with open(wld_path, "w") as f:
            f.write(f"{new_gt[1]}\n{new_gt[2]}\n{new_gt[4]}\n{new_gt[5]}\n{new_gt[0]}\n{new_gt[3]}\n")

        # Création du .prj
        srs = osr.SpatialReference()
        srs.ImportFromWkt(proj)
        prj_path = sr_png.replace(".png", ".prj")
        with open(prj_path, "w") as f:
            f.write(srs.ExportToWkt())

        # -----------------------------
        # Reprojection / création du TIF final
        feedback.pushInfo("Création du TIF final…")
        subprocess.run([
            "gdalwarp",
            "-t_srs", src.GetProjection(),
            sr_png, output_raster
        ], check=True)

        # -----------------------------
        # Nettoyage
        for f in [tmp_png, sr_png, wld_path, prj_path]:
            if os.path.exists(f):
                os.remove(f)

        return {"OUTPUT_RASTER": output_raster}

    def name(self):
        return "esrgan_super_resolution"

    def displayName(self):
        return "Super-résolution ESRGAN (NCNN)"

    def group(self):
        return "Super-Résolution"

    def groupId(self):
        return "super_resolution"

    def createInstance(self):
        return ESRGANSuperResolution()