import json import os import warnings import sys import time import ray import numpy as np import hytools as ht from hytools.io.envi import * from hytools.topo import calc_topo_coeffs from hytools.brdf import calc_brdf_coeffs from hytools.glint import set_glint_parameters from hytools.masks import mask_create warnings.filterwarnings("ignore") np.seterr(divide='ignore', invalid='ignore') # ---------- 全局函数：供 Ray 远程调用 ---------- def export_coeffs(hy_obj, export_dict): """Export correction coefficients to file.""" for correction in hy_obj.corrections: coeff_file = export_dict['output_dir'] coeff_file += os.path.splitext(os.path.basename(hy_obj.file_name))[0] coeff_file += "_%s_coeffs_%s.json" % (correction, export_dict["suffix"]) with open(coeff_file, 'w') as outfile: if correction == 'topo': corr_dict = hy_obj.topo elif correction == 'glint': continue else: corr_dict = hy_obj.brdf json.dump(corr_dict, outfile) def apply_corrections(hy_obj, config_dict): """Apply correction to image and export to file.""" header_dict = hy_obj.get_header() header_dict['data ignore value'] = hy_obj.no_data header_dict['data type'] = 12 output_name = config_dict['export']['output_dir'] output_name += os.path.splitext(os.path.basename(hy_obj.file_name))[0] output_name += "_%s" % config_dict['export']["suffix"] # Export all wavelengths if len(config_dict['export']['subset_waves']) == 0: if config_dict["resample"]: hy_obj.resampler = config_dict['resampler'] waves = hy_obj.resampler['out_waves'] else: waves = hy_obj.wavelengths header_dict['bands'] = len(waves) header_dict['wavelength'] = waves writer = WriteENVI(output_name, header_dict) iterator = hy_obj.iterate(by='line', corrections=hy_obj.corrections, resample=config_dict['resample']) while not iterator.complete: line = iterator.read_next() writer.write_line(line, iterator.current_line) writer.close() # Export subset of wavelengths else: waves = config_dict['export']['subset_waves'] bands = [hy_obj.wave_to_band(x) for x in waves] waves = [round(hy_obj.wavelengths[x], 2) for x in bands] header_dict['bands'] = len(bands) header_dict['wavelength'] = waves writer = WriteENVI(output_name, header_dict) for b, band_num in enumerate(bands): band = hy_obj.get_band(band_num, corrections=hy_obj.corrections) writer.write_band(band, b) writer.close() # Export masks if (config_dict['export']['masks']) and (len(config_dict["corrections"]) > 0): masks = [] mask_names = [] for correction in config_dict["corrections"]: for mask_type in config_dict[correction]['apply_mask']: mask_names.append(correction + '_' + mask_type[0]) masks.append(mask_create(hy_obj, [mask_type])) header_dict['data type'] = 1 header_dict['bands'] = len(masks) header_dict['band names'] = mask_names header_dict['samples'] = hy_obj.columns header_dict['lines'] = hy_obj.lines header_dict['wavelength'] = [] header_dict['fwhm'] = [] header_dict['wavelength units'] = '' header_dict['data ignore value'] = 255 output_name = config_dict['export']['output_dir'] output_name += os.path.splitext(os.path.basename(hy_obj.file_name))[0] output_name += "_%s_mask" % config_dict['export']["suffix"] writer = WriteENVI(output_name, header_dict) for band_num, mask in enumerate(masks): mask = mask.astype(int) mask[~hy_obj.mask['no_data']] = 255 writer.write_band(mask, band_num) del masks # ---------- 主类 ---------- class HyToolsCorrector: """ 高光谱图像校正处理器 用法: corrector = HyToolsCorrector() corrector.run("path/to/config.json") # 或者单次调用 HyToolsCorrector.process_file("path/to/config.json") """ def __init__(self): """初始化处理器（不自动启动 Ray）""" self.config_dict = None self.ray_initialized = False def load_config(self, config_file: str) -> dict: """加载 JSON 配置文件""" with open(config_file, 'r') as f: self.config_dict = json.load(f) return self.config_dict def _init_ray(self): """初始化 Ray（如果尚未初始化）""" if not ray.is_initialized(): num_cpus = self.config_dict.get('num_cpus', None) print(f"Initializing Ray with {num_cpus} CPUs.") ray.init(num_cpus=num_cpus) self.ray_initialized = True else: print("Ray already initialized, reusing existing instance.") self.ray_initialized = True # 假设外部已初始化 def _shutdown_ray(self): """关闭 Ray（仅在由本类启动时关闭）""" if self.ray_initialized and ray.is_initialized(): ray.shutdown() self.ray_initialized = False def run(self, config_file: str = None): """ 执行完整的校正流程 :param config_file: JSON 配置文件路径，若为 None 则需先调用 load_config() """ if config_file is not None: self.load_config(config_file) if self.config_dict is None: raise ValueError("配置文件未加载，请先调用 load_config() 或传入文件路径。") print("Starting image correction process...") total_start = time.time() images = self.config_dict["input_files"] # 初始化 Ray self._init_ray() # 创建远程 actors HyTools = ray.remote(ht.HyTools) actors = [HyTools.remote() for _ in images] # 加载图像 print(f"Loading {len(images)} image(s)...") load_start = time.time() if self.config_dict['file_type'] == 'envi': anc_files = self.config_dict.get("anc_files", {}) _ = ray.get([a.read_file.remote(img, self.config_dict['file_type'], anc_files.get(img, None)) for a, img in zip(actors, images)]) elif self.config_dict['file_type'] == 'neon': _ = ray.get([a.read_file.remote(img, self.config_dict['file_type']) for a, img in zip(actors, images)]) else: raise ValueError(f"不支持的文件类型: {self.config_dict['file_type']}") _ = ray.get([a.create_bad_bands.remote(self.config_dict['bad_bands']) for a in actors]) load_time = time.time() - load_start print(f"加载完成，耗时 {load_time:.2f} 秒") # 应用校正 if self.config_dict.get("corrections"): corr_list = self.config_dict["corrections"] print(f"Applying {len(corr_list)} correction(s): {', '.join(corr_list)}") correction_start = time.time() for correction in corr_list: if correction == 'topo': calc_topo_coeffs(actors, self.config_dict['topo']) elif correction == 'brdf': calc_brdf_coeffs(actors, self.config_dict) elif correction == 'glint': set_glint_parameters(actors, self.config_dict) correction_time = time.time() - correction_start print(f"校正完成，耗时 {correction_time:.2f} 秒") # 导出系数 export_cfg = self.config_dict.get('export', {}) if export_cfg.get('coeffs') and self.config_dict.get("corrections"): print("Exporting correction coefficients...") coeff_start = time.time() _ = ray.get([a.do.remote(export_coeffs, export_cfg) for a in actors]) print(f"系数导出完成，耗时 {time.time() - coeff_start:.2f} 秒") # 导出图像 if export_cfg.get('image'): print("Exporting corrected images...") export_start = time.time() _ = ray.get([a.do.remote(apply_corrections, self.config_dict) for a in actors]) print(f"图像导出完成，耗时 {time.time() - export_start:.2f} 秒") total_time = time.time() - total_start print(f"全部任务完成，总耗时 {total_time:.2f} 秒") # 关闭 Ray（可选，若希望复用则不关闭） self._shutdown_ray() @staticmethod def process_file(config_file: str): """静态方法：单次处理一个配置文件（自动管理 Ray 生命周期）""" corrector = HyToolsCorrector() corrector.run(config_file) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self._shutdown_ray() # ---------- 向后兼容的 main 入口 ---------- def main(): """保持与原始脚本的兼容性""" if len(sys.argv) > 1: config_path = sys.argv[1] else: config_path = r"E:\code\hytools-master\hytools-master\examples\configs\aviris.json" HyToolsCorrector.process_file(config_path) if __name__ == "__main__": main()