duxingchen/WQ_GUI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 各步骤产物改为原子化落盘，复用前校验 .done 完成标记

Browse Source 
跳过条件从 Path.exists() 改为 is_file_complete()，产物成功落盘后统一补 mark_file_complete()。涉及 step1 水域掩膜、step3 去耀斑、SHP 栅格化缓存三处快路径。

step3 去耀斑新增清理分支：无 .done 标记的旧文件一律视为半成品，先删除再重跑（Kutser/Goodman/Hedley/SUGAR 四种方法一致）。

模型落盘改为 atomic_filepath：modeling_batch 批量建模与 automl_trainer 的 AutoML 结果均先写 .__wip 再原子替换，避免中断产生半个 joblib 文件。
This commit is contained in:
duxin
2026-09-15 10:34:34 +08:00
parent 75c74c644d
commit f97527598d
5 changed files with 70 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/core/modeling/modeling_batch.py
View File

@ -824,7 +824,9 @@ class WaterQualityModelingBatch:
'metadata': metadata
}
joblib.dump(save_data, filepath)
from src.utils.util import atomic_filepath
with atomic_filepath(str(filepath)) as _tmp:
joblib.dump(save_data, _tmp)
print(f"模型已保存: {filepath}")
# ═══════════════════════════════════════════════════════════

9
src/core/prediction/automl_trainer.py
View File

