From 892a2574f6fb4e6de692dfed4f6ae8d90a12837b Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 17 Apr 2026 16:58:04 +0300 Subject: [PATCH] Fix: ensure safetensors works with PNG-only datasets - Force save_npy=True when save_safetensors=True (lossless intermediate) - Auto-enable cleanup_npy to remove intermediates after consolidation - Save segmentation PNG in palette mode "P" (class IDs recoverable) - Consolidation now works correctly from PNG-only previous runs Co-Authored-By: Claude Opus 4.6 (1M context) --- src/augmentor/io_utils.py | 13 ++++++++----- src/main.py | 11 +++++++++++ src/tests/test_io_utils.py | 18 ++++++++++++++++++ 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/src/augmentor/io_utils.py b/src/augmentor/io_utils.py index 5b2a317..e89a704 100644 --- a/src/augmentor/io_utils.py +++ b/src/augmentor/io_utils.py @@ -175,11 +175,14 @@ def save_segmentation( _atomic_save_npy(arr, output_dir / f"{stem}_segm.npy") if save_vis: palette = make_palette(num_classes) - seg_np = arr.squeeze(0).astype(np.int32) - h, w = seg_np.shape - seg_clamped = np.clip(seg_np.flatten(), 0, num_classes - 1) - vis = palette[seg_clamped].reshape(h, w, 3) - Image.fromarray(vis).save(output_dir / f"{stem}_segm.png") + seg_np = arr.squeeze(0).astype(np.uint8) + seg_clamped = np.clip(seg_np, 0, num_classes - 1).astype(np.uint8) + img = Image.fromarray(seg_clamped).convert("P") + # PIL palette: flat list of R, G, B, ... (256 entries × 3 = 768 values). + flat_pal = np.zeros(768, dtype=np.uint8) + flat_pal[: num_classes * 3] = palette.flatten() + img.putpalette(flat_pal.tolist()) + img.save(output_dir / f"{stem}_segm.png") def save_segmentation_async( diff --git a/src/main.py b/src/main.py index fdfb4e5..bccb320 100644 --- a/src/main.py +++ b/src/main.py @@ -380,6 +380,17 @@ def run_pipeline( r.output_dir.mkdir(parents=True, exist_ok=True) seen_dirs.add(d) + # When save_safetensors is on, always save intermediate .npy + # (consolidation reads .npy for lossless float16; PNG is lossy uint8). + if pipeline_conf.save_safetensors and not pipeline_conf.save_npy: + pipeline_conf.save_npy = True + if not pipeline_conf.cleanup_npy: + pipeline_conf.cleanup_npy = True + logger.info( + "📦 save_safetensors=True → forcing save_npy=True + cleanup_npy=True " + "(intermediate .npy will be removed after consolidation)." + ) + # Process each stage sequentially. stage_times: dict[str, float] = {} stage_counts: dict[str, int] = {} diff --git a/src/tests/test_io_utils.py b/src/tests/test_io_utils.py index 253d302..7d4bf41 100644 --- a/src/tests/test_io_utils.py +++ b/src/tests/test_io_utils.py @@ -274,6 +274,24 @@ class TestConsolidateSafetensors: shutdown_io_pool() assert (tmp_path / "img01.safetensors").exists() + def test_from_png_only(self, tmp_path: Path) -> None: + """Consolidation works when only .png exist (no .npy).""" + H, W = 32, 32 + # Save depth/edge/chm as vis-only PNG (no npy). + save_depth(torch.rand(1, H, W), tmp_path, "img04", save_npy=False, save_vis=True) + save_edges(torch.rand(1, H, W), tmp_path, "img04", save_npy=False, save_vis=True) + save_chmv2(torch.rand(1, H, W), tmp_path, "img04", save_npy=False, save_vis=True) + # Save segm as palette PNG. + save_segmentation( + torch.randint(0, 11, (1, H, W), dtype=torch.uint8), + tmp_path, "img04", save_npy=False, save_vis=True, num_classes=11, + ) + ok = consolidate_safetensors(tmp_path, "img04") + assert ok + data = st_load_file(tmp_path / "img04.safetensors") + assert set(data.keys()) == {"depth", "edge", "chm", "segm"} + assert data["segm"].dtype == torch.uint8 + def test_values_preserved(self, tmp_path: Path) -> None: """Verify tensor values survive round-trip.""" depth = torch.rand(1, 16, 16)