235 lines
8.2 KiB
Python
235 lines
8.2 KiB
Python
import json
|
|
import os
|
|
import warnings
|
|
import sys
|
|
import time
|
|
import argparse
|
|
import ray
|
|
import numpy as np
|
|
import Flexbrdf.hytools as ht
|
|
from Flexbrdf.hytools.io.envi import *
|
|
from Flexbrdf.hytools.topo import calc_topo_coeffs
|
|
from Flexbrdf.hytools.brdf import calc_brdf_coeffs
|
|
from Flexbrdf.hytools.glint import set_glint_parameters
|
|
from Flexbrdf.hytools.masks import mask_create
|
|
|
|
warnings.filterwarnings("ignore")
|
|
np.seterr(divide='ignore', invalid='ignore')
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="FlexBRDF - 只接受JSON配置文件输入")
|
|
parser.add_argument('config', help='JSON配置文件路径 (例如: config.json)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# 检查配置文件是否存在
|
|
if not os.path.exists(args.config):
|
|
print(f"错误: 配置文件不存在: {args.config}")
|
|
sys.exit(1)
|
|
|
|
# 检查是否为JSON文件
|
|
if not args.config.endswith('.json'):
|
|
print("错误: 只接受JSON格式的配置文件")
|
|
sys.exit(1)
|
|
|
|
# 加载JSON配置
|
|
try:
|
|
with open(args.config, 'r', encoding='utf-8') as f:
|
|
config_dict = json.load(f)
|
|
except json.JSONDecodeError as e:
|
|
print(f"错误: JSON配置文件格式错误: {e}")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"错误: 无法读取配置文件: {e}")
|
|
sys.exit(1)
|
|
|
|
# 验证必需的配置字段
|
|
required_fields = ['input_files', 'anc_files', 'corrections', 'export']
|
|
missing_fields = [f for f in required_fields if f not in config_dict]
|
|
if missing_fields:
|
|
print(f"错误: 配置文件缺少必需字段: {', '.join(missing_fields)}")
|
|
print("必需字段包括: input_files, anc_files, corrections, export")
|
|
sys.exit(1)
|
|
|
|
# 设置默认值
|
|
config_dict.setdefault('file_type', 'envi')
|
|
config_dict.setdefault('num_cpus', 16)
|
|
config_dict.setdefault('bad_bands', [])
|
|
config_dict.setdefault('resample', False)
|
|
config_dict.setdefault('resampler', {'type': 'cubic', 'out_waves': []})
|
|
|
|
print(f"=" * 60)
|
|
print(f"FlexBRDF 图像校正工具")
|
|
print(f"=" * 60)
|
|
print(f"配置文件: {args.config}")
|
|
print(f"输入文件数量: {len(config_dict['input_files'])}")
|
|
print(f"校正类型: {', '.join(config_dict['corrections'])}")
|
|
print(f"输出目录: {config_dict['export']['output_dir']}")
|
|
print(f"=" * 60)
|
|
print("Starting image correction process...")
|
|
total_start = time.time()
|
|
|
|
images = config_dict["input_files"]
|
|
|
|
if ray.is_initialized():
|
|
ray.shutdown()
|
|
print("Using %s CPUs." % config_dict['num_cpus'])
|
|
ray.init(num_cpus = config_dict['num_cpus'])
|
|
|
|
HyTools = ray.remote(ht.HyTools)
|
|
actors = [HyTools.remote() for image in images]
|
|
|
|
print(f"Loading {len(images)} image(s)...")
|
|
load_start = time.time()
|
|
|
|
if config_dict['file_type'] == 'envi':
|
|
anc_files = config_dict["anc_files"]
|
|
_ = ray.get([a.read_file.remote(image,config_dict['file_type'],
|
|
anc_files[image]) for a,image in zip(actors,images)])
|
|
|
|
elif config_dict['file_type'] == 'neon':
|
|
_ = ray.get([a.read_file.remote(image,config_dict['file_type']) for a,image in zip(actors,images)])
|
|
|
|
_ = ray.get([a.create_bad_bands.remote(config_dict['bad_bands']) for a in actors])
|
|
|
|
load_time = time.time() - load_start
|
|
print(".2f")
|
|
|
|
# Apply corrections
|
|
if config_dict["corrections"]:
|
|
print(f"Applying {len(config_dict['corrections'])} correction(s): {', '.join(config_dict['corrections'])}")
|
|
correction_start = time.time()
|
|
|
|
for correction in config_dict["corrections"]:
|
|
if correction =='topo':
|
|
calc_topo_coeffs(actors,config_dict['topo'])
|
|
elif correction == 'brdf':
|
|
calc_brdf_coeffs(actors,config_dict)
|
|
elif correction == 'glint':
|
|
set_glint_parameters(actors,config_dict)
|
|
|
|
correction_time = time.time() - correction_start
|
|
print(".2f")
|
|
|
|
if config_dict['export']['coeffs'] and len(config_dict["corrections"]) > 0:
|
|
print("Exporting correction coefficients...")
|
|
coeff_start = time.time()
|
|
_ = ray.get([a.do.remote(export_coeffs,config_dict['export']) for a in actors])
|
|
coeff_time = time.time() - coeff_start
|
|
print(".2f")
|
|
|
|
if config_dict['export']['image']:
|
|
print("Exporting corrected images...")
|
|
export_start = time.time()
|
|
_ = ray.get([a.do.remote(apply_corrections,config_dict) for a in actors])
|
|
export_time = time.time() - export_start
|
|
print(".2f")
|
|
|
|
total_time = time.time() - total_start
|
|
print(".2f")
|
|
|
|
ray.shutdown()
|
|
|
|
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"] == True:
|
|
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
|
|
|
|
if __name__== "__main__":
|
|
main()
|