xhong/NASA_EarthData_Script
1
0
Fork 0
Code Issues Pull Requests Projects Releases 1 Wiki Activity

refactor(sr2rgb): 使用xarray重构地表反射率合成RGB图像功能.

Browse Source
This commit is contained in:
谢泓 2025-09-05 21:07:51 +08:00
parent 3d5d159d7f
commit bda6f0a1ef
1 changed files with 62 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

103
utils/sr2rgb.py
View File

@ -1,5 +1,5 @@
""" """
将 COG 格式的 Red, Green, Blue 单波段地表反射率图像合成 RGB 图像 将 COG 格式的 Red, Green, Blue 单波段地表反射率图像合成为 RGB 图像
1. 对比度线性拉伸至 0-255; 1. 对比度线性拉伸至 0-255;
2. 合并波段; 2. 合并波段;
@ -7,10 +7,12 @@
""" """
import os import os
import rasterio as rio import xarray as xr
from rioxarray import open_rasterio
import numpy as np import numpy as np
def sr2rgb(red_path, green_path, blue_path, output_path):
def sr2rgb(red_path: str, green_path: str, blue_path: str, output_path: str) -> None:
""" """
将红、绿、蓝三个单波段地表反射率图像合成为RGB图像 将红、绿、蓝三个单波段地表反射率图像合成为RGB图像
@ -25,73 +27,78 @@ def sr2rgb(red_path, green_path, blue_path, output_path):
if not os.path.exists(path): if not os.path.exists(path):
raise FileNotFoundError(f"文件不存在: {path}") raise FileNotFoundError(f"文件不存在: {path}")
print(f"正在处理文件:")
print(f" 红波段: {red_path}")
print(f" 绿波段: {green_path}")
print(f" 蓝波段: {blue_path}")
# 读取三个波段数据 # 读取三个波段数据
with rio.open(red_path) as red_src: red_band = open_rasterio(red_path, masked=True).squeeze(dim="band", drop=True)
red_band = red_src.read(1) green_band = open_rasterio(green_path, masked=True).squeeze(dim="band", drop=True)
profile = red_src.profile blue_band = open_rasterio(blue_path, masked=True).squeeze(dim="band", drop=True)
print(f"红波段形状: {red_band.shape}, 数据类型: {red_band.dtype}") # 暂存元数据
y_coords = red_band.y
x_coords = red_band.x
crs = red_band.rio.crs
transform = red_band.rio.transform()
with rio.open(green_path) as green_src:
green_band = green_src.read(1)
print(f"绿波段形状: {green_band.shape}, 数据类型: {green_band.dtype}")
with rio.open(blue_path) as blue_src:
blue_band = blue_src.read(1)
print(f"蓝波段形状: {blue_band.shape}, 数据类型: {blue_band.dtype}")
# 线性拉伸到0-255范围
def stretch_band(band): def stretch_band(band):
# 处理NaN值 """
线性拉伸到0-255范围
"""
# 处理NaN值与负值
band_no_nan = np.where(np.isnan(band), 0, band) band_no_nan = np.where(np.isnan(band), 0, band)
band_no_nan = np.where(band_no_nan < 0, 0, band_no_nan)
band_min = np.min(band_no_nan) band_min = np.min(band_no_nan)
band_max = np.max(band_no_nan) band_max = np.max(band_no_nan)
# 避免除零错误 # 避免除零错误
if band_max == band_min: if band_max == band_min:
stretched = np.zeros_like(band_no_nan, dtype=np.uint8) stretched = np.zeros_like(band_no_nan, dtype=np.uint8)
else: else:
stretched = ((band_no_nan - band_min) / (band_max - band_min) * 255).astype(np.uint8) stretched = ((band_no_nan - band_min) / (band_max - band_min) * 255).astype(
np.uint8
)
return stretched return stretched
red_stretched = stretch_band(red_band) red_stretched = stretch_band(red_band.values)
green_stretched = stretch_band(green_band) green_stretched = stretch_band(green_band.values)
blue_stretched = stretch_band(blue_band) blue_stretched = stretch_band(blue_band.values)
# 合并三个波段为RGB图像 # 合并三个波段为RGB图像
rgb_array = np.dstack((red_stretched, green_stretched, blue_stretched)) rgb_array = xr.DataArray(
np.dstack((red_stretched, green_stretched, blue_stretched)),
# 更新元数据 dims=("y", "x", "band"),
profile.update( coords={"band": [1, 2, 3], "y": y_coords, "x": x_coords},
dtype=rio.uint8,
count=3,
photometric='RGB',
nodata=0 # 设置nodata为0，因为uint8的范围是0-255
) )
# 转置维度顺序以符合rioxarray要求
# 保存RGB图像 rgb_array = rgb_array.transpose("band", "y", "x")
with rio.open(output_path, 'w', **profile) as dst: # 写入元数据
for i in range(3): rgb_array.rio.write_crs(crs, inplace=True)
dst.write(rgb_array[:, :, i], i + 1) rgb_array.rio.write_transform(transform, inplace=True)
rgb_array.rio.write_nodata(0, inplace=True)
# 保存为TIFF文件
rgb_array.rio.to_raster(output_path, dtype="uint8")
print(f"RGB图像已保存到: {output_path}") print(f"RGB图像已保存到: {output_path}")
return
if __name__ == "__main__": if __name__ == "__main__":
# tif_dir = "D:\\NASA_EarthData_Script\\data\\HLS\\2024\\2024012" # tif_dir = "D:\\NASA_EarthData_Script\\data\\HLS\\2024\\2024012"
# red_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.RED.subset.tif") # red_path = os.path.join(
# green_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.GREEN.subset.tif") # tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.RED.subset.tif"
# blue_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.BLUE.subset.tif") # )
# green_path = os.path.join(
# tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.GREEN.subset.tif"
# )
# blue_path = os.path.join(
# tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.BLUE.subset.tif"
# )
# output_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.RGB.tif") # output_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.RGB.tif")
tif_dir = "D:\\NASA_EarthData_Script\\data\\HLS\\2025\\2025011" tif_dir = "D:\\NASA_EarthData_Script\\data\\HLS\\2025\\2025011"
red_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.RED.subset.tif") red_path = os.path.join(
green_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.GREEN.subset.tif") tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.RED.subset.tif"
blue_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.BLUE.subset.tif") )
green_path = os.path.join(
tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.GREEN.subset.tif"
)
blue_path = os.path.join(
tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.BLUE.subset.tif"
)
output_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.RGB.tif") output_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.RGB.tif")
sr2rgb(red_path, green_path, blue_path, output_path) sr2rgb(red_path, green_path, blue_path, output_path)