pip install opencv-python rasterio numpy tensorflow

import rasterio
import numpy as np
import cv2
from tensorflow.keras.models import load_model
import os
import tempfile

def load_geotiff(image_path):
    """加载GeoTIFF图像并返回数据和元数据"""
    with rasterio.open(image_path) as src:
        image = src.read()
        meta = src.meta.copy()
        profile = src.profile
    return image, meta, profile

def save_geotiff(output_path, data, meta):
    """保存为GeoTIFF，保留原始地理信息"""
    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(data)

def preprocess_image(image):
    """预处理图像数据"""
    # 将数据归一化到0-1范围
    image = image.astype('float32')
    if image.max() > 1:
        image /= 255.0
    return image

def postprocess_image(image):
    """后处理图像数据"""
    image = (image * 255).clip(0, 255).astype('uint8')
    return image

def upscale_image_fsrcnn(image, model_path='FSRCNN_x2.h5', scale=2):
    """使用FSRCNN模型进行超分辨率处理"""
    # 加载预训练的FSRCNN模型
    try:
        model = load_model(model_path)
    except:
        print(f"无法加载模型 {model_path}")
        print("请从 https://github.com/Saafke/FSRCNN_Tensorflow/tree/master/models 下载预训练模型")
        return None

    # 获取原始图像尺寸
    original_shape = image.shape

    # 调整尺寸以便模型处理 (FSRCNN通常需要特定尺寸)
    if len(image.shape) == 2:  # 单通道
        h, w = image.shape
        # 填充图像到可被scale整除的尺寸
        pad_h = (scale - h % scale) % scale
        pad_w = (scale - w % scale) % scale
        image = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect')
        image = np.expand_dims(image, axis=-1)
        image = np.expand_dims(image, axis=0)
    else:  # 多通道
        c, h, w = image.shape
        # 填充图像到可被scale整除的尺寸
        pad_h = (scale - h % scale) % scale
        pad_w = (scale - w % scale) % scale
        image = np.pad(image, ((0, 0), (0, pad_h), (0, pad_w)), mode='reflect')
        # 调整通道顺序为HWC
        image = np.transpose(image, (1, 2, 0))
        image = np.expand_dims(image, axis=0)

    # 使用模型进行超分辨率
    sr_image = model.predict(image)

    # 移除填充并恢复原始形状
    sr_image = sr_image[0]
    if len(original_shape) == 2:  # 单通道
        sr_image = sr_image[:original_shape[0]*scale, :original_shape[1]*scale, 0]
    else:  # 多通道
        sr_image = sr_image[:original_shape[1]*scale, :original_shape[2]*scale, :]
        sr_image = np.transpose(sr_image, (2, 0, 1))

    return sr_image

def geotiff_super_resolution(input_path, output_path, model_path='FSRCNN_x2.h5', scale=2):
    """GeoTIFF图像超分辨率主函数"""
    # 1. 加载GeoTIFF图像
    image, meta, profile = load_geotiff(input_path)

    # 2. 预处理
    processed_image = preprocess_image(image)

    # 3. 超分辨率处理
    sr_image = upscale_image_fsrcnn(processed_image, model_path, scale)
    if sr_image is None:
        return False

    # 4. 后处理
    sr_image = postprocess_image(sr_image)

    # 5. 更新元数据以反映新的尺寸
    meta['width'] = meta['width'] * scale
    meta['height'] = meta['height'] * scale
    meta['transform'] = rasterio.Affine(
        meta['transform'].a / scale, meta['transform'].b, meta['transform'].c,
        meta['transform'].d, meta['transform'].e / scale, meta['transform'].f
    )

    # 6. 保存结果
    save_geotiff(output_path, sr_image, meta)

    return True

if __name__ == "__main__":
    # 示例用法
    input_geotiff = "input.tif"  # 替换为你的输入GeoTIFF路径
    output_geotiff = "output_sr.tif"  # 输出文件路径
    model_path = "FSRCNN_x2.h5"  # 预训练模型路径

    # 执行超分辨率
    success = geotiff_super_resolution(input_geotiff, output_geotiff, model_path)

    if success:
        print(f"超分辨率处理完成，结果已保存到 {output_geotiff}")
    else:
        print("处理失败")