WQ_GUI/src/core/glint_removal/Goodman.py

import numpy as np
import sys
# import preprocessing

try:
    from osgeo import gdal
    GDAL_AVAILABLE = True
except ImportError:
    GDAL_AVAILABLE = False
    print("警告: GDAL未安装，将使用numpy处理模式")

try:
    from tqdm import tqdm
    TQDM_AVAILABLE = True
except ImportError:
    TQDM_AVAILABLE = False
    # 如果tqdm不可用，定义一个简单的包装器
    def tqdm(iterable, desc=None, total=None, disable=None):
        return iterable

# 检测是否在 PyInstaller 打包环境（无控制台）
_is_frozen_gui = getattr(sys, "frozen", False) and (not hasattr(sys, 'stdout') or sys.stdout is None)

class Goodman:
    def __init__(self, im_aligned, NIR_lower = 25, NIR_upper = 37, A = 0.000019, B = 0.1, 
                 use_gdal=True, chunk_size=None, water_mask=None, output_path=None):
        """
        :param im_aligned (np.ndarray or str): band aligned and calibrated & corrected reflectance image
            可以是numpy数组或GDAL可读取的文件路径
        :param NIR_lower (int): band index which corresponds to 641.93nm, closest band to 640nm
        :param NIR_upper (int): band index which corresponds to 751.49nm, closest band to 750nm
        :param A (float): the values in Goodman et al's paper, using AVIRIS reflectance (rather than radiance) data
        :param B (float): the values in Goodman et al's paper, using AVIRIS reflectance (rather than radiance) data
            see Goodman et al, which corrects each pixel independently. The NIR radiance is subtracted from the radiance at each wavelength,
            but a wavelength-independent offset is also added. 
            it is not clear how A and B were chosen, but an optimization for a case where in situ data is
            available would enable values to be found
        :param use_gdal (bool): 是否使用GDAL加速处理（需要GDAL可用且输入为文件路径或大数组）
        :param chunk_size (int): 已废弃，不再使用分块处理，改为逐波段处理
        :param water_mask (np.ndarray or str or None): 水域掩膜，1表示水域，0表示非水域
            可以是numpy数组、栅格文件路径(.dat/.tif)或shapefile路径(.shp)
            如果为None，则处理全图
        :param output_path (str or None): 输出文件路径，如果提供则保存校正后的图像
            如果为None，则不保存
        """
        self.im_aligned = im_aligned
        self.NIR_lower = NIR_lower
        self.NIR_upper = NIR_upper
        self.A = A
        self.B = B
        self.use_gdal = use_gdal and GDAL_AVAILABLE
        self.chunk_size = chunk_size
        self.is_file_path = isinstance(im_aligned, str)
        self.output_path = output_path
        
        # 获取图像信息（需要在加载掩膜之前获取尺寸）
        if self.is_file_path:
            if not self.use_gdal:
                raise ValueError("输入为文件路径时，必须安装GDAL")
            self.dataset = gdal.Open(im_aligned, gdal.GA_ReadOnly)
            if self.dataset is None:
                raise ValueError(f"无法打开影像文件: {im_aligned}")
            self.height = self.dataset.RasterYSize
            self.width = self.dataset.RasterXSize
            self.n_bands = self.dataset.RasterCount
        else:
            self.dataset = None
            self.height = im_aligned.shape[0]
            self.width = im_aligned.shape[1]
            self.n_bands = im_aligned.shape[-1]
        
        # 加载水域掩膜（在获取图像尺寸之后）
        self.water_mask = self._load_water_mask(water_mask)
    
    def _load_water_mask(self, water_mask):
        """
        加载水域掩膜
        
        :param water_mask: 可以是None、numpy数组、文件路径(.dat/.tif)或shapefile路径(.shp)
        :return: numpy数组或None，1表示水域，0表示非水域
        """
        if water_mask is None:
            return None
        
        # 如果已经是numpy数组
        if isinstance(water_mask, np.ndarray):
            if water_mask.shape[:2] != (self.height, self.width):
                raise ValueError(f"掩膜尺寸 {water_mask.shape[:2]} 与图像尺寸 {(self.height, self.width)} 不匹配")
            return (water_mask > 0).astype(np.uint8)  # 确保是0/1掩膜
        
        # 如果是文件路径
        if isinstance(water_mask, str):
            if not GDAL_AVAILABLE:
                raise ValueError("使用文件路径作为掩膜时，必须安装GDAL")
            
            # 检查是否为shapefile
            if water_mask.lower().endswith('.shp'):
                # 从shp文件创建掩膜
                if self.is_file_path:
                    ref_path = self.im_aligned
                else:
                    raise ValueError("输入为numpy数组时，无法从shp文件创建掩膜（需要参考栅格）")
                
                try:
                    from osgeo import ogr
                    ref_dataset = gdal.Open(ref_path, gdal.GA_ReadOnly)
                    if ref_dataset is None:
                        raise ValueError(f"无法打开参考栅格文件: {ref_path}")
                    
                    geotransform = ref_dataset.GetGeoTransform()
                    projection = ref_dataset.GetProjection()
                    width = ref_dataset.RasterXSize
                    height = ref_dataset.RasterYSize
                    
                    # 创建内存中的栅格数据集
                    mem_driver = gdal.GetDriverByName('MEM')
                    mask_dataset = mem_driver.Create('', width, height, 1, gdal.GDT_Byte)
                    mask_dataset.SetGeoTransform(geotransform)
                    mask_dataset.SetProjection(projection)
                    
                    mask_band = mask_dataset.GetRasterBand(1)
                    mask_band.Fill(0)
                    
                    # 打开shp文件
                    shp_dataset = ogr.Open(water_mask)
                    if shp_dataset is None:
                        raise ValueError(f"无法打开shp文件: {water_mask}")
                    
                    layer = shp_dataset.GetLayer()
                    gdal.RasterizeLayer(mask_dataset, [1], layer, burn_values=[1])
                    
                    water_mask_array = mask_band.ReadAsArray()
                    
                    ref_dataset = None
                    mask_dataset = None
                    shp_dataset = None
                    
                    return (water_mask_array > 0).astype(np.uint8)
                except Exception as e:
                    raise ValueError(f"从shp文件创建掩膜时出错: {e}")
            else:
                # 栅格文件
                mask_dataset = gdal.Open(water_mask, gdal.GA_ReadOnly)
                if mask_dataset is None:
                    raise ValueError(f"无法打开掩膜文件: {water_mask}")
                
                mask_array = mask_dataset.GetRasterBand(1).ReadAsArray()
                mask_dataset = None
                
                if mask_array.shape != (self.height, self.width):
                    raise ValueError(f"掩膜尺寸 {mask_array.shape} 与图像尺寸 {(self.height, self.width)} 不匹配")
                
                return (mask_array > 0).astype(np.uint8)
        
        raise ValueError(f"不支持的掩膜类型: {type(water_mask)}")

    def _get_corrected_bands_numpy(self):
        """
        使用numpy处理（用于小图像或GDAL不可用时）
        
        注意：由于输入已经是numpy数组，数据已在内存中。
        此方法通过逐波段处理，避免同时创建多个校正后的波段数组。
        内存峰值 = 原始数组 + NIR波段(2个) + 当前处理的波段(1个)
        """
        # 预提取重复使用的NIR波段，避免在循环中重复访问
        # 这些波段会一直保存在内存中，因为它们需要用于所有波段的校正
        R_640 = self.im_aligned[:,:,self.NIR_lower]
        R_750 = self.im_aligned[:,:,self.NIR_upper]
        # 预计算常量部分
        diff_640_750 = R_640 - R_750
        corrected_bands = []
        
        # 获取水域掩膜（如果存在）
        water_mask_bool = self.water_mask.astype(bool) if self.water_mask is not None else None
        
        # 逐波段处理：每次只处理一个波段，处理完后立即添加到结果列表
        for i in tqdm(range(self.n_bands), desc="处理波段 (numpy)", total=self.n_bands, disable=_is_frozen_gui):
            # 获取当前波段（这是数组视图，不是复制）
            R = self.im_aligned[:,:,i]
            # 优化计算：减少中间数组创建
            corrected_band = R - R_750 + self.A + self.B * diff_640_750
            # 使用np.maximum原地操作，将负值设为0
            np.maximum(corrected_band, 0, out=corrected_band)
            
            # 如果存在水域掩膜，只对水域区域应用校正
            if water_mask_bool is not None:
                corrected_band = np.where(water_mask_bool, corrected_band, R)
            
            # 立即添加到结果列表（corrected_band会保留在列表中）
            corrected_bands.append(corrected_band)
        return corrected_bands

    def _get_corrected_bands_gdal(self, out_dataset=None):
        """
        使用GDAL逐波段处理。

        Args:
            out_dataset: 若提供，每处理完一个波段立即写入此数据集（流式模式）；
                        若为 None，累积到列表中返回（传统模式，大图像慎用）。

        Returns:
            若 out_dataset 为 None，返回波段列表；
            若 out_dataset 不为 None，返回 None（波段已直接写入输出文件）
        """
        # 获取NIR波段对象（用于所有波段的校正）
        band_640 = self.dataset.GetRasterBand(self.NIR_lower + 1)
        band_750 = self.dataset.GetRasterBand(self.NIR_upper + 1)

        # 读取NIR波段（全程保留在内存中）
        R_640 = band_640.ReadAsArray().astype(np.float32)
        R_750 = band_750.ReadAsArray().astype(np.float32)
        diff_640_750 = R_640 - R_750
        del R_640  # 释放不再需要的 R_640

        # 获取水域掩膜
        water_mask_bool = self.water_mask.astype(bool) if self.water_mask is not None else None

        # 输出模式
        streaming = (out_dataset is not None)
        corrected_bands = [] if not streaming else None

        # 逐波段处理
        for i in tqdm(range(self.n_bands), desc="处理波段 (GDAL)",
                       total=self.n_bands, disable=_is_frozen_gui):
            # 读取当前波段
            current_band = self.dataset.GetRasterBand(i + 1)
            R = current_band.ReadAsArray().astype(np.float32)

            # 校正计算（原地操作以减少临时分配）
            corrected = R - R_750
            corrected += self.A
            corrected += self.B * diff_640_750
            np.maximum(corrected, 0, out=corrected)

            # 水域掩膜：只在有水的地方用校正值，陆地保持原值
            if water_mask_bool is not None:
                corrected = np.where(water_mask_bool, corrected, R)

            if streaming:
                # 流式模式：立即写入磁盘并释放
                out_band = out_dataset.GetRasterBand(i + 1)
                out_band.WriteArray(corrected)
                out_band.FlushCache()
                del corrected, R
            else:
                # 传统模式：累积到列表
                corrected_bands.append(corrected)
                del R

        # 清理
        del R_750, diff_640_750

        return corrected_bands

    def _get_corrected_bands_gdal_mem(self):
        """使用GDAL内存驱动处理numpy数组，逐波段处理"""
        # 创建内存数据集
        driver = gdal.GetDriverByName('MEM')
        mem_dataset = driver.Create('', self.width, self.height, self.n_bands, gdal.GDT_Float32)
        
        # 将numpy数组写入内存数据集（显示进度）
        for i in tqdm(range(self.n_bands), desc="加载波段到内存", total=self.n_bands, disable=_is_frozen_gui):
            band = mem_dataset.GetRasterBand(i + 1)
            band.WriteArray(self.im_aligned[:,:,i])
            band.FlushCache()
        
        # 临时保存原始dataset引用
        original_dataset = self.dataset
        self.dataset = mem_dataset
        
        try:
            # 使用逐波段处理方法
            result = self._get_corrected_bands_gdal()
        finally:
            # 恢复原始dataset
            self.dataset = original_dataset
            mem_dataset = None
        
        return result

    def _save_corrected_bands(self, corrected_bands):
        """
        保存校正后的波段到文件（BSQ格式，ENVI格式）
        
        注意：为了节省内存，直接逐波段写入，不先堆叠成完整数组
        
        :param corrected_bands: 校正后的波段列表
        """
        if not GDAL_AVAILABLE:
            raise ImportError("GDAL未安装，无法保存影像文件")
        
        if self.output_path is None:
            return
        
        import os
        # 确保输出目录存在
        output_dir = os.path.dirname(self.output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)
        
        # 从第一个波段获取尺寸信息（避免堆叠所有波段）
        if not corrected_bands:
            raise ValueError("校正后的波段列表为空")
        first_band = corrected_bands[0]
        height, width = first_band.shape
        n_bands = len(corrected_bands)
        
        # 获取地理变换和投影信息
        if self.is_file_path and self.dataset is not None:
            geotransform = self.dataset.GetGeoTransform()
            projection = self.dataset.GetProjection()
        else:
            # 如果没有地理信息，使用默认值
            geotransform = (0, 1, 0, 0, 0, -1)
            projection = ""
        
        # 强制使用ENVI格式（BSQ格式），确保文件扩展名为.bsq
        base_path, ext = os.path.splitext(self.output_path)
        # 如果扩展名不是.bsq，使用基础路径添加.bsq
        if ext.lower() != '.bsq':
            bsq_path = base_path + '.bsq'
        else:
            bsq_path = self.output_path
        
        # 使用ENVI驱动（默认就是BSQ格式）
        driver = gdal.GetDriverByName('ENVI')
        if driver is None:
            raise ValueError("无法创建ENVI格式文件，ENVI驱动不可用")
        
        # 创建ENVI格式数据集（会自动生成.hdr文件）
        dataset = driver.Create(bsq_path, width, height, n_bands, gdal.GDT_Float32)
        if dataset is None:
            raise ValueError(f"无法创建输出文件: {bsq_path}")
        
        try:
            # 设置地理变换和投影
            if geotransform:
                dataset.SetGeoTransform(geotransform)
            if projection:
                dataset.SetProjection(projection)
            
            # 直接逐波段写入（不先堆叠，节省内存）
            for i in tqdm(range(n_bands), desc="保存波段", total=n_bands, disable=_is_frozen_gui):
                band = dataset.GetRasterBand(i + 1)
                # 直接从列表中获取波段并写入，避免创建完整数组
                band.WriteArray(corrected_bands[i])
                band.FlushCache()
        finally:
            dataset = None
        
        # 检查.hdr文件是否已创建
        hdr_path = bsq_path + '.hdr'
        if os.path.exists(hdr_path):
            print(f"校正后的图像已保存至: {bsq_path} (BSQ格式)")
            print(f"头文件已保存至: {hdr_path}")
        else:
            print(f"校正后的图像已保存至: {bsq_path} (BSQ格式)")
            print(f"警告: 未检测到.hdr文件，但GDAL应该已自动创建")

    def get_corrected_bands(self):
        """
        获取校正后的波段（自动选择最优处理模式）

        内存优化: 当 output_path 已设置时，使用流式模式逐波段直接写入磁盘，
        避免在内存中累积全部波段（对 6522×13215×150 的大影像可节省 ~46 GB）。

        :return: 校正后的波段列表（流式模式返回 None，波段已在输出文件中）
        """
        # ── 流式模式：output_path 已设置 → 逐波段处理+立即写入 → 零累积 ──
        if self.output_path is not None:
            return self._get_corrected_bands_streaming()

        # ── 传统模式：output_path 为空 → 返回波段列表 ──
        if self.is_file_path:
            if self.use_gdal:
                return self._get_corrected_bands_gdal(out_dataset=None)
            else:
                raise ValueError("输入为文件路径时，必须安装GDAL")
        else:
            if self.use_gdal and self.height * self.width * self.n_bands > 100000000:
                return self._get_corrected_bands_gdal_mem()
            else:
                return self._get_corrected_bands_numpy()
    
    def _get_corrected_bands_streaming(self):
        """流式处理：逐波段校正并直接写入输出文件，不累积内存

        适用于大尺度影像（如 6522×13215×150）。
        内存峰值 ≈ 3 个全波段数组（NIR×2 + 当前波段），而非全部 150 个波段。

        :return: None（波段已在输出文件中）
        """
        import os

        # ── 创建输出文件 ──
        base_path, ext = os.path.splitext(self.output_path)
        bsq_path = base_path + '.bsq' if ext.lower() != '.bsq' else self.output_path
        output_dir = os.path.dirname(bsq_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)

        driver = gdal.GetDriverByName('ENVI')
        out_ds = driver.Create(
            bsq_path, self.width, self.height, self.n_bands, gdal.GDT_Float32
        )
        if out_ds is None:
            raise ValueError(f"无法创建输出文件: {bsq_path}")

        # ── 设置地理参考 ──
        if self.is_file_path and self.dataset is not None:
            out_ds.SetGeoTransform(self.dataset.GetGeoTransform())
            out_ds.SetProjection(self.dataset.GetProjection())

        try:
            # 逐波段处理 + 立即写入（波段在 _get_corrected_bands_gdal 的循环中
            # 由 WriteArray 写入 → FlushCache → del，不会累积）
            self._get_corrected_bands_gdal(out_dataset=out_ds)
        finally:
            out_ds = None  # 关闭文件，确保数据落盘

        # ── 日志 ──
        hdr_path = bsq_path + '.hdr'
        if os.path.exists(hdr_path):
            print(f"校正后的图像已保存至: {bsq_path} (BSQ格式, 流式写入)")
        else:
            print(f"校正后的图像已保存至: {bsq_path} (BSQ格式)")
            print("警告: 未检测到.hdr文件，但GDAL应该已自动创建")

        return None

    def __del__(self):
        """清理资源"""
        if self.dataset is not None and self.is_file_path:
            self.dataset = None