feat(utils): 重构原始影像数据转换RGBA图像的工具函数.
This commit is contained in:
parent
7347cd60bb
commit
32d629aede
@ -1,19 +1,26 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
将地表反射率图像转换为 8bit RGB 图像
|
||||
===============================================================================
|
||||
将原始数值影像转换为 8bit RGBA 图像
|
||||
|
||||
- Red, Green, Blue 单波段图像
|
||||
- 多波段图像
|
||||
- 支持 Red, Green, Blue 等单波段影像
|
||||
- 支持多波段影像
|
||||
|
||||
1. 对比度线性拉伸至 0-255;
|
||||
2. 合并波段;
|
||||
3. 保存为 uint8 格式 RGB 图像;
|
||||
1. 将原始数据中 NoData 值替换为 NaN, 并将其设置为 Alpha 通道
|
||||
2. 对比度线性拉伸至 0-255;
|
||||
3. 合并包含 Alpha 通道在内的 4 个波段;
|
||||
4. 保存为 uint8 格式 RGBA 图像
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
Authors: CVEO Team
|
||||
Last Updated: 2025-01-08
|
||||
===============================================================================
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from osgeo import gdal
|
||||
@ -82,7 +89,7 @@ def render_image_rgb(
|
||||
output_suffix: str = "_rgb.png",
|
||||
) -> str:
|
||||
"""
|
||||
通用 GDAL 数据源处理函数, 读取多波段图像并转换为 RGB 图像.
|
||||
通用 GDAL 数据源处理函数, 读取多波段图像并转换为 RGBA 图像.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@ -95,7 +102,7 @@ def render_image_rgb(
|
||||
percentile : float, optional
|
||||
百分比截断参数 (仅用于 'percentile' 方法), by default 0.
|
||||
save_tiff : bool, optional
|
||||
是否保存为 RGB TIFF 文件, by default True.
|
||||
是否保存为 RGBA TIFF 文件, by default True.
|
||||
output_suffix : str, optional
|
||||
输出文件后缀, by default "_rgb.png".
|
||||
|
||||
@ -116,17 +123,32 @@ def render_image_rgb(
|
||||
bands = bands_lst[:3]
|
||||
|
||||
chans = []
|
||||
# 初始化 Alpha 掩膜 (全 True 表示全不透明/有效)
|
||||
alpha_mask = np.ones((ysize, xsize), dtype=bool)
|
||||
|
||||
for b in bands:
|
||||
band = ds.GetRasterBand(int(b))
|
||||
nodata = band.GetNoDataValue()
|
||||
arr = band.ReadAsArray(0, 0, xsize, ysize)
|
||||
|
||||
# 更新 Alpha 掩膜: 任何波段为 NoData 或 NaN，则该像素透明
|
||||
if nodata is not None:
|
||||
alpha_mask &= ~np.isclose(arr, nodata)
|
||||
|
||||
if np.issubdtype(arr.dtype, np.floating):
|
||||
alpha_mask &= np.isfinite(arr)
|
||||
|
||||
chans.append(
|
||||
normalize_band(
|
||||
arr, method=normalization_method, percentile=percentile, nodata=nodata
|
||||
)
|
||||
)
|
||||
|
||||
rgb = np.dstack(chans)
|
||||
# 生成包含 Alpha 通道的 RGBA
|
||||
alpha_channel = np.where(alpha_mask, 255, 0).astype(np.uint8)
|
||||
chans.append(alpha_channel)
|
||||
|
||||
rgba = np.dstack(chans)
|
||||
|
||||
# 构建输出路径
|
||||
base_path = str(tiff_path)
|
||||
@ -138,7 +160,8 @@ def render_image_rgb(
|
||||
if save_tiff:
|
||||
tiff_output_path = base_path + "_rgb.tif"
|
||||
driver = gdal.GetDriverByName("GTiff")
|
||||
out_ds = driver.Create(tiff_output_path, xsize, ysize, 3, gdal.GDT_Byte)
|
||||
# 创建 4 波段 (RGBA)
|
||||
out_ds = driver.Create(tiff_output_path, xsize, ysize, 4, gdal.GDT_Byte)
|
||||
out_ds.SetProjection(ds.GetProjection())
|
||||
out_ds.SetGeoTransform(ds.GetGeoTransform())
|
||||
for i, chan in enumerate(chans):
|
||||
@ -147,8 +170,8 @@ def render_image_rgb(
|
||||
out_ds = None
|
||||
|
||||
ds = None
|
||||
Image.fromarray(rgb, mode="RGB").save(png_path, format="PNG")
|
||||
print(f"已完成RGB图像转换, 保存至: {png_path}")
|
||||
Image.fromarray(rgba, mode="RGBA").save(png_path, format="PNG")
|
||||
print(f"已完成RGBA图像转换, 保存至: {png_path}")
|
||||
return png_path
|
||||
|
||||
|
||||
@ -156,7 +179,7 @@ def combine_bands_to_rgb(
|
||||
red_path: str, green_path: str, blue_path: str, output_path: str
|
||||
) -> None:
|
||||
"""
|
||||
将红, 绿, 蓝三个单波段图像文件合并为 RGB 图像 (GeoTIFF).
|
||||
将红, 绿, 蓝三个单波段图像文件合并为 RGBA 图像 (GeoTIFF).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@ -167,7 +190,7 @@ def combine_bands_to_rgb(
|
||||
blue_path : str
|
||||
蓝色波段文件路径.
|
||||
output_path : str
|
||||
输出 RGB 图像路径.
|
||||
输出 RGBA 图像路径.
|
||||
"""
|
||||
paths = [red_path, green_path, blue_path]
|
||||
for path in paths:
|
||||
@ -177,29 +200,41 @@ def combine_bands_to_rgb(
|
||||
# 读取三个波段数据
|
||||
bands_data = []
|
||||
ref_band = None
|
||||
mask = None
|
||||
|
||||
for path in paths:
|
||||
da = open_rasterio(path, masked=True).squeeze(dim="band", drop=True)
|
||||
if ref_band is None:
|
||||
ref_band = da
|
||||
# 使用 minmax 方法 (对应原 stretch_band)
|
||||
|
||||
# 更新掩膜: masked=True 时 nodata 会被转换为 NaN
|
||||
current_mask = np.isfinite(da.values)
|
||||
if mask is None:
|
||||
mask = current_mask
|
||||
else:
|
||||
mask &= current_mask
|
||||
|
||||
# 使用 minmax 方法
|
||||
bands_data.append(normalize_band(da.values, method="minmax"))
|
||||
|
||||
# 合并三个波段为RGB图像
|
||||
# 不再设置 NoData 值，使用 Alpha 通道表示透明度
|
||||
alpha_channel = np.where(mask, 255, 0).astype(np.uint8)
|
||||
bands_data.append(alpha_channel)
|
||||
|
||||
# 合并四个波段为 RGBA 图像
|
||||
rgb_array = xr.DataArray(
|
||||
np.dstack(bands_data),
|
||||
dims=("y", "x", "band"),
|
||||
coords={"band": [1, 2, 3], "y": ref_band.y, "x": ref_band.x},
|
||||
coords={"band": [1, 2, 3, 4], "y": ref_band.y, "x": ref_band.x},
|
||||
)
|
||||
# 转置维度顺序以符合rioxarray要求
|
||||
rgb_array = rgb_array.transpose("band", "y", "x")
|
||||
# 写入元数据
|
||||
rgb_array.rio.write_crs(ref_band.rio.crs, inplace=True)
|
||||
rgb_array.rio.write_transform(ref_band.rio.transform(), inplace=True)
|
||||
rgb_array.rio.write_nodata(0, inplace=True)
|
||||
# 保存为TIFF文件
|
||||
rgb_array.rio.to_raster(output_path, driver="COG", dtype="uint8")
|
||||
print(f"RGB图像已保存到: {output_path}")
|
||||
print(f"RGBA图像已保存到: {output_path}")
|
||||
return
|
||||
|
||||
|
||||
@ -231,8 +266,8 @@ 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 = "Target4"
|
||||
year = 2025
|
||||
for region_name in ["Target1", "Target2", "Target3", "Target4"]:
|
||||
for year in range(2015, 2025 + 1):
|
||||
img_name = f"S1_S2_{region_name}_{year}.tif"
|
||||
img_tif_path = imgs_dir / region_name / img_name
|
||||
bands_lst = [3, 2, 1] # GDAL 波段索引从1开始
|
||||
Loading…
x
Reference in New Issue
Block a user