256 lines
9.2 KiB
Python
256 lines
9.2 KiB
Python
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() |