diff --git a/src/core/modeling/modeling_batch.py b/src/core/modeling/modeling_batch.py index 9d26ded..e129bfd 100644 --- a/src/core/modeling/modeling_batch.py +++ b/src/core/modeling/modeling_batch.py @@ -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}") # ═══════════════════════════════════════════════════════════ diff --git a/src/core/prediction/automl_trainer.py b/src/core/prediction/automl_trainer.py index 4ed55a4..96e3590 100644 --- a/src/core/prediction/automl_trainer.py +++ b/src/core/prediction/automl_trainer.py @@ -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, diff --git a/src/core/steps/glint_removal_step.py b/src/core/steps/glint_removal_step.py index 59f578e..9d8245e 100644 --- a/src/core/steps/glint_removal_step.py +++ b/src/core/steps/glint_removal_step.py @@ -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}") diff --git a/src/core/steps/water_mask_step.py b/src/core/steps/water_mask_step.py index ae59e34..d704c76 100644 --- a/src/core/steps/water_mask_step.py +++ b/src/core/steps/water_mask_step.py @@ -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" diff --git a/src/core/utils/mask_converter.py b/src/core/utils/mask_converter.py index ac1c766..dab8db7 100644 --- a/src/core/utils/mask_converter.py +++ b/src/core/utils/mask_converter.py @@ -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])