@ -455,11 +455,12 @@ def train_with_automl(
if final_model is None:
continue
# 保存
# 保存（★ 原子写入：先 .__wip，成功后同卷替换 + .done）
import joblib
from src.utils.util import atomic_filepath
fname = f"{tgt}_{preproc}_{model_name}_AUTOML.joblib"
fpath = preproc_dir / fname
joblib.dump({
_save_data = {
"model": final_model,
"target_column_name": tgt,
"preprocess_method": preproc,
@ -475,7 +476,9 @@ def train_with_automl(
"was_subsampled": was_sub,
"split_method": split_method,
},
}, fpath)
}
with atomic_filepath(str(fpath)) as _tmp:
joblib.dump(_save_data, _tmp)
cand = AutoMLResult(
success=True,

57
src/core/steps/glint_removal_step.py
View File

@ -300,6 +300,7 @@ class GlintRemovalStep:
from src.core.utils.mask_converter import (
prepare_water_mask_for_algorithm as _default_prepare,
)
from src.utils.util import is_file_complete, mark_file_complete
# 使用提供的函数或默认函数
if _get_image_geo_info is None:
@ -428,10 +429,19 @@ class GlintRemovalStep:
final_bsq = hardcoded_bsq
final_hdr = hardcoded_hdr
if Path(hardcoded_bsq).exists():
print(f"检测到已存在的去耀斑影像文件，直接使用: {hardcoded_bsq}")
if is_file_complete(hardcoded_bsq):
print(f"检测到完整且已完成的去耀斑影像，直接使用: {hardcoded_bsq}")
notify("skipped", f"去耀斑影像已设置: {hardcoded_bsq}")
return _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
else:
# 清理上次中断遗留的半成品脏数据（无 .done 标记一律视为不完整）
for _stale in (hardcoded_bsq, hardcoded_hdr, hardcoded_bsq + '.hdr'):
if _stale and Path(_stale).exists():
print(f"清理上次中断遗留的半成品数据: {_stale}")
try:
Path(_stale).unlink()
except Exception:
pass
kutser = Kutser(
img_path,
@ -448,6 +458,7 @@ class GlintRemovalStep:
if Path(hardcoded_bsq).exists():
_copy_hdr_info(img_path, hardcoded_bsq)
final = _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
mark_file_complete(final) # ★ 原子化成功标记，防半成品复用
notify("completed", f"去耀斑影像已生成: {final}")
return final
raise RuntimeError(f"Kutser算法未生成输出文件: {hardcoded_bsq}")
@ -466,10 +477,19 @@ class GlintRemovalStep:
final_bsq = hardcoded_bsq
final_hdr = hardcoded_hdr
if Path(hardcoded_bsq).exists():
print(f"检测到已存在的去耀斑影像文件，直接使用: {hardcoded_bsq}")
if is_file_complete(hardcoded_bsq):
print(f"检测到完整且已完成的去耀斑影像，直接使用: {hardcoded_bsq}")
notify("skipped", f"去耀斑影像已设置: {hardcoded_bsq}")
return _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
else:
# 清理上次中断遗留的半成品脏数据（无 .done 标记一律视为不完整）
for _stale in (hardcoded_bsq, hardcoded_hdr, hardcoded_bsq + '.hdr'):
if _stale and Path(_stale).exists():
print(f"清理上次中断遗留的半成品数据: {_stale}")
try:
Path(_stale).unlink()
except Exception:
pass
goodman = Goodman(
img_path,
@ -489,6 +509,7 @@ class GlintRemovalStep:
del corrected_bands
final = _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
mark_file_complete(final) # ★ 原子化成功标记，防半成品复用
notify("completed", f"去耀斑影像已生成: {final}")
return final
@ -504,10 +525,19 @@ class GlintRemovalStep:
final_bsq = hardcoded_bsq
final_hdr = hardcoded_hdr
if Path(hardcoded_bsq).exists():
print(f"检测到已存在的去耀斑影像文件，直接使用: {hardcoded_bsq}")
if is_file_complete(hardcoded_bsq):
print(f"检测到完整且已完成的去耀斑影像，直接使用: {hardcoded_bsq}")
notify("skipped", f"去耀斑影像已设置: {hardcoded_bsq}")
return _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
else:
# 清理上次中断遗留的半成品脏数据（无 .done 标记一律视为不完整）
for _stale in (hardcoded_bsq, hardcoded_hdr, hardcoded_bsq + '.hdr'):
if _stale and Path(_stale).exists():
print(f"清理上次中断遗留的半成品数据: {_stale}")
try:
Path(_stale).unlink()
except Exception:
pass
hedley = Hedley(
img_path,
@ -521,6 +551,7 @@ class GlintRemovalStep:
if Path(hardcoded_bsq).exists():
_copy_hdr_info(img_path, hardcoded_bsq)
final = _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
mark_file_complete(final) # ★ 原子化成功标记，防半成品复用
notify("completed", f"去耀斑影像已生成: {final}")
return final
raise RuntimeError(f"Hedley算法未生成输出文件: {hardcoded_bsq}")
@ -547,10 +578,19 @@ class GlintRemovalStep:
final_bsq = hardcoded_bsq
final_hdr = hardcoded_hdr
if Path(hardcoded_bsq).exists():
print(f"检测到已存在的去耀斑影像文件，直接使用: {hardcoded_bsq}")
if is_file_complete(hardcoded_bsq):
print(f"检测到完整且已完成的去耀斑影像，直接使用: {hardcoded_bsq}")
notify("skipped", f"去耀斑影像已设置: {hardcoded_bsq}")
return _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
else:
# 清理上次中断遗留的半成品脏数据（无 .done 标记一律视为不完整）
for _stale in (hardcoded_bsq, hardcoded_hdr, hardcoded_bsq + '.hdr'):
if _stale and Path(_stale).exists():
print(f"清理上次中断遗留的半成品数据: {_stale}")
try:
Path(_stale).unlink()
except Exception:
pass
if sugar_bounds is None:
sugar_bounds = [(1, 2)]
@ -569,6 +609,7 @@ class GlintRemovalStep:
if Path(hardcoded_bsq).exists():
_copy_hdr_info(img_path, hardcoded_bsq)
final = _safe_rename(hardcoded_bsq, hardcoded_hdr, final_bsq, final_hdr)
mark_file_complete(final) # ★ 原子化成功标记，防半成品复用
notify("completed", f"去耀斑影像已生成: {final}")
return final
raise RuntimeError(f"SUGAR算法未生成输出文件: {hardcoded_bsq}")

11
src/core/steps/water_mask_step.py
View File

@ -46,6 +46,7 @@ class WaterMaskStep:
dat 格式的水域掩膜文件路径
"""
from src.utils.extract_water_area import rasterize_shp, ndwi
from src.utils.util import is_file_complete, mark_file_complete
from src.core.utils.preview_generator import (
generate_image_preview,
generate_water_mask_overlay,
@ -85,12 +86,13 @@ class WaterMaskStep:
ndwi_output_path = output_path or str(water_mask_dir / "water_mask_from_ndwi.dat")
os.makedirs(Path(ndwi_output_path).parent, exist_ok=True)
if Path(ndwi_output_path).exists():
print(f"检测到已存在的NDWI掩膜文件，直接使用: {ndwi_output_path}")
if is_file_complete(ndwi_output_path):
print(f"检测到完整且已完成的NDWI掩膜文件，直接使用: {ndwi_output_path}")
notify("skipped", f"水域掩膜已设置: {ndwi_output_path}")
return ndwi_output_path
ndwi(img_path, ndwi_threshold, ndwi_output_path)
mark_file_complete(ndwi_output_path)
if generate_png:
overlay_path = water_mask_dir / "water_mask_overlay.png"
@ -119,8 +121,8 @@ class WaterMaskStep:
shp_output_path = output_path or str(water_mask_dir / "water_mask_from_shp.dat")
os.makedirs(Path(shp_output_path).parent, exist_ok=True)
if Path(shp_output_path).exists():
print(f"检测到已存在的栅格化掩膜文件，直接使用: {shp_output_path}")
if is_file_complete(shp_output_path):
print(f"检测到完整且已完成的栅格化掩膜文件，直接使用: {shp_output_path}")
notify("skipped", f"水域掩膜已设置: {shp_output_path}")
if generate_png:
overlay_path = water_mask_dir / "water_mask_overlay.png"
@ -130,6 +132,7 @@ class WaterMaskStep:
safe_mask_path = os.path.abspath(mask_path).replace("\\", "/")
rasterize_shp(safe_mask_path, shp_output_path, img_path)
mark_file_complete(shp_output_path)
if generate_png:
overlay_path = water_mask_dir / "water_mask_overlay.png"

8
src/core/utils/mask_converter.py
View File

@ -94,6 +94,7 @@ def _convert_shp_to_mask(shp_path: str, img_path: str,
callback=None) -> np.ndarray:
"""将 shapefile 栅格化为掩膜数组"""
from src.utils.extract_water_area import rasterize_shp
from src.utils.util import is_file_complete, mark_file_complete
safe_shp_path = os.path.abspath(shp_path).replace('\\', '/')
shp_name = Path(safe_shp_path).stem
@ -103,9 +104,9 @@ def _convert_shp_to_mask(shp_path: str, img_path: str,
else:
temp_mask_path = f"/tmp/water_mask_{shp_name}.dat"
# 缓存：已栅格化则直接读取
if Path(temp_mask_path).exists():
print(f"使用已存在的栅格化掩膜: {temp_mask_path}")
# 缓存：已完成（带 .done 标记）才直接读取，否则重新栅格化
if is_file_complete(temp_mask_path):
print(f"使用已完成的栅格化掩膜: {temp_mask_path}")
return _load_raster_mask(temp_mask_path, image_shape[0], image_shape[1])
# 需要栅格化
@ -114,6 +115,7 @@ def _convert_shp_to_mask(shp_path: str, img_path: str,
print(f"正在将 SHP 栅格化: {safe_shp_path}")
rasterize_shp(safe_shp_path, temp_mask_path, img_path)
mark_file_complete(temp_mask_path)
return _load_raster_mask(temp_mask_path, image_shape[0], image_shape[1])