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) <noreply@anthropic.com>

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)
+        seg_np = arr.squeeze(0).astype(np.uint8)
-        h, w = seg_np.shape
+        seg_clamped = np.clip(seg_np, 0, num_classes - 1).astype(np.uint8)
-        seg_clamped = np.clip(seg_np.flatten(), 0, num_classes - 1)
+        img = Image.fromarray(seg_clamped).convert("P")
-        vis = palette[seg_clamped].reshape(h, w, 3)
+        # PIL palette: flat list of R, G, B, ... (256 entries × 3 = 768 values).
-        Image.fromarray(vis).save(output_dir / f"{stem}_segm.png")
+        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(

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] = {}

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)