import gc import os import numpy as np import sys import xarray as xr import rasterio from rasterio.windows import Window 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 # 尝试导入 tqdm 进度条库，如果不可用则设为 None try: from tqdm import tqdm except ImportError: tqdm = None 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 7–11 = 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 _block_windows(src, block_size=512): """Yield raster windows for block-wise processing.""" rows, cols = src.height, src.width for row in range(0, rows, block_size): row_end = min(row + block_size, rows) for col in range(0, cols, block_size): col_end = min(col + block_size, cols) window = Window(col, row, col_end - col, row_end - row) yield window def _resolve_output_envi_path(output_file): """Return the ENVI image path used for rasterio writes.""" if output_file.lower().endswith(".img"): return output_file if output_file.lower().endswith(".hdr"): return output_file[:-4] + ".img" return output_file + ".img" def _read_wavelengths_from_header(hyperspectral_file): """Read wavelength metadata from ENVI header if available.""" img = spectral.open_image(_resolve_envi_header(hyperspectral_file)) if hasattr(img, 'bands') and img.bands.centers is not None: wavelengths = img.bands.centers else: wavelengths = img.metadata.get('wavelength') if wavelengths is None: return None return np.asarray(wavelengths, dtype=np.float32) 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 def _get_output_filename(base_output_file: str, var_name: str) -> str: """Generate output filename with variable suffix. Examples: input.hdr + Rw_brdf -> input_Rw_brdf.hdr input.hdr + brdf_unc -> input_brdf_unc.hdr """ base, ext = os.path.splitext(base_output_file) if ext.lower() in ('.hdr', '.img'): base = base # keep base without extension suffix = f"_{var_name}" if var_name.lower() != "rw_brdf" else "" return f"{base}{suffix}{ext}" def save_multiple_brdf_results( ds_out, output_file, source_file=None, output_vars=None, output_format="ENVI", dtype=np.float32, ): """Save multiple output variables to separate files.""" if output_vars is None or len(output_vars) == 0: output_vars = ["Rw_brdf"] saved_files = [] for var in output_vars: out_file = _get_output_filename(output_file, var) try: saved_path = save_brdf_result( ds_out=ds_out, output_file=out_file, source_file=source_file, output_var=var, output_format=output_format, dtype=dtype, ) saved_files.append((var, saved_path)) print(f" Saved {var} to: {saved_path}") except KeyError as e: print(f" Warning: {e}. Skipping {var}.") return saved_files def run_brdf_correction_block( hyperspectral_file, angle_file, mask_file, output_file, wavelengths=None, brdf_model='L11', output_vars=None, # 支持列表 block_size=512, compute_uncertainty=False, progress_bar=True, ): """ Run BRDF correction block-wise and write multiple variables to separate ENVI files. """ if output_vars is None: output_vars = ["Rw_brdf"] if isinstance(output_vars, str): output_vars = [output_vars] # 1. Open all inputs src_rw = rasterio.open(hyperspectral_file) n_bands = src_rw.count n_rows = src_rw.height n_cols = src_rw.width src_ang = rasterio.open(angle_file) if src_ang.height != n_rows or src_ang.width != n_cols: raise ValueError("Angle file dimensions do not match hyperspectral file.") src_mask = rasterio.open(mask_file) if src_mask.height != n_rows or src_mask.width != n_cols: raise ValueError("Mask file dimensions do not match hyperspectral file.") # 2. Wavelengths if wavelengths is None: wavelengths = _read_wavelengths_from_header(hyperspectral_file) if wavelengths is None: wavelengths = np.arange(1, n_bands + 1, dtype=np.float32) bands_xr = xr.DataArray(wavelengths, dims=('bands',), coords={'bands': wavelengths}, name='bands') # 3. Prepare multiple output files (one per variable) out_meta = src_rw.meta.copy() out_meta.update({ 'driver': 'ENVI', 'dtype': 'float32', 'count': n_bands, 'interleave': 'bsq', }) # Open one output file for each variable output_dsts = {} output_paths = {} for var_name in output_vars: var_file = _get_output_filename(output_file, var_name) out_path = _resolve_output_envi_path(var_file) output_dsts[var_name] = rasterio.open(out_path, 'w', **out_meta) output_paths[var_name] = out_path print(f" Writing {var_name} to {out_path}") # 4. Band indices in angle file (1-based) band_sza = 7 band_vza = 9 band_raa = 11 # 5. Map output_var to ds_corr field (support uncertainty variables) var_map = { 'Rw_brdf': 'rho_ex_w', 'rho_ex_w': 'rho_ex_w', 'nrrs': 'nrrs', 'C_brdf': 'C_brdf', 'brdf_unc': 'brdf_unc', 'nrrs_unc': 'nrrs_unc', } # Validate all requested output variables for v in output_vars: if v not in var_map and v not in ['Rw_brdf']: raise ValueError( f"Unknown output_var '{v}'. Allowed: {list(var_map.keys())}" ) # 6. Progress bar setup iterator = _block_windows(src_rw, block_size=block_size) if progress_bar and tqdm is not None: total_blocks = ((n_rows + block_size - 1) // block_size) * ((n_cols + block_size - 1) // block_size) iterator = tqdm(iterator, total=total_blocks, desc="Processing blocks", unit="block") elif progress_bar: print("Warning: tqdm not installed, progress bar disabled.") # Statistics for performance monitoring blocks_processed = 0 blocks_skipped = 0 total_water_pixels = 0 # 7. Block-wise processing with early filtering for window in iterator: # Step 1: Read mask FIRST to enable early exit for pure land blocks mask_block = src_mask.read(1, window=window).astype(np.int8) water_mask = (mask_block == 1) n_water = np.count_nonzero(water_mask) if n_water == 0: # Fast path: pure land/no-water block - minimal processing blocks_skipped += 1 # Still need to read hyperspectral for Rw_brdf output, but skip angles rw_block = src_rw.read(window=window).astype(np.float32) output_blocks = {} for var_name in output_vars: if var_name == 'Rw_brdf': output_blocks[var_name] = rw_block.copy() else: output_blocks[var_name] = np.full_like(rw_block, np.nan, dtype=np.float32) for var_name, block_data in output_blocks.items(): output_dsts[var_name].write(block_data, window=window) del rw_block, mask_block, water_mask gc.collect() continue # Step 2: Has water pixels - full processing blocks_processed += 1 total_water_pixels += n_water # Read remaining data only for blocks that need processing rw_block = src_rw.read(window=window).astype(np.float32) sza_block = src_ang.read(band_sza, window=window).astype(np.float32) vza_block = src_ang.read(band_vza, window=window).astype(np.float32) raa_block = src_ang.read(band_raa, window=window).astype(np.float32) water_idx = np.where(water_mask) # Prepare output blocks for all variables output_blocks = {} for var_name in output_vars: if var_name == 'Rw_brdf': output_blocks[var_name] = rw_block.copy() else: output_blocks[var_name] = np.full_like(rw_block, np.nan, dtype=np.float32) # Process water pixels rw_water = rw_block[:, water_idx[0], water_idx[1]].T # (n_water, bands) sza_water = sza_block[water_idx] vza_water = vza_block[water_idx] raa_water = raa_block[water_idx] pixel_ids = np.arange(n_water) ds_block = xr.Dataset( { 'Rw': xr.DataArray(rw_water, dims=('n', 'bands'), coords={'n': pixel_ids, 'bands': bands_xr}), 'sza': xr.DataArray(sza_water, dims='n', coords={'n': pixel_ids}), 'vza': xr.DataArray(vza_water, dims='n', coords={'n': pixel_ids}), 'raa': xr.DataArray(raa_water, dims='n', coords={'n': pixel_ids}), } ) ds_corr = brdf_prototype( ds_block, brdf_model=brdf_model, compute_uncertainty=compute_uncertainty, ) # Write each requested variable to the water pixels for var_name in output_vars: field_name = var_map.get(var_name, 'rho_ex_w') if var_name == 'Rw_brdf': field_name = 'rho_ex_w' if field_name in ds_corr: values = ds_corr[field_name].values.T output_blocks[var_name][:, water_idx[0], water_idx[1]] = values # Write all variables for this block for var_name, block_data in output_blocks.items(): output_dsts[var_name].write(block_data, window=window) del (rw_block, sza_block, vza_block, raa_block, mask_block, output_blocks, water_mask, water_idx, rw_water, sza_water, vza_water, raa_water, pixel_ids, ds_block, ds_corr) gc.collect() # Close all files src_rw.close() src_ang.close() src_mask.close() for dst in output_dsts.values(): dst.close() # Print performance statistics total_blocks = blocks_processed + blocks_skipped if total_blocks > 0: skip_ratio = blocks_skipped / total_blocks * 100 print(f"Block processing completed:") print(f" - Total blocks: {total_blocks}") print(f" - Processed blocks (with water): {blocks_processed}") print(f" - Skipped blocks (no water): {blocks_skipped} ({skip_ratio:.1f}%)") print(f" - Total water pixels processed: {total_water_pixels:,}") return list(output_paths.values()) # --------------------------------------------------------------------------- # Top-level pipeline # --------------------------------------------------------------------------- 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 brdf_prototype(ds, adf=None, brdf_model='L11', compute_uncertainty=True): # 测试 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) 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 only if requested if 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 run_brdf_correction( hyperspectral_file, angle_file, mask_file, wavelengths=None, brdf_model='L11', chunk_size=4096, output_var=None, # if None, keep full output (original behaviour) compute_uncertainty=True, # compute uncertainty only if requested progress_bar=True, # 是否显示进度条 ): """BRDF correction pipeline for water pixels in a hyperspectral scene. ... (docstring unchanged, update description) ... """ # ------------------------------------------------------------------ # 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. Determine which output variable(s) we need to keep # ------------------------------------------------------------------ # If output_var is None, we keep everything (original behaviour) keep_all = (output_var is None) # Map output_var to the field in ds_corr var_map = { 'Rw_brdf': 'rho_ex_w', 'rho_ex_w': 'rho_ex_w', 'nrrs': 'nrrs', 'C_brdf': 'C_brdf' } if not keep_all: if output_var not in var_map: raise ValueError(f"Unknown output_var '{output_var}'. Allowed: {list(var_map.keys())}") field_name = var_map[output_var] else: field_name = None # ------------------------------------------------------------------ # 4. Allocate output arrays (only those needed) # ------------------------------------------------------------------ shape3 = (n_bands, n_rows, n_cols) shape2 = (n_rows, n_cols) if keep_all: # Original behaviour: allocate all arrays 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() else: # Only allocate the requested output array output_full = np.full(shape3, np.nan, dtype=np.float32) # We may also need C_brdf_fail if requested? Not for output_var, but could be used internally? Skip. # For simplicity, we do not allocate extra arrays. # ------------------------------------------------------------------ # 5–6. Chunked BRDF correction # ------------------------------------------------------------------ n_chunks = int(np.ceil(n_water / chunk_size)) # 初始化进度条 iterator = range(n_chunks) if progress_bar and tqdm is not None: iterator = tqdm(iterator, desc="Processing chunks", unit="chunk") elif progress_bar: print("Warning: tqdm not installed, progress bar disabled.") for chunk_idx in iterator: 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}), } ) # Run BRDF model with uncertainty control ds_corr = brdf_prototype(ds_chunk, brdf_model=brdf_model, compute_uncertainty=compute_uncertainty) if keep_all: # Write all results back 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 if compute_uncertainty: 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 else: # Only write the requested output field values = ds_corr[field_name].values.T output_full[:, rows_c, cols_c] = values # 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') if keep_all: # Original behaviour: return all variables 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), '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), } ) if compute_uncertainty: ds_out['brdf_unc'] = xr.DataArray(brdf_unc_full, dims=dims3, coords=full_coords) ds_out['nrrs_unc'] = xr.DataArray(nrrs_unc_full, dims=dims3, coords=full_coords) else: # Return only the requested output variable + geometry + mask ds_out = xr.Dataset( { output_var: xr.DataArray(output_full, dims=dims3, coords=full_coords), 'Rw': xr.DataArray(rw_data, dims=dims3, coords=full_coords), '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), } ) # If uncertainties were computed, we could add them optionally, but user didn't request them. # For memory efficiency, we skip them here. return ds_out def process_brdf_files( hyperspectral_file, angle_file, mask_file, output_file, wavelengths=None, brdf_model="L11", output_vars=None, output_format="ENVI", chunk_size=4096, block_size=512, compute_uncertainty=False, progress_bar=True, ): """Run BRDF correction from three input files and save multiple variables to separate files.""" if output_format.upper() != "ENVI": raise NotImplementedError("Only ENVI output format is supported in block-wise version.") if output_vars is None: output_vars = ["Rw_brdf"] if isinstance(output_vars, str): output_vars = [output_vars] print(f"Computing BRDF correction with variables: {output_vars}") print(f"Output base name: {output_file}") saved_paths = run_brdf_correction_block( hyperspectral_file=hyperspectral_file, angle_file=angle_file, mask_file=mask_file, output_file=output_file, wavelengths=wavelengths, brdf_model=brdf_model, output_vars=output_vars, block_size=block_size, compute_uncertainty=compute_uncertainty, progress_bar=progress_bar, ) return None, saved_paths 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", nargs='+', default=["Rw_brdf"], help="Output variables to save (one or more). " "Supported: Rw_brdf, rho_ex_w, nrrs, C_brdf, brdf_unc, nrrs_unc. " "Example: --output-var Rw_brdf nrrs brdf_unc") parser.add_argument("--output-format", default="ENVI", choices=["ENVI"]) parser.add_argument("--chunk-size", type=int, default=4096, help="Number of water pixels per processing batch (default: 4096).") parser.add_argument("--block-size", type=int, default=512, help="Spatial block size for block-wise processing (default: 512).") parser.add_argument("--compute-uncertainty", action="store_true", help="Compute BRDF uncertainty (may increase runtime and output size).") parser.add_argument("--no-progress-bar", action="store_true", help="Disable progress bar (useful for non-interactive environments).") args = parser.parse_args() progress_bar = not args.no_progress_bar _ds_out, saved_paths = 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_vars=args.output_var, output_format=args.output_format, chunk_size=args.chunk_size, block_size=args.block_size, compute_uncertainty=args.compute_uncertainty, progress_bar=progress_bar, ) if isinstance(saved_paths, list): print(f"BRDF correction finished. Saved {len(saved_paths)} files:") for path in saved_paths: print(f" - {path}") else: print(f"BRDF correction finished. Saved to: {saved_paths}") if __name__ == "__main__": main()