feat(raw_to_cog): 改进 NoData 值处理逻辑并添加 Warp 转换支持.

utils/raw_to_cog.py

@@ -8,7 +8,7 @@ COG (Cloud Optimized GeoTIFF) 是一种优化设计的 GeoTIFF 文件格式, 特
 
 -------------------------------------------------------------------------------
 Authors: CVEO Team
-Last Updated: 2026-01-30
+Last Updated: 2026-02-03
 ===============================================================================
 """
 
@@ -31,6 +31,43 @@ gdal.SetConfigOption("GDAL_NUM_THREADS", "ALL_CPUS")
 gdal.SetConfigOption("GDAL_CACHEMAX", "1024")
 
 
+def check_and_get_image_info(ds):
+    """
+    检查影像基本信息以及是否存在 NaN 或 NoData 值
+
+    Returns
+    -------
+
+    tuple: (has_invalid, original_no_data)
+        has_invalid: bool
+            是否存在 NaN 或 NoData 值
+        original_no_data: float or None
+            原影像定义的 NoData 值
+    """
+    width = ds.RasterXSize
+    height = ds.RasterYSize
+    band_count = ds.RasterCount
+
+    logging.info(f"影像尺寸: {width}x{height}, 波段数: {band_count}")
+
+    has_invalid = False
+    original_no_data = None
+
+    # 读取第一个波段的 NoData 值作为参考
+    if band_count > 0:
+        original_no_data = ds.GetRasterBand(1).GetNoDataValue()
+
+    if original_no_data is not None:
+        logging.info(f"波段 1: 原影像定义的 NoData 值为 {original_no_data}")
+        has_invalid = True
+    else:
+        logging.info("波段 1: 原影像未定义 NoData 值")
+
+    # 仅检查是否存在 NaN 值
+
+    return has_invalid, original_no_data
+
+
 def tif_to_cog(
     input_path: str,
     output_path: str,
@@ -65,30 +102,51 @@ def tif_to_cog(
     if ds is None:
         logging.error(f"无法打开输入文件: {input_path}")
         return
-    # 查看原数据是否已设置无效值
+
-    band = ds.GetRasterBand(1)
+    # 1. 检查影像信息和无效值
-    original_no_data = band.GetNoDataValue()
+    has_invalid, original_no_data = check_and_get_image_info(ds)
-    original_data_type = band.DataType
+
-    if np.isnan(no_data):
+    # 2. 确定最终使用的 NoData 值
+    target_no_data = no_data
+
+    # 如果用户没有指定特定的 no_data (即传入的是 nan), 则尝试沿用原数据的 nodata
+    # 如果原数据也没有 nodata, 则根据 output_type 设置默认值
+    if np.isnan(target_no_data):
         if original_no_data is None:
-            # 根据输出数据类型设置无效值
             if output_type == gdal.GDT_Int16:
-                no_data = -32768
+                target_no_data = -32768
             elif output_type == gdal.GDT_Float32:
-                no_data = -9999
+                target_no_data = -9999
             elif output_type == gdal.GDT_Byte:
-                no_data = None
+                target_no_data = None
             else:
-                no_data = np.nan
+                target_no_data = np.nan
-            logging.warning(f"原数据未设无效值, 现已设为: {no_data}")
+            logging.warning(
+                f"原数据未设无效值且未指定新无效值, 现根据输出类型设为: {target_no_data}"
+            )
         else:
-            no_data = original_no_data
+            target_no_data = original_no_data
-            logging.info(f"原数据无效值为: {no_data}")
+            logging.info(f"沿用原数据无效值: {target_no_data}")
-    translate_options = gdal.TranslateOptions(
+    else:
-        format="COG",  # 输出为 COG 格式, 自动构建金字塔
+        logging.info(f"使用指定的新无效值: {target_no_data}")
-        outputType=output_type,
+
-        noData=no_data,
+    # 3. 转换逻辑
-        creationOptions=[
+    # 若原始影像存在无效值(NaN 或 NoData) 且 指定了新的 NoData 值, 且 新旧 NoData 值不一致 / 存在 NaN 值时使用重映射 (gdal.Warp)
+    use_warp = False
+    if target_no_data is not None and not np.isnan(target_no_data):
+        if has_invalid:
+            if original_no_data != target_no_data:
+                use_warp = True
+                logging.info(
+                    f"原始影像存在无效值且已指定新无效值 ({original_no_data} -> {target_no_data}), 将使用 Warp 进行处理"
+                )
+            else:
+                pass
+
+    # 开始转换
+    logging.info(f"开始转换为 COG 格式...")
+
+    creation_options = [
         f"COMPRESS={compress}",
         "PREDICTOR=2",  # 差值预测, 利于影像压缩
         "NUM_THREADS=ALL_CPUS",
@@ -96,12 +154,36 @@ def tif_to_cog(
         # "ZLEVEL=4",  # 压缩级别, 支持 1-9, 默认为 6, COG 不支持设置
         # "TILED=YES", # 分块参数, COG 格式自带分块, 不支持手动设置
         # "PHOTOMETRIC=RGB", # 可视化参数, COG 格式不支持设置
-        ],
+    ]
+
+    if use_warp:
+        # 使用 Warp
+        warp_options = gdal.WarpOptions(
+            format="COG",  # 输出为 COG 格式, 自动构建金字塔
+            outputType=output_type,
+            dstNodata=target_no_data,
+            # srcNodata 未指定时, 将自动读取源数据的 nodata
+            # 若源数据存在 NaN 却未定义 nodata, 将自动处理 float 的 NaN
+            # srcNodata=original_no_data,
+            creationOptions=creation_options,
             callback=gdal.TermProgress_nocb,  # 进度回调, 不显示进度条
+            multithread=True,
+        )
+        gdal.Warp(str(output_path), ds, options=warp_options)
+    else:
+        # 使用 Translate (无需修改原数据时使用, 更高效)
+        # 但 gdal.Translate 无法将原始 NoData 直接转为新的 nodata
+        translate_options = gdal.TranslateOptions(
+            format="COG",
+            outputType=output_type,
+            noData=target_no_data,
+            creationOptions=creation_options,
+            callback=gdal.TermProgress_nocb,
         )
-    logging.info(f"开始转换为 COG 格式...")
         gdal.Translate(str(output_path), ds, options=translate_options)
+
     logging.info(f"结果已保存到: {output_path}")
+    ds = None
     return
 
 
@@ -123,10 +205,9 @@ def main(input_path, output_path, output_type=gdal.GDT_Float32):
 
 if __name__ == "__main__":
     # 输入目录: 包含分块tif影像的根目录
-    # input_root = Path(r"D:\CVEOdata\RS_Data\Terrain")
+    input_root = Path(r"D:\CVEOdata\RS_Data\Terrain")
     # 输出目录: 存放最终COG结果的目录
-    # output_root = Path(r"D:\CVEOdata\RS_Data\Terrain")
+    output_root = Path(r"D:\CVEOdata\RS_Data\Terrain")
     input_path = input_root / "GLO30.DSM.2014.Hubei.30m.tif"
     output_path = output_root / "GLO30.DSM.2014.Hubei.30m.tif"
-    # output_type = gdal.GDT_Float32
+    output_type = gdal.GDT_Byte
-    # main(input_path, output_path, output_type)