From a4c81d08139137d4286d7763046795e821f4681f Mon Sep 17 00:00:00 2001
From: xhong <xhong_17@outlook.com>
Date: Wed, 7 Jan 2026 11:51:31 +0800
Subject: [PATCH] =?UTF-8?q?fix(sr2rgb):=20=E6=B7=BB=E5=8A=A0=E8=BD=AC?=
 =?UTF-8?q?=E6=8D=A2=E5=89=8D=E5=90=8E=E5=AF=B9NoData=E5=80=BC=E7=9A=84?=
 =?UTF-8?q?=E5=A4=84=E7=90=86.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 utils/sr2rgb.py | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/utils/sr2rgb.py b/utils/sr2rgb.py
index 43a92e3..4498390 100644
--- a/utils/sr2rgb.py
+++ b/utils/sr2rgb.py
@@ -24,7 +24,10 @@ gdal.UseExceptions()
 
 
 def normalize_band(
-    band: np.ndarray, method: str = "minmax", percentile: float = 2.0
+    band: np.ndarray,
+    method: str = "minmax",
+    percentile: float = 2.0,
+    nodata: Optional[float] = None,
 ) -> np.ndarray:
     """
     将波段数据归一化到 0-255 uint8
@@ -37,16 +40,22 @@ def normalize_band(
         归一化方法 ('minmax', 'percentile', 'clip'), by default "minmax"
     percentile: float, optional
         百分比截断参数 (仅用于 'percentile' 方法), by default 2.0
+    nodata: float, optional
+        NoData 值, 如果不为 None, 则将该值替换为 NaN, by default None
     """
     band = np.asarray(band, dtype=np.float32)
 
+    # 处理 NoData 值
+    if nodata is not None:
+        band[np.isclose(band, nodata)] = np.nan
+
     if method == "clip":
-        # 对应原 render8bit 逻辑: 直接截断并转换
+        # 直接截断并转换
         band[~np.isfinite(band)] = 0.0
         return np.clip(band, 0.0, 255.0).astype(np.uint8)
 
     if method == "percentile":
-        # 对应原 linear_stretch 逻辑
+        # 百分比截断并线性拉伸
         lower = np.nanpercentile(band, percentile)
         upper = np.nanpercentile(band, 100 - percentile)
     else:  # minmax
@@ -59,6 +68,8 @@ def normalize_band(
         return np.zeros_like(band, dtype=np.uint8)
 
     stretched = (band - lower) / (upper - lower) * 255.0
+    # 处理拉伸后的 NaN 值 (原始数据中的 NaN 会传播)
+    stretched[~np.isfinite(stretched)] = 0.0
     return np.clip(stretched, 0, 255).astype(np.uint8)
 
 
@@ -107,9 +118,12 @@ def render_image_rgb(
     chans = []
     for b in bands:
         band = ds.GetRasterBand(int(b))
+        nodata = band.GetNoDataValue()
         arr = band.ReadAsArray(0, 0, xsize, ysize)
         chans.append(
-            normalize_band(arr, method=normalization_method, percentile=percentile)
+            normalize_band(
+                arr, method=normalization_method, percentile=percentile, nodata=nodata
+            )
         )
 
     rgb = np.dstack(chans)
@@ -134,6 +148,7 @@ def render_image_rgb(
 
     ds = None
     Image.fromarray(rgb, mode="RGB").save(png_path, format="PNG")
+    print(f"已完成RGB图像转换, 保存至: {png_path}")
     return png_path
 
 
@@ -216,7 +231,7 @@ if __name__ == "__main__":
     # combine_bands_to_rgb(red_path, green_path, blue_path, output_path)
 
     imgs_dir = Path(r"D:\CVEOProjects\prjaef\aef-backend-demo\media\imgs")
-    region_name = "Target3"
+    region_name = "Target4"
     year = 2025
     img_name = f"S1_S2_{region_name}_{year}.tif"
     img_tif_path = imgs_dir / region_name / img_name