Files
BRDF/ocbrdf/ocbrdf_main.py
2026-04-10 16:46:45 +08:00

578 lines
21 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import gc
import os
import numpy as np
import sys
import xarray as xr
import rasterio
import spectral
from brdf_model_M02 import M02
from brdf_model_M02SeaDAS import M02SeaDAS
from brdf_model_L11 import L11
from brdf_model_O25 import O25
from brdf_utils import ADF_OCP, squeeze_trivial_dims
UNC_LUT_CACHE = {}
"""
主 BRDF 校正模块
输入为 xarray 数据集
所需的光谱维度为 "bands",其他维度自由
输入数据集中必填的字段:
Rw: 方向性海洋反射率
sza: 太阳天顶角
vza: 观测天顶角
raa: 相对方位角当太阳和观测在同一侧时raa=0
输入数据集中可选的字段:
Rw_unc: Rw 的不确定度(若缺失,设为零)
输出数据集中的字段:
nrrs: 完全归一化的遥感反射率
rho_ex_w: nrrs * PI
omega_b: bb/(a+bb)
eta_b: bbw/bb
C_brdf: BRDF 校正因子
brdf_unc: C_brdf 的不确定度
nrrs_unc: nrrs 的不确定度
使用和适配 brdf_hypercp 模块需保留的信息:
Brdf_hypercp 是 EUMETSAT 研究项目的一部分
"BRDF correction of S3 OLCI water reflectance products"S3 OLCI 水体反射率产品的 BRDF 校正),
合同号RB_EUM-CO-21-4600002626-JIG。
研究团队成员Davide D'Alimonte (davide.dalimonte@aequora.org)
Tamito Kajiyama (tamito.kajiyama@aequora.org)
Jaime Pitarch (jaime.pitarchportero@artov.ismar.cnr.it)
Vittorio Brando (vittorio.brando@cnr.it)
Marco Talone (talone@icm.csic.es) 以及
Constant Mazeran (constant.mazeran@solvo.fr)。
BRDF 查找表中的相对方位角遵循 OLCI 约定。详见 https://www.eumetsat.int/media/50720图 6。
"""
# ---------------------------------------------------------------------------
# File I/O helpers
# ---------------------------------------------------------------------------
def _resolve_envi_header(filepath):
"""Return the ENVI header path accepted by spectral."""
if filepath.lower().endswith(".hdr"):
return filepath
hdr_path = filepath + ".hdr"
if os.path.exists(hdr_path):
return hdr_path
base, _ext = os.path.splitext(filepath)
hdr_path = base + ".hdr"
if os.path.exists(hdr_path):
return hdr_path
return filepath
def read_bsq(filepath, wavelengths=None, dtype=np.float32):
"""Read an ENVI BSQ hyperspectral image.
Parameters
----------
filepath : str
Path to the ENVI header (.hdr) or image file.
wavelengths : array-like, optional
Band centre wavelengths in nm. When *None*, wavelengths are read
from the ENVI header if available.
Returns
-------
data : ndarray, shape (bands, rows, cols), float64
wavelengths : ndarray, shape (bands,)
img : spectral image object (metadata / transform access)
"""
img = spectral.open_image(_resolve_envi_header(filepath))
data = np.asarray(img.open_memmap(interleave="source"), dtype=dtype)
# data = np.transpose(data, (2, 0, 1)) # -> (bands, rows, cols)
if wavelengths is None:
if hasattr(img, 'bands') and img.bands.centers is not None:
wavelengths = np.array(img.bands.centers, dtype=np.float32)
else:
n_bands = data.shape[0]
wavelengths = np.arange(1, n_bands + 1, dtype=np.float32)
return data, np.asarray(wavelengths, dtype=dtype), img
def read_bip_angles(filepath, dtype=np.float32):
"""Read an ENVI BIP angle file and extract the five geometry bands.
Band layout (1-indexed):
band 7 -> SZA Solar Zenith Angle
band 8 -> SAA Solar Azimuth Angle
band 9 -> VZA Sensor Zenith Angle
band 10 -> VAA Sensor Azimuth Angle
band 11 -> RAA Relative Azimuth Angle
Parameters
----------
filepath : str
Path to the ENVI BIP file (header or image file).
Returns
-------
dict with keys 'sza', 'saa', 'vza', 'vaa', 'raa'.
Each value is a 2-D ndarray of shape (rows, cols), float64.
"""
img = spectral.open_image(_resolve_envi_header(filepath))
data = np.asarray(img.open_memmap(interleave="source"), dtype=dtype)
n_bands = data.shape[2]
if n_bands < 11:
raise ValueError(
f"Angle file has only {n_bands} bands; at least 11 are required "
"(bands 711 = SZA, SAA, VZA, VAA, RAA)."
)
return {
'sza': data[:, :, 6].copy(),
'saa': data[:, :, 7].copy(),
'vza': data[:, :, 8].copy(),
'vaa': data[:, :, 9].copy(),
'raa': data[:, :, 10].copy(),
}
def read_water_mask(filepath):
"""Read a single-band GeoTIFF water mask.
Parameters
----------
filepath : str
Path to the GeoTIFF file.
Returns
-------
mask : ndarray, shape (rows, cols), int8
Pixels with value 1 indicate water; all other values are non-water.
"""
with rasterio.open(filepath) as src:
mask = src.read(1).astype(np.int8)
return mask
def save_brdf_result(
ds_out,
output_file,
source_file=None,
output_var="Rw_brdf",
output_format="ENVI",
dtype=np.float32,
):
"""Save the BRDF result to ENVI or NetCDF."""
output_format = output_format.upper()
if output_var not in ds_out:
raise KeyError(f"Output variable '{output_var}' not found in result dataset.")
if output_format == "NETCDF":
if not output_file.lower().endswith(".nc"):
output_file = output_file + ".nc"
ds_out.to_netcdf(output_file)
return output_file
if output_format != "ENVI":
raise ValueError("output_format must be 'ENVI' or 'NETCDF'.")
cube = np.asarray(ds_out[output_var].values, dtype=dtype)
cube = np.transpose(cube, (1, 2, 0))
if output_file.lower().endswith(".hdr"):
hdr_path = output_file
else:
hdr_path = output_file + ".hdr"
metadata = {
"interleave": "bsq",
"description": f"BRDF corrected result: {output_var}",
"bands": cube.shape[2],
"lines": cube.shape[0],
"samples": cube.shape[1],
}
if "bands" in ds_out.coords:
metadata["wavelength"] = [str(v) for v in ds_out.coords["bands"].values.tolist()]
if source_file is not None:
metadata["source_file"] = source_file
spectral.envi.save_image(
hdr_path,
cube,
dtype=dtype,
interleave="bsq",
metadata=metadata,
force=True,
)
return hdr_path
# ---------------------------------------------------------------------------
# Top-level pipeline
# ---------------------------------------------------------------------------
def run_brdf_correction(
hyperspectral_file,
angle_file,
mask_file,
wavelengths=None,
brdf_model='L11',
chunk_size=4096,
):
"""BRDF correction pipeline for water pixels in a hyperspectral scene.
Reads the three input files, identifies water pixels via the mask,
runs ``brdf_prototype`` only on those pixels, and returns a full-scene
xarray Dataset where water pixels carry corrected values and non-water
pixels retain the original Rw (BRDF-derived fields are NaN outside
the water mask).
Parameters
----------
hyperspectral_file : str
Path to the ENVI BSQ hyperspectral image (.hdr sidecar required).
angle_file : str
Path to the ENVI BIP angle file. Bands 711 must be
SZA, SAA, VZA, VAA, RAA respectively.
mask_file : str
Path to the GeoTIFF water mask. Pixels == 1 are treated as water.
wavelengths : array-like, optional
Band centre wavelengths in nm. Overrides the header value when given.
brdf_model : str
BRDF model identifier: ``'L11'`` (default), ``'M02'``,
``'M02SeaDAS'``, or ``'O25'``.
chunk_size : int, optional
Number of water pixels processed per batch (default 4096).
Smaller values reduce peak memory at the cost of more loop
iterations.
Returns
-------
xr.Dataset
Full-scene dataset with dimensions ``(bands, y, x)``.
Variables:
* ``Rw`` original directional water reflectance
* ``nrrs`` normalised remote-sensing reflectance (water only)
* ``rho_ex_w`` nrrs × π (water only)
* ``C_brdf`` BRDF correction factor (water only)
* ``brdf_unc`` uncertainty of C_brdf (water only)
* ``nrrs_unc`` uncertainty of nrrs (water only)
* ``C_brdf_fail`` flag: True where correction failed (water only)
* ``water_mask`` the input mask (1 = water)
* ``sza``, ``saa``, ``vza``, ``vaa``, ``raa`` viewing geometry
"""
# ------------------------------------------------------------------
# 1. Read all inputs
# ------------------------------------------------------------------
rw_data, wvl, _img_meta = read_bsq(hyperspectral_file, wavelengths, dtype=np.float32)
angles = read_bip_angles(angle_file, dtype=np.float32)
water_mask = read_water_mask(mask_file)
n_bands, n_rows, n_cols = rw_data.shape
# Sanity-check spatial dimensions
if angles['sza'].shape != (n_rows, n_cols):
raise ValueError(
f"Angle file spatial shape {angles['sza'].shape} does not match "
f"hyperspectral image shape ({n_rows}, {n_cols})."
)
if water_mask.shape != (n_rows, n_cols):
raise ValueError(
f"Mask shape {water_mask.shape} does not match "
f"hyperspectral image shape ({n_rows}, {n_cols})."
)
# ------------------------------------------------------------------
# 2. Identify water pixels
# ------------------------------------------------------------------
water_rows, water_cols = np.where(water_mask == 1)
n_water = water_rows.size
if n_water == 0:
raise RuntimeError("No water pixels found in the mask (no pixels == 1).")
# ------------------------------------------------------------------
# 3. Pre-allocate full-scene output arrays
# Water pixels will be filled in chunks; non-water pixels stay NaN.
# ------------------------------------------------------------------
shape3 = (n_bands, n_rows, n_cols)
shape2 = (n_rows, n_cols)
nrrs_full = np.full(shape3, np.nan, dtype=np.float32)
rho_ex_w_full = np.full(shape3, np.nan, dtype=np.float32)
C_brdf_full = np.full(shape3, np.nan, dtype=np.float32)
brdf_unc_full = np.full(shape3, np.nan, dtype=np.float32)
nrrs_unc_full = np.full(shape3, np.nan, dtype=np.float32)
C_brdf_fail_full = np.zeros(shape2, dtype=bool)
rw_brdf_full = rw_data.copy()
# ------------------------------------------------------------------
# 46. Chunked BRDF correction
# Process water pixels in batches of `chunk_size` to bound
# peak memory. Each batch builds a minimal xarray Dataset,
# calls brdf_prototype, writes results back to the full-scene
# arrays, then releases the batch objects.
# ------------------------------------------------------------------
n_chunks = int(np.ceil(n_water / chunk_size))
for chunk_idx in range(n_chunks):
lo = chunk_idx * chunk_size
hi = min(lo + chunk_size, n_water)
rows_c = water_rows[lo:hi]
cols_c = water_cols[lo:hi]
# Extract batch vectors
rw_c = rw_data[:, rows_c, cols_c].T # (chunk, bands)
sza_c = angles['sza'][rows_c, cols_c] # (chunk,)
vza_c = angles['vza'][rows_c, cols_c]
raa_c = angles['raa'][rows_c, cols_c]
pixel_idx_c = np.arange(hi - lo)
ds_chunk = xr.Dataset(
{
'Rw': xr.DataArray(rw_c, dims=('n', 'bands'),
coords={'n': pixel_idx_c, 'bands': wvl}),
'sza': xr.DataArray(sza_c, dims='n',
coords={'n': pixel_idx_c}),
'vza': xr.DataArray(vza_c, dims='n',
coords={'n': pixel_idx_c}),
'raa': xr.DataArray(raa_c, dims='n',
coords={'n': pixel_idx_c}),
}
)
ds_corr = brdf_prototype(ds_chunk, brdf_model=brdf_model)
# Write results back; ds_corr arrays have dims (n, bands) → transpose
nrrs_full[:, rows_c, cols_c] = ds_corr['nrrs'].values.T
rho_ex_w_full[:, rows_c, cols_c] = ds_corr['rho_ex_w'].values.T
C_brdf_full[:, rows_c, cols_c] = ds_corr['C_brdf'].values.T
brdf_unc_full[:, rows_c, cols_c] = ds_corr['brdf_unc'].values.T
nrrs_unc_full[:, rows_c, cols_c] = ds_corr['nrrs_unc'].values.T
C_brdf_fail_full[rows_c, cols_c] = ds_corr['C_brdf_fail'].values
rw_brdf_full[:, rows_c, cols_c] = ds_corr['rho_ex_w'].values.T
# Release batch objects so GC can reclaim memory
del rw_c, sza_c, vza_c, raa_c, pixel_idx_c, ds_chunk, ds_corr
gc.collect()
# ------------------------------------------------------------------
# 7. Assemble full-scene output Dataset
# ------------------------------------------------------------------
full_coords = {'bands': wvl, 'y': np.arange(n_rows), 'x': np.arange(n_cols)}
dims3 = ('bands', 'y', 'x')
dims2 = ('y', 'x')
ds_out = xr.Dataset(
{
'Rw': xr.DataArray(rw_data, dims=dims3, coords=full_coords),
'Rw_brdf': xr.DataArray(rw_brdf_full, dims=dims3, coords=full_coords),
'nrrs': xr.DataArray(nrrs_full, dims=dims3, coords=full_coords),
'rho_ex_w': xr.DataArray(rho_ex_w_full, dims=dims3, coords=full_coords),
'C_brdf': xr.DataArray(C_brdf_full, dims=dims3, coords=full_coords),
'brdf_unc': xr.DataArray(brdf_unc_full, dims=dims3, coords=full_coords),
'nrrs_unc': xr.DataArray(nrrs_unc_full, dims=dims3, coords=full_coords),
'C_brdf_fail': xr.DataArray(C_brdf_fail_full, dims=dims2),
'water_mask': xr.DataArray(water_mask, dims=dims2),
'sza': xr.DataArray(angles['sza'], dims=dims2),
'saa': xr.DataArray(angles['saa'], dims=dims2),
'vza': xr.DataArray(angles['vza'], dims=dims2),
'vaa': xr.DataArray(angles['vaa'], dims=dims2),
'raa': xr.DataArray(angles['raa'], dims=dims2),
}
)
return ds_out
def process_brdf_files(
hyperspectral_file,
angle_file,
mask_file,
output_file,
wavelengths=None,
brdf_model="L11",
output_var="Rw_brdf",
output_format="ENVI",
chunk_size=4096,
):
"""Run BRDF correction from three input files and save the result."""
ds_out = run_brdf_correction(
hyperspectral_file=hyperspectral_file,
angle_file=angle_file,
mask_file=mask_file,
wavelengths=wavelengths,
brdf_model=brdf_model,
chunk_size=chunk_size,
)
saved_path = save_brdf_result(
ds_out=ds_out,
output_file=output_file,
source_file=hyperspectral_file,
output_var=output_var,
output_format=output_format,
)
return ds_out, saved_path
# ---------------------------------------------------------------------------
# Core BRDF correction (unchanged from original)
# ---------------------------------------------------------------------------
def brdf_prototype(ds, adf=None, brdf_model='L11'):
# 测试 BRDF 模型在 GUI 中不支持:强制覆盖
# brdf_model = 'M02SeaDAS'
# 压缩单一维度(例如,投放次数、提取次数等)以避免插值问题
ds, squeezedDims = squeeze_trivial_dims(ds)
# Initialise model
if brdf_model == 'M02':
BRDF_model = M02(bands=ds.bands, aot=ds.aot, wind=ds.wind, adf=None) # Don't use brdf_py.ADF context
elif brdf_model == 'M02SeaDAS':
BRDF_model = M02SeaDAS(bands=ds.bands, adf=None) # Don't use brdf_py.ADF context
elif brdf_model == 'L11':
BRDF_model = L11(bands=ds.bands, adf=None) # Don't use brdf_py.ADF context
elif brdf_model == 'O25':
BRDF_model = O25(bands=ds.bands, adf=None) # Don't use brdf_py.ADF context
else:
print("BRDF model %s not supported" % brdf_model)
sys.exit(1)
# Init pixel
BRDF_model.init_pixels(ds['sza'], ds['vza'], ds['raa'])
# Compute IOP and normalize by iterating
ds['nrrs'] = ds['Rw'] / np.pi
ds['convergeFlag'] = (0 * ds['sza']).astype(bool)
ds['C_brdf'] = 0 * ds['nrrs'] + 1
for iter_brdf in range(int(BRDF_model.niter)):
# M02: Initialise chl_iter
if brdf_model in ['M02', 'M02SeaDAS'] and (iter_brdf == 0):
chl_iter = {}
ds['log10_chl'] = 0 * ds['sza'] + float(np.log10(BRDF_model.OC4MEchl0))
chl_iter[-1] = 0 * ds['sza'] + float(BRDF_model.OC4MEchl0)
ds = BRDF_model.backward(ds, iter_brdf)
# M02: Check convergence (dummy for M02SeaDAS for the moment... epsilon set to 0)
if brdf_model in ['M02', 'M02SeaDAS']:
chl_iter[iter_brdf] = 10 ** ds['log10_chl']
# Check if convergence is reached |chl_old-chl_new| < epsilon * chl_new
ds['convergeFlag'] = (ds['convergeFlag']) | (
(np.abs(chl_iter[iter_brdf - 1] - chl_iter[iter_brdf]) < float(BRDF_model.OC4MEepsilon) * chl_iter[
iter_brdf]))
# Apply forward model in both geometries
# forward_mod = BRDF_model.forward(ds).transpose('n', 'bands')
# forward_mod0 = BRDF_model.forward(ds, normalized=True).transpose('n', 'bands')
forward_mod = BRDF_model.forward(ds)
forward_mod0 = BRDF_model.forward(ds, normalized=True)
ratio = forward_mod0 / forward_mod
# Drop remnant coordinates to avoid ambiguities in the update of the BRDF factor.
for coord in ratio.coords:
if coord not in ds['C_brdf'].coords:
ratio = ratio.drop(coord)
# Normalize reflectance
ds['C_brdf'] = xr.where(ds['convergeFlag'], ds['C_brdf'], ratio)
ds['nrrs'] = ds['Rw'] / np.pi * ds['C_brdf']
# Flag BRDF where NaN and set to 1 (no correction applied).
ds['C_brdf_fail'] = np.isnan(ds['C_brdf'])
ds['C_brdf'] = xr.where(ds['C_brdf_fail'], 1, ds['C_brdf'])
ds['nrrs'] = xr.where(ds['C_brdf_fail'], ds['Rw'] / np.pi, ds['nrrs'])
# If QAA_fail is raised, raise C_brdf_fail (but still apply C_brdf).
if 'QAA_fail' in ds:
ds['C_brdf_fail'] = (ds['C_brdf_fail']) | (ds['QAA_fail'])
# Compute uncertainty
ds = brdf_uncertainty(ds)
# Compute flag
ds['flags_level2'] = ds['Rw'] * 0 # TODO
# Convert to reflectance unit
ds['rho_ex_w'] = ds['nrrs'] * np.pi
# Expand squeezed trivial dimensions
for dim,d0 in squeezedDims.items():
ds = ds.expand_dims(dim,axis=d0)
return ds
def brdf_uncertainty(ds, adf=None):
''' Compute uncertainty of BRDF factor and propagate to nrrs '''
# Read LUT
if adf is None:
adf = ADF_OCP
unc_lut_path = adf % 'UNC'
if unc_lut_path not in UNC_LUT_CACHE:
UNC_LUT_CACHE[unc_lut_path] = xr.open_dataset(unc_lut_path, engine='netcdf4')
LUT = UNC_LUT_CACHE[unc_lut_path]
# Interpolate relative uncertainty
unc = LUT['unc'].interp(lambda_unc=ds.bands, theta_s_unc=ds.sza, theta_v_unc=ds.vza,
delta_phi_unc=ds.raa)
# Compute absolute uncertainty of factor
ds['brdf_unc'] = unc * ds['C_brdf']
# Flag BRDF_unc where NaN and set to 0
ds['brdf_unc_fail'] = np.isnan(ds['brdf_unc'])
ds['brdf_unc'] = xr.where(ds['brdf_unc_fail'], 0, ds['brdf_unc'])
# Propagate to nrrs
Rwex_unc2 = ds['brdf_unc'] * ds['brdf_unc'] * ds['Rw'] * ds['Rw']
if 'Rw_unc' in ds:
Rwex_unc2 += ds['C_brdf'] * ds['C_brdf'] * ds['Rw_unc'] * ds['Rw_unc']
ds['nrrs_unc'] = np.sqrt(Rwex_unc2)/np.pi
return ds
def main():
import argparse
parser = argparse.ArgumentParser(description="Water-region BRDF correction for hyperspectral data.")
parser.add_argument("hyperspectral_file", help="Input ENVI BSQ hyperspectral file or .hdr path.")
parser.add_argument("angle_file", help="Input ENVI BIP angle file or .hdr path.")
parser.add_argument("mask_file", help="Input water mask GeoTIFF; pixels == 1 are corrected.")
parser.add_argument("output_file", help="Output path prefix. ENVI writes .hdr/.img, NetCDF writes .nc.")
parser.add_argument("--brdf-model", default="L11", choices=["L11", "M02", "M02SeaDAS", "O25"])
parser.add_argument("--output-var", default="Rw_brdf", choices=["Rw_brdf", "rho_ex_w", "nrrs", "C_brdf"])
parser.add_argument("--output-format", default="ENVI", choices=["ENVI", "NETCDF"])
parser.add_argument("--chunk-size", type=int, default=4096,
help="Number of water pixels per processing batch (default: 4096).")
args = parser.parse_args()
_ds_out, saved_path = process_brdf_files(
hyperspectral_file=args.hyperspectral_file,
angle_file=args.angle_file,
mask_file=args.mask_file,
output_file=args.output_file,
brdf_model=args.brdf_model,
output_var=args.output_var,
output_format=args.output_format,
chunk_size=args.chunk_size,
)
print(f"BRDF correction finished. Saved to: {saved_path}")
if __name__ == "__main__":
main()