From 13ff07989118a3f7c6478d4f839ddc88b8ff9d82 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 17 Apr 2026 17:11:01 +0300 Subject: [PATCH] Refactor output to directory-based layout + migration script Replace prefix-based naming (crop_12_4_depth.png) with directory-based layout where modality is determined by folder (depth/crop_12_4.png). New structure: output_root/{modality}/{rel_parent}/{stem}.png (vis) output_root/npy/{modality}/{rel_parent}/{stem}.npy (intermediate) output_root/safetensors/{rel_parent}/{stem}.safetensors (training) - Rewrite io_utils.py save functions: (output_root, rel_parent, stem) - Update ImageRecord: output_root + rel_parent instead of output_dir - Add path helpers: npy_path(), vis_path(), safetensors_path() - Add scripts/migrate_layout.py for converting existing datasets - Update all tests (143 passing) - Update README with new layout docs Co-Authored-By: Claude Opus 4.6 (1M context) --- README.md | 49 ++++-- scripts/migrate_layout.py | 153 +++++++++++++++++ src/augmentor/dataset.py | 46 +++--- src/augmentor/io_utils.py | 217 ++++++++++++++----------- src/main.py | 65 +++----- src/tests/test_dataset.py | 17 +- src/tests/test_io_utils.py | 212 ++++++++++++------------ src/tests/test_pipeline_integration.py | 65 +++----- 8 files changed, 502 insertions(+), 322 deletions(-) create mode 100644 scripts/migrate_layout.py diff --git a/README.md b/README.md index efb462e..b5bcfa9 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,9 @@ python -m pytest src/tests/ -v │ │ └── dataset.py # Discovery, filtering, PyTorch Dataset │ ├── conf/ # Gin-configurable dataclasses │ ├── utils/ # Profiler, benchmark, GPU utils -│ └── tests/ # 141 тест (pytest) +│ └── tests/ # 143 теста (pytest) +├── scripts/ +│ └── migrate_layout.py # Миграция со старого prefix-формата └── docs/ ├── segmentation_class_analysis.md # Анализ классов сегментации (11 классов) ├── segearth_ov3_architecture.md # Архитектура SegEarth-OV3 + SAM 3.1 @@ -145,22 +147,24 @@ free_vram = total - reserved batch = round_down_pow2(free_vram / act_per_sample * 0.7) ``` -**Resume** проверяет существование `{stem}_{suffix}.png` (или `.npy`) для каждого изображения и `{stem}.safetensors` для этапа консолидации. Пайплайн можно прервать Ctrl+C и перезапустить -- готовые пропускаются. +**Resume** проверяет существование файлов в соответствующих директориях модальностей. Пайплайн можно прервать Ctrl+C и перезапустить -- готовые пропускаются. ## Формат выхода -Структура директорий **зеркалит** исходный датасет. Исходные изображения не копируются: +Модальность определяется **папкой**, а не суффиксом файла: ``` World-UAV-aug/ -├── Rot/SouthernSuburbs/DB/img/ -│ ├── crop_12_4.safetensors # ВСЕ модальности (для обучения, zero-copy mmap) -│ ├── crop_12_4_depth.png # grayscale визуализация -│ ├── crop_12_4_edge.png # grayscale визуализация -│ ├── crop_12_4_segm.png # RGB palette визуализация (11 классов) -│ └── crop_12_4_chm.png # grayscale визуализация -├── Country/... -└── Terrain/... +├── depth/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis +├── edge/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis +├── segm/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis (palette mode P) +├── chm/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis +├── npy/depth/Rot/SouthernSuburbs/DB/img/crop_12_4.npy # float16 intermediate +├── npy/edge/... +├── npy/segm/... +├── npy/chm/... +├── safetensors/Rot/SouthernSuburbs/DB/img/crop_12_4.safetensors # для обучения +└── manifest.json ``` ### SafeTensors (рекомендуемый формат для обучения) @@ -213,10 +217,11 @@ World-UAV-aug/ from safetensors.torch import load_file stem = "crop_12_4" -aug_dir = Path("World-UAV-aug/Rot/SouthernSuburbs/DB/img") +aug_root = Path("World-UAV-aug") +rel_parent = "Rot/SouthernSuburbs/DB/img" # Zero-copy чтение всех модальностей за ~0.1ms -data = load_file(aug_dir / f"{stem}.safetensors", device="cpu") +data = load_file(aug_root / "safetensors" / rel_parent / f"{stem}.safetensors", device="cpu") depth = data["depth"] # [1, 256, 256] float16, [0, 1] edge = data["edge"] # [1, 256, 256] float16, [0, 1] @@ -236,13 +241,25 @@ from PIL import Image import numpy as np # Depth / Edge / CHM -- grayscale float [0, 1] -depth = np.array(Image.open(aug_dir / f"{stem}_depth.png")) / 255.0 # [H, W] -edge = np.array(Image.open(aug_dir / f"{stem}_edge.png")) / 255.0 -chm = np.array(Image.open(aug_dir / f"{stem}_chm.png")) / 255.0 +depth = np.array(Image.open(aug_root / "depth" / rel_parent / f"{stem}.png")) / 255.0 +edge = np.array(Image.open(aug_root / "edge" / rel_parent / f"{stem}.png")) / 255.0 +chm = np.array(Image.open(aug_root / "chm" / rel_parent / f"{stem}.png")) / 255.0 ``` > PNG визуализации квантуют float16 в uint8 (256 уровней). Для обучения используйте SafeTensors. +### Миграция со старого формата + +Если данные сгенерированы в старом prefix-формате (`crop_12_4_depth.png`), мигрируйте: + +```bash +# Сначала проверить (dry-run) +python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug --dry-run + +# Выполнить миграцию +python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug +``` + ## Скачивание весов Веса скачиваются один раз в `in/weights/` (~10 GB суммарно): diff --git a/scripts/migrate_layout.py b/scripts/migrate_layout.py new file mode 100644 index 0000000..fc7655d --- /dev/null +++ b/scripts/migrate_layout.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Migrate World-UAV-aug from flat prefix layout to directory-based layout. + +Old layout (prefix-based): + World-UAV-aug/Rot/scene/DB/img/crop_12_4_depth.png + World-UAV-aug/Rot/scene/DB/img/crop_12_4_edge.png + World-UAV-aug/Rot/scene/DB/img/crop_12_4_segm.png + World-UAV-aug/Rot/scene/DB/img/crop_12_4_chm.png + World-UAV-aug/Rot/scene/DB/img/crop_12_4_depth.npy + ... + +New layout (directory-based): + World-UAV-aug/depth/Rot/scene/DB/img/crop_12_4.png + World-UAV-aug/edge/Rot/scene/DB/img/crop_12_4.png + World-UAV-aug/segm/Rot/scene/DB/img/crop_12_4.png + World-UAV-aug/chm/Rot/scene/DB/img/crop_12_4.png + World-UAV-aug/npy/depth/Rot/scene/DB/img/crop_12_4.npy + ... + +Usage: + python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug + python scripts/migrate_layout.py /mnt/data1tb/cvgl_datasets/World-UAV-aug --dry-run +""" + +from __future__ import annotations + +import argparse +import os +import shutil +from pathlib import Path + +# Suffix → (modality, extension) +_SUFFIX_MAP = { + "_depth.png": ("depth", ".png"), + "_depth.npy": ("depth", ".npy"), + "_edge.png": ("edge", ".png"), + "_edge.npy": ("edge", ".npy"), + "_segm.png": ("segm", ".png"), + "_segm.npy": ("segm", ".npy"), + "_chm.png": ("chm", ".png"), + "_chm.npy": ("chm", ".npy"), +} + +# Top-level dirs to skip (they belong to the new layout). +_SKIP_DIRS = {"depth", "edge", "segm", "chm", "npy", "safetensors"} + + +def find_old_files(root: Path) -> list[tuple[Path, str, str, str]]: + """Find all files with old prefix-based naming. + + Returns list of (old_path, rel_parent, stem, suffix_key). + """ + results = [] + for p in sorted(root.rglob("*")): + if not p.is_file(): + continue + # Skip files already in new-layout dirs. + try: + rel = p.relative_to(root) + except ValueError: + continue + if rel.parts and rel.parts[0] in _SKIP_DIRS: + continue + + name = p.name + for suffix_key, (modality, ext) in _SUFFIX_MAP.items(): + if name.endswith(suffix_key): + stem = name[: -len(suffix_key)] + rel_parent = str(p.parent.relative_to(root)) + results.append((p, rel_parent, stem, suffix_key)) + break + return results + + +def compute_new_path( + root: Path, rel_parent: str, stem: str, suffix_key: str, +) -> Path: + """Compute the new path for a file.""" + modality, ext = _SUFFIX_MAP[suffix_key] + if ext == ".npy": + return root / "npy" / modality / rel_parent / f"{stem}{ext}" + else: + return root / modality / rel_parent / f"{stem}{ext}" + + +def migrate(root: Path, dry_run: bool = False) -> None: + files = find_old_files(root) + if not files: + print(f"No old-layout files found in {root}") + return + + print(f"Found {len(files)} files to migrate in {root}") + + moved = 0 + for old_path, rel_parent, stem, suffix_key in files: + new_path = compute_new_path(root, rel_parent, stem, suffix_key) + + if dry_run: + print(f" {old_path.relative_to(root)} → {new_path.relative_to(root)}") + else: + new_path.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(old_path), str(new_path)) + moved += 1 + + if dry_run: + print(f"\nDry run: {len(files)} files would be moved.") + else: + print(f"Moved {moved} files.") + + # Clean up empty directories left behind. + if not dry_run: + _cleanup_empty_dirs(root) + + +def _cleanup_empty_dirs(root: Path) -> None: + """Remove empty directories (bottom-up).""" + removed = 0 + for dirpath, dirnames, filenames in os.walk(str(root), topdown=False): + d = Path(dirpath) + if d == root: + continue + # Skip new-layout dirs. + try: + rel = d.relative_to(root) + except ValueError: + continue + if rel.parts and rel.parts[0] in _SKIP_DIRS: + continue + if not any(d.iterdir()): + d.rmdir() + removed += 1 + if removed: + print(f"Cleaned up {removed} empty directories.") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Migrate World-UAV-aug from prefix to directory layout", + ) + parser.add_argument("root", type=Path, help="Path to World-UAV-aug directory") + parser.add_argument("--dry-run", action="store_true", + help="Show what would be moved without doing it") + args = parser.parse_args() + + if not args.root.is_dir(): + print(f"Error: {args.root} is not a directory") + return + + migrate(args.root, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/src/augmentor/dataset.py b/src/augmentor/dataset.py index 480bd45..8823aa0 100644 --- a/src/augmentor/dataset.py +++ b/src/augmentor/dataset.py @@ -11,6 +11,8 @@ from PIL import Image from torch.utils.data import Dataset from torchvision import transforms +from src.augmentor.io_utils import npy_path, vis_path, safetensors_path + logger = logging.getLogger(__name__) EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp"} @@ -33,7 +35,8 @@ class ImageRecord(NamedTuple): abs_path: Path rel_path: str stem: str - output_dir: Path + output_root: Path + rel_parent: str def is_query_record(record: ImageRecord) -> bool: @@ -104,7 +107,11 @@ def discover_images( rel = p.relative_to(root) records.append(ImageRecord( - abs_path=p, rel_path=str(rel), stem=p.stem, output_dir=Path(), + abs_path=p, + rel_path=str(rel), + stem=p.stem, + output_root=Path(), + rel_parent=str(rel.parent), )) if n_skipped_incomplete > 0: @@ -119,17 +126,12 @@ def attach_output_dirs( records: list[ImageRecord], output_root: Path, ) -> list[ImageRecord]: - """Set output_dir for each record: output_root / /.""" - out: list[ImageRecord] = [] - for r in records: - rel = Path(r.rel_path) - odir = output_root / rel.parent - out.append(r._replace(output_dir=odir)) - return out + """Set output_root for each record.""" + return [r._replace(output_root=output_root) for r in records] -# Suffix appended to stem for each stage: {stem}_{suffix}.npy -STAGE_SUFFIX: dict[str, str] = { +# Modality name for each stage (used for folder names). +STAGE_MODALITY: dict[str, str] = { "depth": "depth", "edges": "edge", "segmentation": "segm", @@ -137,12 +139,6 @@ STAGE_SUFFIX: dict[str, str] = { } -def stage_filename(stem: str, stage: str, ext: str = ".npy") -> str: - """Build output filename: e.g. crop_12_4_depth.npy""" - suffix = STAGE_SUFFIX.get(stage, stage) - return f"{stem}_{suffix}{ext}" - - def filter_completed( records: list[ImageRecord], stage: str, @@ -150,16 +146,15 @@ def filter_completed( """Return (pending_records, n_skipped) for a given stage.""" if stage == "consolidate": return filter_consolidated(records) - suffix = STAGE_SUFFIX.get(stage) - if suffix is None: + modality = STAGE_MODALITY.get(stage) + if modality is None: return records, 0 pending: list[ImageRecord] = [] skipped = 0 for r in records: - # Check both .npy and .png — either means the stage is done. - npy = r.output_dir / f"{r.stem}_{suffix}.npy" - png = r.output_dir / f"{r.stem}_{suffix}.png" - if npy.exists() or png.exists(): + np_p = npy_path(r.output_root, modality, r.rel_parent, r.stem) + vis_p = vis_path(r.output_root, modality, r.rel_parent, r.stem) + if np_p.exists() or vis_p.exists(): skipped += 1 else: pending.append(r) @@ -173,7 +168,7 @@ def filter_consolidated( pending: list[ImageRecord] = [] skipped = 0 for r in records: - st = r.output_dir / f"{r.stem}.safetensors" + st = safetensors_path(r.output_root, r.rel_parent, r.stem) if st.exists(): skipped += 1 else: @@ -212,5 +207,6 @@ class AugmentDataset(Dataset): "image_raw": tensor, "rel_path": r.rel_path, "stem": r.stem, - "output_dir": str(r.output_dir), + "output_root": str(r.output_root), + "rel_parent": r.rel_parent, } diff --git a/src/augmentor/io_utils.py b/src/augmentor/io_utils.py index e89a704..5f7f94e 100644 --- a/src/augmentor/io_utils.py +++ b/src/augmentor/io_utils.py @@ -1,4 +1,17 @@ -"""I/O utilities: saving depth / edges / segmentation / 6-ch concat. +"""I/O utilities: saving depth / edges / segmentation / safetensors. + +Directory-based output layout — modality determines the folder, not file suffix: + + output_root/ + ├── depth/{rel_parent}/{stem}.png # vis + ├── edge/{rel_parent}/{stem}.png + ├── segm/{rel_parent}/{stem}.png + ├── chm/{rel_parent}/{stem}.png + ├── npy/depth/{rel_parent}/{stem}.npy # intermediate float16/uint8 + ├── npy/edge/{rel_parent}/{stem}.npy + ├── npy/segm/{rel_parent}/{stem}.npy + ├── npy/chm/{rel_parent}/{stem}.npy + └── safetensors/{rel_parent}/{stem}.safetensors No global config imports — all parameters passed explicitly. """ @@ -45,6 +58,29 @@ def shutdown_io_pool() -> None: _io_pool = None +# --------------------------------------------------------------------------- +# Path helpers +# --------------------------------------------------------------------------- + +def vis_path(output_root: Path, modality: str, rel_parent: str, stem: str) -> Path: + """Build: output_root / modality / rel_parent / stem.png""" + return output_root / modality / rel_parent / f"{stem}.png" + + +def npy_path(output_root: Path, modality: str, rel_parent: str, stem: str) -> Path: + """Build: output_root / npy / modality / rel_parent / stem.npy""" + return output_root / "npy" / modality / rel_parent / f"{stem}.npy" + + +def safetensors_path(output_root: Path, rel_parent: str, stem: str) -> Path: + """Build: output_root / safetensors / rel_parent / stem.safetensors""" + return output_root / "safetensors" / rel_parent / f"{stem}.safetensors" + + +# --------------------------------------------------------------------------- +# Palette +# --------------------------------------------------------------------------- + # Intuitive RS segmentation palette: index → RGB. _FIXED_PALETTE = np.array([ [0, 0, 0], # 0: background — black @@ -75,6 +111,10 @@ def make_palette(num_classes: int, seed: int = 42) -> np.ndarray: return palette +# --------------------------------------------------------------------------- +# Low-level atomic save +# --------------------------------------------------------------------------- + def _atomic_save_npy(arr: np.ndarray, path: Path) -> None: """Write .npy atomically via temp file + rename.""" path.parent.mkdir(parents=True, exist_ok=True) @@ -104,130 +144,118 @@ def _apply_colormap(gray: np.ndarray, cmap_name: str = "turbo") -> np.ndarray: return lut[idx] +# --------------------------------------------------------------------------- +# Save float16 maps (depth, edge, chm) +# --------------------------------------------------------------------------- + def _save_float16_map( data: torch.Tensor, - output_dir: Path, + output_root: Path, + rel_parent: str, stem: str, - suffix: str, + modality: str, save_npy: bool = True, save_vis: bool = True, colormap: str | None = None, ) -> None: - """Save a [1, H, W] float tensor as {stem}_{suffix}.npy (float16) + optional vis. - - Args: - colormap: If set (e.g. "turbo"), apply colormap for RGB visualization. - If None, save grayscale. - """ + """Save a [1, H, W] float tensor as .npy (float16) + optional vis .png.""" arr = data.half().numpy() if save_npy: - _atomic_save_npy(arr, output_dir / f"{stem}_{suffix}.npy") + p = npy_path(output_root, modality, rel_parent, stem) + _atomic_save_npy(arr, p) if save_vis: gray = arr.squeeze(0).astype(np.float32) if colormap: vis = _apply_colormap(gray, colormap) else: vis = (gray * 255).clip(0, 255).astype(np.uint8) - Image.fromarray(vis).save(output_dir / f"{stem}_{suffix}.png") + p = vis_path(output_root, modality, rel_parent, stem) + p.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray(vis).save(p) -def save_depth(depth: torch.Tensor, output_dir: Path, stem: str, - save_npy: bool = True, save_vis: bool = True) -> None: - _save_float16_map(depth, output_dir, stem, "depth", save_npy, save_vis) +def save_depth(depth: torch.Tensor, output_root: Path, rel_parent: str, + stem: str, save_npy: bool = True, save_vis: bool = True) -> None: + _save_float16_map(depth, output_root, rel_parent, stem, "depth", save_npy, save_vis) -def save_depth_async(depth: torch.Tensor, output_dir: Path, stem: str, - save_npy: bool = True, save_vis: bool = True) -> None: - get_io_pool().submit(save_depth, depth.clone().cpu(), output_dir, stem, save_npy, save_vis) +def save_depth_async(depth: torch.Tensor, output_root: Path, rel_parent: str, + stem: str, save_npy: bool = True, save_vis: bool = True) -> None: + get_io_pool().submit(save_depth, depth.clone().cpu(), output_root, rel_parent, + stem, save_npy, save_vis) -def save_chmv2(depth: torch.Tensor, output_dir: Path, stem: str, - save_npy: bool = True, save_vis: bool = True) -> None: - _save_float16_map(depth, output_dir, stem, "chm", save_npy, save_vis) +def save_chmv2(depth: torch.Tensor, output_root: Path, rel_parent: str, + stem: str, save_npy: bool = True, save_vis: bool = True) -> None: + _save_float16_map(depth, output_root, rel_parent, stem, "chm", save_npy, save_vis) -def save_chmv2_async(depth: torch.Tensor, output_dir: Path, stem: str, - save_npy: bool = True, save_vis: bool = True) -> None: - get_io_pool().submit(save_chmv2, depth.clone().cpu(), output_dir, stem, save_npy, save_vis) +def save_chmv2_async(depth: torch.Tensor, output_root: Path, rel_parent: str, + stem: str, save_npy: bool = True, save_vis: bool = True) -> None: + get_io_pool().submit(save_chmv2, depth.clone().cpu(), output_root, rel_parent, + stem, save_npy, save_vis) -def save_edges(edges: torch.Tensor, output_dir: Path, stem: str, - save_npy: bool = True, save_vis: bool = True) -> None: - _save_float16_map(edges, output_dir, stem, "edge", save_npy, save_vis) +def save_edges(edges: torch.Tensor, output_root: Path, rel_parent: str, + stem: str, save_npy: bool = True, save_vis: bool = True) -> None: + _save_float16_map(edges, output_root, rel_parent, stem, "edge", save_npy, save_vis) -def save_edges_async(edges: torch.Tensor, output_dir: Path, stem: str, - save_npy: bool = True, save_vis: bool = True) -> None: - get_io_pool().submit(save_edges, edges.clone().cpu(), output_dir, stem, save_npy, save_vis) +def save_edges_async(edges: torch.Tensor, output_root: Path, rel_parent: str, + stem: str, save_npy: bool = True, save_vis: bool = True) -> None: + get_io_pool().submit(save_edges, edges.clone().cpu(), output_root, rel_parent, + stem, save_npy, save_vis) +# --------------------------------------------------------------------------- +# Save segmentation +# --------------------------------------------------------------------------- + def save_segmentation( seg_ids: torch.Tensor, - output_dir: Path, + output_root: Path, + rel_parent: str, stem: str, save_npy: bool = True, save_vis: bool = True, num_classes: int = 150, ) -> None: - """Save segmentation map [1, H, W] uint8 as {stem}_segm.npy.""" + """Save segmentation map [1, H, W] uint8.""" arr = seg_ids.byte().numpy() if save_npy: - _atomic_save_npy(arr, output_dir / f"{stem}_segm.npy") + _atomic_save_npy(arr, npy_path(output_root, "segm", rel_parent, stem)) if save_vis: palette = make_palette(num_classes) 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") + p = vis_path(output_root, "segm", rel_parent, stem) + p.parent.mkdir(parents=True, exist_ok=True) + img.save(p) def save_segmentation_async( seg_ids: torch.Tensor, - output_dir: Path, + output_root: Path, + rel_parent: str, stem: str, save_npy: bool = True, save_vis: bool = True, num_classes: int = 150, ) -> None: get_io_pool().submit( - save_segmentation, seg_ids.clone().cpu(), output_dir, stem, - save_npy, save_vis, num_classes, + save_segmentation, seg_ids.clone().cpu(), output_root, rel_parent, + stem, save_npy, save_vis, num_classes, ) -def save_concat_6ch( - rgb: torch.Tensor, - output_dir: Path, - stem: str, - num_classes: int = 150, -) -> None: - """Assemble 6-ch tensor from saved .npy files.""" - depth_path = output_dir / f"{stem}_depth.npy" - edges_path = output_dir / f"{stem}_edge.npy" - seg_path = output_dir / f"{stem}_segm.npy" - - if not (depth_path.exists() and edges_path.exists() and seg_path.exists()): - logger.warning("⚠️ Missing modality .npy for %s, skipping concat.", stem) - return - - depth = torch.from_numpy(np.load(depth_path).astype(np.float32)) - edges = torch.from_numpy(np.load(edges_path).astype(np.float32)) - seg_ids = torch.from_numpy(np.load(seg_path).astype(np.float32)) - seg_float = seg_ids / max(float(num_classes), 1.0) - - concat = torch.cat([rgb, depth, edges, seg_float], dim=0) - _atomic_save_npy(concat.numpy(), output_dir / f"{stem}_concat.npy") - - # --------------------------------------------------------------------------- # SafeTensors: consolidate all modalities into one file per image # --------------------------------------------------------------------------- -# Mapping from stage suffix → (dtype to store, source extension priority) _MODALITY_SPEC: dict[str, tuple[torch.dtype, str]] = { "depth": (torch.float16, "depth"), "edge": (torch.float16, "edge"), @@ -237,41 +265,38 @@ _MODALITY_SPEC: dict[str, tuple[torch.dtype, str]] = { def _load_modality_tensor( - output_dir: Path, stem: str, suffix: str, dtype: torch.dtype, + output_root: Path, rel_parent: str, stem: str, + modality: str, dtype: torch.dtype, ) -> torch.Tensor | None: """Load a single modality from .npy or .png, return [1, H, W] tensor or None.""" - npy_path = output_dir / f"{stem}_{suffix}.npy" - png_path = output_dir / f"{stem}_{suffix}.png" + np_p = npy_path(output_root, modality, rel_parent, stem) + vis_p = vis_path(output_root, modality, rel_parent, stem) - if npy_path.exists(): - arr = np.load(npy_path) + if np_p.exists(): + arr = np.load(np_p) t = torch.from_numpy(arr.astype(np.float32 if dtype != torch.uint8 else np.uint8)) if t.ndim == 2: t = t.unsqueeze(0) return t.to(dtype) - if png_path.exists(): - img = np.array(Image.open(png_path)) - if suffix == "segm": - # Palette PNG → need to read class IDs, not RGB. - # If saved as palette mode (P), PIL gives indices directly. - pil = Image.open(png_path) + if vis_p.exists(): + if modality == "segm": + pil = Image.open(vis_p) if pil.mode == "P": img = np.array(pil) else: - # RGB palette render — can't recover class IDs reliably, skip. - logger.debug("Skipping %s_%s.png (RGB palette, no class IDs).", stem, suffix) + logger.debug("Skipping %s segm.png (RGB, no class IDs).", stem) return None t = torch.from_numpy(img.astype(np.uint8)) if t.ndim == 2: t = t.unsqueeze(0) return t else: + img = np.array(Image.open(vis_p)) arr = img.astype(np.float32) / 255.0 if arr.ndim == 2: arr = arr[np.newaxis] elif arr.ndim == 3: - # Grayscale saved as RGB — take first channel. arr = arr[:, :, 0:1].transpose(2, 0, 1) return torch.from_numpy(arr).to(dtype) @@ -279,57 +304,63 @@ def _load_modality_tensor( def consolidate_safetensors( - output_dir: Path, + output_root: Path, + rel_parent: str, stem: str, cleanup_npy: bool = False, ) -> bool: - """Bundle available modalities into {stem}.safetensors. + """Bundle available modalities into one .safetensors file. Returns True if the file was written, False if no modalities found. """ tensors: dict[str, torch.Tensor] = {} - npy_paths: list[Path] = [] + npy_paths_to_clean: list[Path] = [] - for suffix, (dtype, _) in _MODALITY_SPEC.items(): - t = _load_modality_tensor(output_dir, stem, suffix, dtype) + for modality, (dtype, _) in _MODALITY_SPEC.items(): + t = _load_modality_tensor(output_root, rel_parent, stem, modality, dtype) if t is not None: - tensors[suffix] = t - npy_path = output_dir / f"{stem}_{suffix}.npy" - if npy_path.exists(): - npy_paths.append(npy_path) + tensors[modality] = t + np_p = npy_path(output_root, modality, rel_parent, stem) + if np_p.exists(): + npy_paths_to_clean.append(np_p) if not tensors: return False - st_path = output_dir / f"{stem}.safetensors" - output_dir.mkdir(parents=True, exist_ok=True) + st_p = safetensors_path(output_root, rel_parent, stem) + st_p.parent.mkdir(parents=True, exist_ok=True) - # Atomic write via temp file. - fd, tmp = tempfile.mkstemp(suffix=".safetensors", dir=output_dir) + fd, tmp = tempfile.mkstemp(suffix=".safetensors", dir=st_p.parent) os.close(fd) try: _st_save_file(tensors, tmp) - os.replace(tmp, st_path) + os.replace(tmp, st_p) except BaseException: if os.path.exists(tmp): os.remove(tmp) raise if cleanup_npy: - for p in npy_paths: + for p in npy_paths_to_clean: p.unlink(missing_ok=True) return True def consolidate_safetensors_async( - output_dir: Path, + output_root: Path, + rel_parent: str, stem: str, cleanup_npy: bool = False, ) -> None: - get_io_pool().submit(consolidate_safetensors, output_dir, stem, cleanup_npy) + get_io_pool().submit(consolidate_safetensors, output_root, rel_parent, + stem, cleanup_npy) +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- + def setup_logging(log_level: str = "INFO", log_file: Path | None = None) -> None: """Configure root logger with coloredlogs for console + optional file handler.""" import coloredlogs diff --git a/src/main.py b/src/main.py index bccb320..f552be1 100644 --- a/src/main.py +++ b/src/main.py @@ -39,6 +39,7 @@ from src.augmentor.inference import ( infer_segmentation_batch, ) from src.augmentor.io_utils import ( + npy_path, vis_path, save_depth_async, save_chmv2_async, save_edges_async, save_segmentation_async, consolidate_safetensors, setup_logging, shutdown_io_pool, @@ -57,7 +58,6 @@ _STAGE_EMOJI = { "edges": "🔪", "segmentation": "🗺️", "chmv2": "🦕", - "concat": "🧩", "consolidate": "📦", } @@ -97,11 +97,7 @@ def _resolve_image_sizes( records: list[ImageRecord], input_conf: InputConfig, ) -> list[tuple[list[ImageRecord], int, str]]: - """Split records into groups by target resolution. - - Returns list of (records, image_size, label) tuples. When - ``query_image_size == image_size`` a single group is returned (no split). - """ + """Split records into groups by target resolution.""" if input_conf.query_image_size == input_conf.image_size: return [(records, input_conf.image_size, "all")] db_recs, query_recs = split_by_view(records) @@ -149,7 +145,9 @@ def run_depth_stage( for batch in pbar: depths = infer_depth_batch(model, batch["image_raw"], device) for i in range(depths.shape[0]): - save_depth_async(depths[i], Path(batch["output_dir"][i]), + save_depth_async(depths[i], + Path(batch["output_root"][i]), + batch["rel_parent"][i], stem=batch["stem"][i], save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis) @@ -169,9 +167,9 @@ def run_edges_stage( """🔪 Compute Sobel edges from saved depth (CPU, batched).""" valid: list[ImageRecord] = [] for r in records: - depth_png = r.output_dir / f"{r.stem}_depth.png" - depth_npy = r.output_dir / f"{r.stem}_depth.npy" - if depth_png.exists() or depth_npy.exists(): + np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem) + vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem) + if np_p.exists() or vis_p.exists(): valid.append(r) else: logger.warning("⚠️ No depth for %s, skipping edges.", r.rel_path) @@ -185,13 +183,13 @@ def run_edges_stage( chunk = valid[start : start + batch_size] depth_tensors = [] for r in chunk: - npy_path = r.output_dir / f"{r.stem}_depth.npy" - png_path = r.output_dir / f"{r.stem}_depth.png" - if npy_path.exists(): - d = np.load(npy_path).astype(np.float32) + np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem) + vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem) + if np_p.exists(): + d = np.load(np_p).astype(np.float32) else: from PIL import Image - d = np.array(Image.open(png_path)).astype(np.float32) / 255.0 + d = np.array(Image.open(vis_p)).astype(np.float32) / 255.0 if d.ndim == 2: d = d[np.newaxis] depth_tensors.append(torch.from_numpy(d)) @@ -200,7 +198,8 @@ def run_edges_stage( depths = depths.unsqueeze(1) edges_batch = compute_edges_from_depth(depths) for j, r in enumerate(chunk): - save_edges_async(edges_batch[j], r.output_dir, stem=r.stem, + save_edges_async(edges_batch[j], r.output_root, r.rel_parent, + stem=r.stem, save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis) processed += len(chunk) @@ -245,7 +244,9 @@ def run_chmv2_stage( for batch in pbar: depths = infer_chmv2_batch(model, processor, batch["image_raw"], device) for i in range(depths.shape[0]): - save_chmv2_async(depths[i], Path(batch["output_dir"][i]), + save_chmv2_async(depths[i], + Path(batch["output_root"][i]), + batch["rel_parent"][i], stem=batch["stem"][i], save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis) @@ -303,7 +304,9 @@ def run_segmentation_stage( ) for j in range(segs.shape[0]): save_segmentation_async( - segs[j], Path(batch["output_dir"][j]), stem=batch["stem"][j], + segs[j], Path(batch["output_root"][j]), + batch["rel_parent"][j], + stem=batch["stem"][j], save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis, num_classes=num_classes, ) @@ -326,7 +329,8 @@ def run_consolidate_stage( colour="magenta", bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]") for r in pbar: - ok = consolidate_safetensors(r.output_dir, r.stem, cleanup_npy=cleanup_npy) + ok = consolidate_safetensors(r.output_root, r.rel_parent, r.stem, + cleanup_npy=cleanup_npy) if ok: written += 1 pbar.set_postfix(written=f"{written}/{total}") @@ -371,15 +375,6 @@ def run_pipeline( logger.error("❌ No images found. Check input_root in pipeline.gin.") return - # Pre-create all output directories in one pass. - logger.info("📁 Pre-creating output directories...") - seen_dirs: set[str] = set() - for r in all_records: - d = str(r.output_dir) - if d not in seen_dirs: - 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: @@ -457,7 +452,7 @@ def run_pipeline( # Manifest. manifest = { - "pipeline_version": "3.3.0-safetensors", + "pipeline_version": "4.0.0-dir-layout", "image_size_db": input_conf.image_size, "image_size_query": input_conf.query_image_size, "profile": hw_conf.profile_name, @@ -501,16 +496,7 @@ def run_pipeline( # --------------------------------------------------------------------------- def main() -> None: - """Load all gin configs and run the augmentation pipeline. - - Supports CLI gin overrides for quick mode switches:: - - # Process only query (drone) images: - python -m src.main --gin "PipelineConfig.source = 'query'" - - # Process only db (satellite) images: - python -m src.main --gin "PipelineConfig.source = 'db'" - """ + """Load all gin configs and run the augmentation pipeline.""" import argparse parser = argparse.ArgumentParser(description="Augmentation pipeline") @@ -523,7 +509,6 @@ def main() -> None: proj_dir = get_proj_dir() path2cfg = f"{proj_dir}in/config_files/" - # Load configs with optional CLI overrides. if args.gin: import gin as _gin cfg_dir = Path(path2cfg) diff --git a/src/tests/test_dataset.py b/src/tests/test_dataset.py index 9d7a1cc..c45b538 100644 --- a/src/tests/test_dataset.py +++ b/src/tests/test_dataset.py @@ -16,6 +16,7 @@ from src.augmentor.dataset import ( filter_completed, INCOMPLETE_SCENES, ) +from src.augmentor.io_utils import npy_path class TestDiscoverImages: @@ -42,9 +43,7 @@ class TestDiscoverImages: (tmp_path / "good.png").write_bytes( Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).tobytes() ) - # Only actually valid images with correct extensions are found. records = discover_images(tmp_path) - # good.png is not a valid PNG file (raw bytes), but discover only checks extension. assert all(r.abs_path.suffix in {".png", ".jpg", ".jpeg", ".bmp"} for r in records) def test_subset_filter(self, tmp_path: Path) -> None: @@ -65,14 +64,13 @@ class TestDiscoverImages: class TestAttachOutputDirs: - def test_output_dir_structure(self, fake_image_dir: Path, tmp_path: Path) -> None: + def test_output_root_structure(self, fake_image_dir: Path, tmp_path: Path) -> None: records = discover_images(fake_image_dir) out_root = tmp_path / "output" attached = attach_output_dirs(records, out_root) for r in attached: - assert str(r.output_dir).startswith(str(out_root)) - # output_dir is the parent directory (same structure, no per-image subfolder) - assert r.output_dir == out_root / Path(r.rel_path).parent + assert r.output_root == out_root + assert r.rel_parent == str(Path(r.rel_path).parent) class TestFilterCompleted: @@ -83,8 +81,9 @@ class TestFilterCompleted: def test_skips_completed(self, sample_records: list[ImageRecord]) -> None: r = sample_records[0] - r.output_dir.mkdir(parents=True, exist_ok=True) - np.save(r.output_dir / f"{r.stem}_depth.npy", np.zeros((1, 8, 8))) + p = npy_path(r.output_root, "depth", r.rel_parent, r.stem) + p.parent.mkdir(parents=True, exist_ok=True) + np.save(p, np.zeros((1, 8, 8))) pending, skipped = filter_completed(sample_records, "depth") assert skipped == 1 assert len(pending) == len(sample_records) - 1 @@ -111,4 +110,4 @@ class TestAugmentDataset: def test_getitem_keys(self, sample_records: list[ImageRecord]) -> None: ds = AugmentDataset(sample_records, image_size=32) item = ds[0] - assert set(item.keys()) == {"image_raw", "rel_path", "stem", "output_dir"} + assert set(item.keys()) == {"image_raw", "rel_path", "stem", "output_root", "rel_parent"} diff --git a/src/tests/test_io_utils.py b/src/tests/test_io_utils.py index 7d4bf41..0eaabbf 100644 --- a/src/tests/test_io_utils.py +++ b/src/tests/test_io_utils.py @@ -11,6 +11,7 @@ from safetensors.torch import load_file as st_load_file from src.augmentor.io_utils import ( make_palette, + npy_path, vis_path, safetensors_path, save_chmv2, save_chmv2_async, save_depth, @@ -19,7 +20,6 @@ from src.augmentor.io_utils import ( save_edges_async, save_segmentation, save_segmentation_async, - save_concat_6ch, consolidate_safetensors, consolidate_safetensors_async, shutdown_io_pool, @@ -53,126 +53,143 @@ class TestAtomicSaveNpy: assert len(tmp_files) == 0 +# --------------------------------------------------------------------------- +# Path helpers +# --------------------------------------------------------------------------- + +class TestPathHelpers: + def test_npy_path(self) -> None: + p = npy_path(Path("/out"), "depth", "Rot/scene/DB/img", "crop_1") + assert p == Path("/out/npy/depth/Rot/scene/DB/img/crop_1.npy") + + def test_vis_path(self) -> None: + p = vis_path(Path("/out"), "segm", "Rot/scene/DB/img", "crop_1") + assert p == Path("/out/segm/Rot/scene/DB/img/crop_1.png") + + def test_safetensors_path(self) -> None: + p = safetensors_path(Path("/out"), "Rot/scene/DB/img", "crop_1") + assert p == Path("/out/safetensors/Rot/scene/DB/img/crop_1.safetensors") + + # --------------------------------------------------------------------------- # save_depth # --------------------------------------------------------------------------- +_RP = "sub" # rel_parent for tests + + class TestSaveDepth: def test_saves_float16_npy(self, tmp_path: Path) -> None: - save_depth(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=True, save_vis=False) - arr = np.load(tmp_path / "img01_depth.npy") + save_depth(torch.rand(1, 32, 32), tmp_path, _RP, "img01", + save_npy=True, save_vis=False) + p = npy_path(tmp_path, "depth", _RP, "img01") + arr = np.load(p) assert arr.dtype == np.float16 assert arr.shape == (1, 32, 32) def test_saves_vis_png(self, tmp_path: Path) -> None: - save_depth(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=False, save_vis=True) - assert (tmp_path / "img01_depth.png").exists() - assert not (tmp_path / "img01_depth.npy").exists() + save_depth(torch.rand(1, 32, 32), tmp_path, _RP, "img01", + save_npy=False, save_vis=True) + assert vis_path(tmp_path, "depth", _RP, "img01").exists() + assert not npy_path(tmp_path, "depth", _RP, "img01").exists() def test_npy_values_in_range(self, tmp_path: Path) -> None: - save_depth(torch.rand(1, 16, 16), tmp_path, "img01") - arr = np.load(tmp_path / "img01_depth.npy").astype(np.float32) + save_depth(torch.rand(1, 16, 16), tmp_path, _RP, "img01") + p = npy_path(tmp_path, "depth", _RP, "img01") + arr = np.load(p).astype(np.float32) assert arr.min() >= 0.0 assert arr.max() <= 1.0 class TestSaveChmv2: def test_saves_float16_npy(self, tmp_path: Path) -> None: - save_chmv2(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=True, save_vis=False) - arr = np.load(tmp_path / "img01_chm.npy") + save_chmv2(torch.rand(1, 32, 32), tmp_path, _RP, "img01", + save_npy=True, save_vis=False) + arr = np.load(npy_path(tmp_path, "chm", _RP, "img01")) assert arr.dtype == np.float16 assert arr.shape == (1, 32, 32) def test_saves_vis_png(self, tmp_path: Path) -> None: - save_chmv2(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=False, save_vis=True) - assert (tmp_path / "img01_chm.png").exists() - assert not (tmp_path / "img01_chm.npy").exists() + save_chmv2(torch.rand(1, 32, 32), tmp_path, _RP, "img01", + save_npy=False, save_vis=True) + assert vis_path(tmp_path, "chm", _RP, "img01").exists() + assert not npy_path(tmp_path, "chm", _RP, "img01").exists() def test_npy_values_in_range(self, tmp_path: Path) -> None: - save_chmv2(torch.rand(1, 16, 16), tmp_path, "img01") - arr = np.load(tmp_path / "img01_chm.npy").astype(np.float32) + save_chmv2(torch.rand(1, 16, 16), tmp_path, _RP, "img01") + arr = np.load(npy_path(tmp_path, "chm", _RP, "img01")).astype(np.float32) assert arr.min() >= 0.0 assert arr.max() <= 1.0 class TestSaveEdges: def test_saves_float16_npy(self, tmp_path: Path) -> None: - save_edges(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=True, save_vis=False) - arr = np.load(tmp_path / "img01_edge.npy") + save_edges(torch.rand(1, 32, 32), tmp_path, _RP, "img01", + save_npy=True, save_vis=False) + arr = np.load(npy_path(tmp_path, "edge", _RP, "img01")) assert arr.dtype == np.float16 def test_saves_vis(self, tmp_path: Path) -> None: - save_edges(torch.rand(1, 32, 32), tmp_path, "img01", save_npy=False, save_vis=True) - assert (tmp_path / "img01_edge.png").exists() + save_edges(torch.rand(1, 32, 32), tmp_path, _RP, "img01", + save_npy=False, save_vis=True) + assert vis_path(tmp_path, "edge", _RP, "img01").exists() class TestSaveSegmentation: def test_saves_uint8_npy(self, tmp_path: Path) -> None: seg = torch.randint(0, 5, (1, 32, 32), dtype=torch.uint8) - save_segmentation(seg, tmp_path, "img01", save_npy=True, save_vis=False, num_classes=5) - arr = np.load(tmp_path / "img01_segm.npy") + save_segmentation(seg, tmp_path, _RP, "img01", + save_npy=True, save_vis=False, num_classes=5) + arr = np.load(npy_path(tmp_path, "segm", _RP, "img01")) assert arr.dtype == np.uint8 assert arr.shape == (1, 32, 32) def test_saves_vis(self, tmp_path: Path) -> None: seg = torch.randint(0, 3, (1, 32, 32), dtype=torch.uint8) - save_segmentation(seg, tmp_path, "img01", save_npy=False, save_vis=True, num_classes=3) - assert (tmp_path / "img01_segm.png").exists() + save_segmentation(seg, tmp_path, _RP, "img01", + save_npy=False, save_vis=True, num_classes=3) + assert vis_path(tmp_path, "segm", _RP, "img01").exists() def test_int_tensor_also_works(self, tmp_path: Path) -> None: seg = torch.randint(0, 5, (1, 16, 16), dtype=torch.int64) - save_segmentation(seg, tmp_path, "img01", num_classes=5) - arr = np.load(tmp_path / "img01_segm.npy") + save_segmentation(seg, tmp_path, _RP, "img01", num_classes=5) + arr = np.load(npy_path(tmp_path, "segm", _RP, "img01")) assert arr.dtype == np.uint8 class TestAsyncSaves: def test_depth_async(self, tmp_path: Path) -> None: - save_depth_async(torch.rand(1, 16, 16), tmp_path, "img01", save_npy=True, save_vis=False) + save_depth_async(torch.rand(1, 16, 16), tmp_path, _RP, "img01", + save_npy=True, save_vis=False) shutdown_io_pool() - assert (tmp_path / "img01_depth.npy").exists() + assert npy_path(tmp_path, "depth", _RP, "img01").exists() def test_chmv2_async(self, tmp_path: Path) -> None: - save_chmv2_async(torch.rand(1, 16, 16), tmp_path, "img01", save_npy=True, save_vis=False) + save_chmv2_async(torch.rand(1, 16, 16), tmp_path, _RP, "img01", + save_npy=True, save_vis=False) shutdown_io_pool() - assert (tmp_path / "img01_chm.npy").exists() + assert npy_path(tmp_path, "chm", _RP, "img01").exists() def test_edges_async(self, tmp_path: Path) -> None: - save_edges_async(torch.rand(1, 16, 16), tmp_path, "img01", save_npy=True, save_vis=False) + save_edges_async(torch.rand(1, 16, 16), tmp_path, _RP, "img01", + save_npy=True, save_vis=False) shutdown_io_pool() - assert (tmp_path / "img01_edge.npy").exists() + assert npy_path(tmp_path, "edge", _RP, "img01").exists() def test_segmentation_async(self, tmp_path: Path) -> None: seg = torch.randint(0, 3, (1, 16, 16), dtype=torch.uint8) - save_segmentation_async(seg, tmp_path, "img01", save_npy=True, save_vis=False, num_classes=3) + save_segmentation_async(seg, tmp_path, _RP, "img01", + save_npy=True, save_vis=False, num_classes=3) shutdown_io_pool() - assert (tmp_path / "img01_segm.npy").exists() + assert npy_path(tmp_path, "segm", _RP, "img01").exists() def test_multiple_async_writes(self, tmp_path: Path) -> None: for i in range(8): - save_depth_async(torch.rand(1, 8, 8), tmp_path, f"img_{i}", + save_depth_async(torch.rand(1, 8, 8), tmp_path, _RP, f"img_{i}", save_npy=True, save_vis=False) shutdown_io_pool() for i in range(8): - assert (tmp_path / f"img_{i}_depth.npy").exists() - - -class TestSaveConcat6ch: - def test_concat_shape(self, tmp_path: Path) -> None: - H, W = 16, 16 - rgb = torch.rand(3, H, W) - np.save(tmp_path / "img01_depth.npy", np.random.rand(1, H, W).astype(np.float16)) - np.save(tmp_path / "img01_edge.npy", np.random.rand(1, H, W).astype(np.float16)) - np.save(tmp_path / "img01_segm.npy", np.random.randint(0, 5, (1, H, W)).astype(np.uint8)) - - save_concat_6ch(rgb, tmp_path, stem="img01", num_classes=5) - concat = np.load(tmp_path / "img01_concat.npy") - assert concat.shape == (6, H, W) - - def test_skips_when_missing(self, tmp_path: Path) -> None: - rgb = torch.rand(3, 8, 8) - save_concat_6ch(rgb, tmp_path, stem="img01", num_classes=5) - assert not (tmp_path / "img01_concat.npy").exists() + assert npy_path(tmp_path, "depth", _RP, f"img_{i}").exists() # --------------------------------------------------------------------------- @@ -190,9 +207,9 @@ class TestMakePalette: np.testing.assert_array_equal(pal[0], [0, 0, 0]) def test_cache_returns_same(self) -> None: - p1 = make_palette(7) - p2 = make_palette(7) - assert p1 is p2 + pal = make_palette(7) + pal2 = make_palette(7) + assert pal is pal2 # --------------------------------------------------------------------------- @@ -201,32 +218,31 @@ class TestMakePalette: class TestConsolidateSafetensors: def _save_all_modalities(self, tmp_path: Path, stem: str = "img01") -> None: - """Helper: save all 4 modalities as .npy files.""" H, W = 32, 32 - save_depth(torch.rand(1, H, W), tmp_path, stem, save_npy=True, save_vis=False) - save_edges(torch.rand(1, H, W), tmp_path, stem, save_npy=True, save_vis=False) - save_chmv2(torch.rand(1, H, W), tmp_path, stem, save_npy=True, save_vis=False) + save_depth(torch.rand(1, H, W), tmp_path, _RP, stem, save_npy=True, save_vis=False) + save_edges(torch.rand(1, H, W), tmp_path, _RP, stem, save_npy=True, save_vis=False) + save_chmv2(torch.rand(1, H, W), tmp_path, _RP, stem, save_npy=True, save_vis=False) save_segmentation( torch.randint(0, 11, (1, H, W), dtype=torch.uint8), - tmp_path, stem, save_npy=True, save_vis=False, num_classes=11, + tmp_path, _RP, stem, save_npy=True, save_vis=False, num_classes=11, ) def test_creates_safetensors_file(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - ok = consolidate_safetensors(tmp_path, "img01") + ok = consolidate_safetensors(tmp_path, _RP, "img01") assert ok - assert (tmp_path / "img01.safetensors").exists() + assert safetensors_path(tmp_path, _RP, "img01").exists() def test_contains_all_modalities(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - consolidate_safetensors(tmp_path, "img01") - data = st_load_file(tmp_path / "img01.safetensors") + consolidate_safetensors(tmp_path, _RP, "img01") + data = st_load_file(safetensors_path(tmp_path, _RP, "img01")) assert set(data.keys()) == {"depth", "edge", "chm", "segm"} def test_dtypes_correct(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - consolidate_safetensors(tmp_path, "img01") - data = st_load_file(tmp_path / "img01.safetensors") + consolidate_safetensors(tmp_path, _RP, "img01") + data = st_load_file(safetensors_path(tmp_path, _RP, "img01")) assert data["depth"].dtype == torch.float16 assert data["edge"].dtype == torch.float16 assert data["chm"].dtype == torch.float16 @@ -234,70 +250,64 @@ class TestConsolidateSafetensors: def test_shapes_correct(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - consolidate_safetensors(tmp_path, "img01") - data = st_load_file(tmp_path / "img01.safetensors") + consolidate_safetensors(tmp_path, _RP, "img01") + data = st_load_file(safetensors_path(tmp_path, _RP, "img01")) for key in ("depth", "edge", "chm", "segm"): assert data[key].shape == (1, 32, 32), f"{key} shape mismatch" def test_partial_modalities(self, tmp_path: Path) -> None: - """Only depth + edge available → safetensors has just those.""" - save_depth(torch.rand(1, 16, 16), tmp_path, "img02", save_npy=True, save_vis=False) - save_edges(torch.rand(1, 16, 16), tmp_path, "img02", save_npy=True, save_vis=False) - ok = consolidate_safetensors(tmp_path, "img02") + save_depth(torch.rand(1, 16, 16), tmp_path, _RP, "img02", save_npy=True, save_vis=False) + save_edges(torch.rand(1, 16, 16), tmp_path, _RP, "img02", save_npy=True, save_vis=False) + ok = consolidate_safetensors(tmp_path, _RP, "img02") assert ok - data = st_load_file(tmp_path / "img02.safetensors") + data = st_load_file(safetensors_path(tmp_path, _RP, "img02")) assert set(data.keys()) == {"depth", "edge"} def test_no_modalities_returns_false(self, tmp_path: Path) -> None: - ok = consolidate_safetensors(tmp_path, "missing") + ok = consolidate_safetensors(tmp_path, _RP, "missing") assert not ok - assert not (tmp_path / "missing.safetensors").exists() + assert not safetensors_path(tmp_path, _RP, "missing").exists() def test_cleanup_npy(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - consolidate_safetensors(tmp_path, "img01", cleanup_npy=True) - assert (tmp_path / "img01.safetensors").exists() - assert not (tmp_path / "img01_depth.npy").exists() - assert not (tmp_path / "img01_edge.npy").exists() - assert not (tmp_path / "img01_chm.npy").exists() - assert not (tmp_path / "img01_segm.npy").exists() + consolidate_safetensors(tmp_path, _RP, "img01", cleanup_npy=True) + assert safetensors_path(tmp_path, _RP, "img01").exists() + assert not npy_path(tmp_path, "depth", _RP, "img01").exists() + assert not npy_path(tmp_path, "edge", _RP, "img01").exists() + assert not npy_path(tmp_path, "chm", _RP, "img01").exists() + assert not npy_path(tmp_path, "segm", _RP, "img01").exists() def test_no_cleanup_keeps_npy(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - consolidate_safetensors(tmp_path, "img01", cleanup_npy=False) - assert (tmp_path / "img01.safetensors").exists() - assert (tmp_path / "img01_depth.npy").exists() + consolidate_safetensors(tmp_path, _RP, "img01", cleanup_npy=False) + assert safetensors_path(tmp_path, _RP, "img01").exists() + assert npy_path(tmp_path, "depth", _RP, "img01").exists() def test_async_consolidation(self, tmp_path: Path) -> None: self._save_all_modalities(tmp_path) - consolidate_safetensors_async(tmp_path, "img01") + consolidate_safetensors_async(tmp_path, _RP, "img01") shutdown_io_pool() - assert (tmp_path / "img01.safetensors").exists() + assert safetensors_path(tmp_path, _RP, "img01").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_depth(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True) + save_edges(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True) + save_chmv2(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True) save_segmentation( torch.randint(0, 11, (1, H, W), dtype=torch.uint8), - tmp_path, "img04", save_npy=False, save_vis=True, num_classes=11, + tmp_path, _RP, "img04", save_npy=False, save_vis=True, num_classes=11, ) - ok = consolidate_safetensors(tmp_path, "img04") + ok = consolidate_safetensors(tmp_path, _RP, "img04") assert ok - data = st_load_file(tmp_path / "img04.safetensors") + data = st_load_file(safetensors_path(tmp_path, _RP, "img04")) 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) - save_depth(depth, tmp_path, "img03", save_npy=True, save_vis=False) - consolidate_safetensors(tmp_path, "img03") - data = st_load_file(tmp_path / "img03.safetensors") - # float16 round-trip: compare at fp16 precision + save_depth(depth, tmp_path, _RP, "img03", save_npy=True, save_vis=False) + consolidate_safetensors(tmp_path, _RP, "img03") + data = st_load_file(safetensors_path(tmp_path, _RP, "img03")) expected = depth.half() torch.testing.assert_close(data["depth"], expected, atol=0, rtol=0) diff --git a/src/tests/test_pipeline_integration.py b/src/tests/test_pipeline_integration.py index fc0deb9..3adb2ff 100644 --- a/src/tests/test_pipeline_integration.py +++ b/src/tests/test_pipeline_integration.py @@ -16,7 +16,7 @@ from src.augmentor.dataset import ( discover_images, ) from src.augmentor.inference import compute_edges_from_depth -from src.augmentor.io_utils import shutdown_io_pool +from src.augmentor.io_utils import npy_path, vis_path, shutdown_io_pool from src.conf.hardware_conf import HardwareConfig from src.conf.input_conf import InputConfig from src.conf.models_conf import ModelsConfig @@ -103,9 +103,6 @@ class TestDepthStageIntegration: ) -> None: from src.main import run_depth_stage - for r in sample_records: - r.output_dir.mkdir(parents=True, exist_ok=True) - mock_model = _make_mock_depth_model_da3(input_conf.image_size, input_conf.image_size) with patch("src.main.load_depth_model", return_value=mock_model), \ @@ -116,9 +113,9 @@ class TestDepthStageIntegration: ) for r in sample_records: - npy_path = r.output_dir / f"{r.stem}_depth.npy" - assert npy_path.exists(), f"Missing {r.stem}_depth.npy in {r.output_dir}" - arr = np.load(npy_path) + p = npy_path(r.output_root, "depth", r.rel_parent, r.stem) + assert p.exists(), f"Missing depth npy for {r.stem}" + arr = np.load(p) assert arr.dtype == np.float16 assert arr.shape[0] == 1 @@ -133,9 +130,6 @@ class TestChmv2StageIntegration: ) -> None: from src.main import run_chmv2_stage - for r in sample_records: - r.output_dir.mkdir(parents=True, exist_ok=True) - mock_model, mock_processor = _make_mock_chmv2( input_conf.image_size, input_conf.image_size, ) @@ -148,9 +142,9 @@ class TestChmv2StageIntegration: ) for r in sample_records: - npy_path = r.output_dir / f"{r.stem}_chm.npy" - assert npy_path.exists(), f"Missing {r.stem}_chm.npy in {r.output_dir}" - arr = np.load(npy_path) + p = npy_path(r.output_root, "chm", r.rel_parent, r.stem) + assert p.exists(), f"Missing chm npy for {r.stem}" + arr = np.load(p) assert arr.dtype == np.float16 assert arr.shape[0] == 1 @@ -164,18 +158,16 @@ class TestEdgesStageIntegration: from src.main import run_edges_stage for r in sample_records: - r.output_dir.mkdir(parents=True, exist_ok=True) - np.save( - r.output_dir / f"{r.stem}_depth.npy", - np.random.rand(1, 64, 64).astype(np.float16), - ) + p = npy_path(r.output_root, "depth", r.rel_parent, r.stem) + p.parent.mkdir(parents=True, exist_ok=True) + np.save(p, np.random.rand(1, 64, 64).astype(np.float16)) run_edges_stage(sample_records, pipeline_conf) for r in sample_records: - npy_path = r.output_dir / f"{r.stem}_edge.npy" - assert npy_path.exists(), f"Missing {r.stem}_edge.npy in {r.output_dir}" - arr = np.load(npy_path) + p = npy_path(r.output_root, "edge", r.rel_parent, r.stem) + assert p.exists(), f"Missing edge npy for {r.stem}" + arr = np.load(p) assert arr.dtype == np.float16 def test_skips_missing_depth( @@ -185,11 +177,9 @@ class TestEdgesStageIntegration: ) -> None: from src.main import run_edges_stage - for r in sample_records: - r.output_dir.mkdir(parents=True, exist_ok=True) run_edges_stage(sample_records, pipeline_conf) for r in sample_records: - assert not (r.output_dir / f"{r.stem}_edge.npy").exists() + assert not npy_path(r.output_root, "edge", r.rel_parent, r.stem).exists() class TestSegmentationStageIntegration: @@ -203,9 +193,6 @@ class TestSegmentationStageIntegration: ) -> None: from src.main import run_segmentation_stage - for r in sample_records: - r.output_dir.mkdir(parents=True, exist_ok=True) - num_classes = seg_conf.num_classes mock_model, seg_config_dict = _make_mock_segformer( num_classes, input_conf.image_size, input_conf.image_size, @@ -221,9 +208,9 @@ class TestSegmentationStageIntegration: ) for r in sample_records: - npy_path = r.output_dir / f"{r.stem}_segm.npy" - assert npy_path.exists() - arr = np.load(npy_path) + p = npy_path(r.output_root, "segm", r.rel_parent, r.stem) + assert p.exists() + arr = np.load(p) assert arr.dtype == np.uint8 @@ -275,10 +262,11 @@ class TestFullPipelineSmoke: output_root = Path(pipeline_conf.output_root) assert (output_root / "manifest.json").exists() - found_depth = list(output_root.rglob("*_depth.npy")) - found_edges = list(output_root.rglob("*_edge.npy")) - found_seg = list(output_root.rglob("*_segm.npy")) - found_chmv2 = list(output_root.rglob("*_chm.npy")) + # New dir layout: npy/depth/..., npy/edge/..., etc. + found_depth = list((output_root / "npy" / "depth").rglob("*.npy")) + found_edges = list((output_root / "npy" / "edge").rglob("*.npy")) + found_seg = list((output_root / "npy" / "segm").rglob("*.npy")) + found_chmv2 = list((output_root / "npy" / "chm").rglob("*.npy")) assert len(found_depth) > 0 assert len(found_edges) > 0 assert len(found_seg) > 0 @@ -297,6 +285,7 @@ class TestFullPipelineVisOnly: output_root=str(tmp_path / "output"), stages=["depth", "edges", "segmentation", "chmv2"], save_npy=False, save_vis=True, + save_safetensors=False, save_concat=False, resume=False, log_level="WARNING", ) hw_conf = HardwareConfig( @@ -321,8 +310,8 @@ class TestFullPipelineVisOnly: run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf) output_root = Path(pipeline_conf.output_root) - assert len(list(output_root.rglob("*_depth.png"))) > 0 - assert len(list(output_root.rglob("*_edge.png"))) > 0 - assert len(list(output_root.rglob("*_segm.png"))) > 0 - assert len(list(output_root.rglob("*_chm.png"))) > 0 + assert len(list((output_root / "depth").rglob("*.png"))) > 0 + assert len(list((output_root / "edge").rglob("*.png"))) > 0 + assert len(list((output_root / "segm").rglob("*.png"))) > 0 + assert len(list((output_root / "chm").rglob("*.png"))) > 0 assert len(list(output_root.rglob("*.npy"))) == 0