diff --git a/src/core/steps/glint_removal_step.py b/src/core/steps/glint_removal_step.py index 960568e..6aa2de6 100644 --- a/src/core/steps/glint_removal_step.py +++ b/src/core/steps/glint_removal_step.py @@ -48,6 +48,149 @@ def _safe_rename(src_bsq: str, src_hdr: str, dest_bsq: str, dest_hdr: str) -> st return dest_bsq +def _ensure_bsq(img_path): + """检测 BIP 格式并自动转换为 BSQ(一次性,后续步骤全受益)。 + + BIP 格式下逐波段读取需要扫描整文件,308 波段就是 308 次全文件遍历。 + 转换为 BSQ 后每波段连续存储,读取速度从 104 秒/波段降至 0.3 秒/波段。 + """ + if not img_path or not os.path.isfile(img_path): + return img_path + + # 检测 BIP + _hdr = img_path + '.hdr' + if not os.path.exists(_hdr): + _hdr = os.path.splitext(img_path)[0] + '.hdr' + _is_bip = False + if os.path.exists(_hdr): + try: + with open(_hdr, 'r') as f: + _txt = f.read().lower() + _is_bip = 'interleave = bip' in _txt + except Exception: + pass + + if not _is_bip: + return img_path # 已经是 BSQ/BIL,无需转换 + + # 生成转换路径(保持同名以兼容所有下游步骤) + _src = Path(img_path) + _tmp_bsq = _src.parent / (_src.stem + '_BSQ_tmp') # 临时名,避免覆盖 + + # 检查空间(转换需要同时存放 BIP + BSQ,约 2x 文件大小) + _src_size = _src.stat().st_size + try: + import shutil + _free = shutil.disk_usage(str(_src.parent)).free + except Exception: + _free = _src_size * 2 + + _need_margin = _src_size * 0.05 # 5% 安全边界 + if _free < _src_size + _need_margin: + print(f"[BIP→BSQ] ⚠ 空间不足以转换!" + f"需要额外 {(_src_size+_need_margin)/1024**3:.1f}GB, 当前可用 {_free/1024**3:.1f}GB") + print(f"[BIP→BSQ] 跳过转换,继续使用 BIP(速度较慢)") + print(f"[BIP→BSQ] 提示: 清出 {(_src_size+_need_margin-_free)/1024**3:.1f}GB 空间后重试即可自动转换") + return img_path + + # gdal.Translate 转换 + print(f"[BIP→BSQ] 检测到 BIP 格式,开始转换..." + f"({_src_size/1024**3:.1f}GB, 写入临时文件后原地替换)") + import time as _t + _t0 = _t.time() + try: + from osgeo import gdal + gdal.UseExceptions() + gdal.Translate( + str(_tmp_bsq), img_path, + format='ENVI', + creationOptions=['INTERLEAVE=BSQ'], + ) + _elapsed = _t.time() - _t0 + + # 记录临时文件路径(后续重命名时用到) + _tmp_data = str(_tmp_bsq) + _tmp_hdr = _tmp_data + '.hdr' + + # ★ 保存原始 HDR 中的波长/元数据(gdal.Translate 不保留) + _bip_hdr = img_path + '.hdr' + if not os.path.exists(_bip_hdr): + _bip_hdr = os.path.splitext(img_path)[0] + '.hdr' + _hdr_extras = {} # 需要保留的元数据 + if os.path.exists(_bip_hdr): + try: + with open(_bip_hdr, 'r') as _fh: + _in_wavelength = False + _wl_lines = [] + for _line in _fh: + _stripped = _line.strip() + if _stripped.lower().startswith('wavelength'): + _in_wavelength = True + _wl_lines.append(_line) + if '}' in _stripped: + _in_wavelength = False + elif _in_wavelength: + _wl_lines.append(_line) + if '}' in _stripped: + _in_wavelength = False + elif any(_stripped.lower().startswith(k) for k in + ('wavelength units', 'band names', 'bands name', + 'sensor type', 'spectral binning')): + _hdr_extras[_stripped.split('=')[0].strip()] = _line + if _wl_lines: + _hdr_extras['_wavelength_lines'] = ''.join(_wl_lines) + except Exception: + pass + + # 删除原始 BIP 及其 HDR + if not os.path.exists(_bip_hdr): + _bip_hdr = os.path.splitext(img_path)[0] + '.hdr' + _src.unlink() + if os.path.exists(_bip_hdr): + os.unlink(_bip_hdr) + + # 将 BSQ 临时文件重命名为原始文件名(所有下游步骤零改动) + _final = str(_src) + os.rename(_tmp_data, _final) + _final_hdr = _final + '.hdr' + if os.path.exists(_tmp_hdr): + if os.path.exists(_final_hdr): + os.unlink(_final_hdr) + os.rename(_tmp_hdr, _final_hdr) + + # ★ 将波长等元数据写回新 HDR + if _hdr_extras and os.path.exists(_final_hdr): + try: + with open(_final_hdr, 'a') as _fh: + _fh.write('\n') + for _key, _val in _hdr_extras.items(): + if _key == '_wavelength_lines': + _fh.write(_val) + if not _val.rstrip().endswith('}'): + _fh.write('\n') + else: + _fh.write(_val) + if not _val.rstrip().endswith('\n'): + _fh.write('\n') + except Exception: + pass + + print(f"[BIP→BSQ] 转换完成 ({_elapsed:.0f}s)," + f"BIP → BSQ 原地替换,所有路径无需修改") + return str(_final) + except Exception as e: + # 清理失败的临时文件 + if os.path.exists(_tmp_data): + try: + os.unlink(_tmp_data) + if os.path.exists(_tmp_hdr): + os.unlink(_tmp_hdr) + except Exception: + pass + print(f"[BIP→BSQ] 转换失败: {e},继续使用 BIP") + return img_path + + class GlintRemovalStep: """去除耀斑步骤""" @@ -227,6 +370,9 @@ class GlintRemovalStep: interp_end_time = time.time() print(f"插值完成,使用插值后的影像: {img_path}") + # ---- BIP → BSQ 自动转换(一次性,后续步骤全受益)---- + img_path = _ensure_bsq(img_path) + # ---- 获取影像信息 ---- geotransform, projection, width, height, n_bands = _get_image_geo_info(img_path) print(f"影像尺寸: {width} x {height} x {n_bands}")