Pikaliov/depth_edges_annotate_worlduav
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor output to directory-based layout + migration script

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-17 17:11:01 +03:00
parent 892a2574f6
commit 13ff079891
8 changed files with 502 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
README.md
View File

@@ -58,7 +58,9 @@ python -m pytest src/tests/ -v
│ │ └── dataset.py # Discovery, filtering, PyTorch Dataset │ │ └── dataset.py # Discovery, filtering, PyTorch Dataset
│ ├── conf/ # Gin-configurable dataclasses │ ├── conf/ # Gin-configurable dataclasses
│ ├── utils/ # Profiler, benchmark, GPU utils │ ├── utils/ # Profiler, benchmark, GPU utils
│ └── tests/ # 141 тест (pytest) │ └── tests/ # 143 теста (pytest)
├── scripts/
│ └── migrate_layout.py # Миграция со старого prefix-формата
└── docs/ └── docs/
├── segmentation_class_analysis.md # Анализ классов сегментации (11 классов) ├── segmentation_class_analysis.md # Анализ классов сегментации (11 классов)
├── segearth_ov3_architecture.md # Архитектура SegEarth-OV3 + SAM 3.1 ├── 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) 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/ World-UAV-aug/
├── Rot/SouthernSuburbs/DB/img/ ├── depth/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis
│ ├── crop_12_4.safetensors # ВСЕ модальности (для обучения, zero-copy mmap) ├── edge/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis
│ ├── crop_12_4_depth.png # grayscale визуализация ├── segm/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis (palette mode P)
│ ├── crop_12_4_edge.png # grayscale визуализация ├── chm/Rot/SouthernSuburbs/DB/img/crop_12_4.png # vis
│ ├── crop_12_4_segm.png # RGB palette визуализация (11 классов) ├── npy/depth/Rot/SouthernSuburbs/DB/img/crop_12_4.npy # float16 intermediate
│ └── crop_12_4_chm.png # grayscale визуализация ├── npy/edge/...
├── Country/... ├── npy/segm/...
└── Terrain/... ├── npy/chm/...
├── safetensors/Rot/SouthernSuburbs/DB/img/crop_12_4.safetensors # для обучения
└── manifest.json
``` ```
### SafeTensors (рекомендуемый формат для обучения) ### SafeTensors (рекомендуемый формат для обучения)
@@ -213,10 +217,11 @@ World-UAV-aug/
from safetensors.torch import load_file from safetensors.torch import load_file
stem = "crop_12_4" 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 # 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] depth = data["depth"] # [1, 256, 256] float16, [0, 1]
edge = data["edge"] # [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 import numpy as np
# Depth / Edge / CHM -- grayscale float [0, 1] # Depth / Edge / CHM -- grayscale float [0, 1]
depth = np.array(Image.open(aug_dir / f"{stem}_depth.png")) / 255.0 # [H, W] depth = np.array(Image.open(aug_root / "depth" / rel_parent / f"{stem}.png")) / 255.0
edge = np.array(Image.open(aug_dir / f"{stem}_edge.png")) / 255.0 edge = np.array(Image.open(aug_root / "edge" / rel_parent / f"{stem}.png")) / 255.0
chm = np.array(Image.open(aug_dir / f"{stem}_chm.png")) / 255.0 chm = np.array(Image.open(aug_root / "chm" / rel_parent / f"{stem}.png")) / 255.0
``` ```
> PNG визуализации квантуют float16 в uint8 (256 уровней). Для обучения используйте SafeTensors. > 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 суммарно): Веса скачиваются один раз в `in/weights/` (~10 GB суммарно):

153
scripts/migrate_layout.py Normal file
View File

@@ -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()

46
src/augmentor/dataset.py
View File

@@ -11,6 +11,8 @@ from PIL import Image
from torch.utils.data import Dataset from torch.utils.data import Dataset
from torchvision import transforms from torchvision import transforms
from src.augmentor.io_utils import npy_path, vis_path, safetensors_path
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp"} EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp"}
@@ -33,7 +35,8 @@ class ImageRecord(NamedTuple):
abs_path: Path abs_path: Path
rel_path: str rel_path: str
stem: str stem: str
output_dir: Path output_root: Path
rel_parent: str
def is_query_record(record: ImageRecord) -> bool: def is_query_record(record: ImageRecord) -> bool:
@@ -104,7 +107,11 @@ def discover_images(
rel = p.relative_to(root) rel = p.relative_to(root)
records.append(ImageRecord( 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: if n_skipped_incomplete > 0:
@@ -119,17 +126,12 @@ def attach_output_dirs(
records: list[ImageRecord], records: list[ImageRecord],
output_root: Path, output_root: Path,
) -> list[ImageRecord]: ) -> list[ImageRecord]:
"""Set output_dir for each record: output_root / <parent dirs>/.""" """Set output_root for each record."""
out: list[ImageRecord] = [] return [r._replace(output_root=output_root) for r in records]
for r in records:
rel = Path(r.rel_path)
odir = output_root / rel.parent
out.append(r._replace(output_dir=odir))
return out
# Suffix appended to stem for each stage: {stem}_{suffix}.npy # Modality name for each stage (used for folder names).
STAGE_SUFFIX: dict[str, str] = { STAGE_MODALITY: dict[str, str] = {
"depth": "depth", "depth": "depth",
"edges": "edge", "edges": "edge",
"segmentation": "segm", "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( def filter_completed(
records: list[ImageRecord], records: list[ImageRecord],
stage: str, stage: str,
@@ -150,16 +146,15 @@ def filter_completed(
"""Return (pending_records, n_skipped) for a given stage.""" """Return (pending_records, n_skipped) for a given stage."""
if stage == "consolidate": if stage == "consolidate":
return filter_consolidated(records) return filter_consolidated(records)
suffix = STAGE_SUFFIX.get(stage) modality = STAGE_MODALITY.get(stage)
if suffix is None: if modality is None:
return records, 0 return records, 0
pending: list[ImageRecord] = [] pending: list[ImageRecord] = []
skipped = 0 skipped = 0
for r in records: for r in records:
# Check both .npy and .png — either means the stage is done. np_p = npy_path(r.output_root, modality, r.rel_parent, r.stem)
npy = r.output_dir / f"{r.stem}_{suffix}.npy" vis_p = vis_path(r.output_root, modality, r.rel_parent, r.stem)
png = r.output_dir / f"{r.stem}_{suffix}.png" if np_p.exists() or vis_p.exists():
if npy.exists() or png.exists():
skipped += 1 skipped += 1
else: else:
pending.append(r) pending.append(r)
@@ -173,7 +168,7 @@ def filter_consolidated(
pending: list[ImageRecord] = [] pending: list[ImageRecord] = []
skipped = 0 skipped = 0
for r in records: 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(): if st.exists():
skipped += 1 skipped += 1
else: else:
@@ -212,5 +207,6 @@ class AugmentDataset(Dataset):
"image_raw": tensor, "image_raw": tensor,
"rel_path": r.rel_path, "rel_path": r.rel_path,
"stem": r.stem, "stem": r.stem,
"output_dir": str(r.output_dir), "output_root": str(r.output_root),
"rel_parent": r.rel_parent,
} }

217
src/augmentor/io_utils.py
View File

@@ -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. No global config imports — all parameters passed explicitly.
""" """
@@ -45,6 +58,29 @@ def shutdown_io_pool() -> None:
_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. # Intuitive RS segmentation palette: index → RGB.
_FIXED_PALETTE = np.array([ _FIXED_PALETTE = np.array([
[0, 0, 0], # 0: background — black [0, 0, 0], # 0: background — black
@@ -75,6 +111,10 @@ def make_palette(num_classes: int, seed: int = 42) -> np.ndarray:
return palette return palette
# ---------------------------------------------------------------------------
# Low-level atomic save
# ---------------------------------------------------------------------------
def _atomic_save_npy(arr: np.ndarray, path: Path) -> None: def _atomic_save_npy(arr: np.ndarray, path: Path) -> None:
"""Write .npy atomically via temp file + rename.""" """Write .npy atomically via temp file + rename."""
path.parent.mkdir(parents=True, exist_ok=True) 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] return lut[idx]
# ---------------------------------------------------------------------------
# Save float16 maps (depth, edge, chm)
# ---------------------------------------------------------------------------
def _save_float16_map( def _save_float16_map(
data: torch.Tensor, data: torch.Tensor,
output_dir: Path, output_root: Path,
rel_parent: str,
stem: str, stem: str,
suffix: str, modality: str,
save_npy: bool = True, save_npy: bool = True,
save_vis: bool = True, save_vis: bool = True,
colormap: str | None = None, colormap: str | None = None,
) -> None: ) -> None:
"""Save a [1, H, W] float tensor as {stem}_{suffix}.npy (float16) + optional vis. """Save a [1, H, W] float tensor as .npy (float16) + optional vis .png."""
Args:
colormap: If set (e.g. "turbo"), apply colormap for RGB visualization.
If None, save grayscale.
"""
arr = data.half().numpy() arr = data.half().numpy()
if save_npy: 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: if save_vis:
gray = arr.squeeze(0).astype(np.float32) gray = arr.squeeze(0).astype(np.float32)
if colormap: if colormap:
vis = _apply_colormap(gray, colormap) vis = _apply_colormap(gray, colormap)
else: else:
vis = (gray * 255).clip(0, 255).astype(np.uint8) 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, def save_depth(depth: torch.Tensor, output_root: Path, rel_parent: str,
save_npy: bool = True, save_vis: bool = True) -> None: stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
_save_float16_map(depth, output_dir, stem, "depth", save_npy, save_vis) _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, def save_depth_async(depth: torch.Tensor, output_root: Path, rel_parent: str,
save_npy: bool = True, save_vis: bool = True) -> None: 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) 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, def save_chmv2(depth: torch.Tensor, output_root: Path, rel_parent: str,
save_npy: bool = True, save_vis: bool = True) -> None: stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
_save_float16_map(depth, output_dir, stem, "chm", save_npy, save_vis) _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, def save_chmv2_async(depth: torch.Tensor, output_root: Path, rel_parent: str,
save_npy: bool = True, save_vis: bool = True) -> None: 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) 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, def save_edges(edges: torch.Tensor, output_root: Path, rel_parent: str,
save_npy: bool = True, save_vis: bool = True) -> None: stem: str, save_npy: bool = True, save_vis: bool = True) -> None:
_save_float16_map(edges, output_dir, stem, "edge", save_npy, save_vis) _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, def save_edges_async(edges: torch.Tensor, output_root: Path, rel_parent: str,
save_npy: bool = True, save_vis: bool = True) -> None: 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) get_io_pool().submit(save_edges, edges.clone().cpu(), output_root, rel_parent,
stem, save_npy, save_vis)
# ---------------------------------------------------------------------------
# Save segmentation
# ---------------------------------------------------------------------------
def save_segmentation( def save_segmentation(
seg_ids: torch.Tensor, seg_ids: torch.Tensor,
output_dir: Path, output_root: Path,
rel_parent: str,
stem: str, stem: str,
save_npy: bool = True, save_npy: bool = True,
save_vis: bool = True, save_vis: bool = True,
num_classes: int = 150, num_classes: int = 150,
) -> None: ) -> None:
"""Save segmentation map [1, H, W] uint8 as {stem}_segm.npy.""" """Save segmentation map [1, H, W] uint8."""
arr = seg_ids.byte().numpy() arr = seg_ids.byte().numpy()
if save_npy: 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: if save_vis:
palette = make_palette(num_classes) palette = make_palette(num_classes)
seg_np = arr.squeeze(0).astype(np.uint8) seg_np = arr.squeeze(0).astype(np.uint8)
seg_clamped = np.clip(seg_np, 0, num_classes - 1).astype(np.uint8) seg_clamped = np.clip(seg_np, 0, num_classes - 1).astype(np.uint8)
img = Image.fromarray(seg_clamped).convert("P") 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 = np.zeros(768, dtype=np.uint8)
flat_pal[: num_classes * 3] = palette.flatten() flat_pal[: num_classes * 3] = palette.flatten()
img.putpalette(flat_pal.tolist()) 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( def save_segmentation_async(
seg_ids: torch.Tensor, seg_ids: torch.Tensor,
output_dir: Path, output_root: Path,
rel_parent: str,
stem: str, stem: str,
save_npy: bool = True, save_npy: bool = True,
save_vis: bool = True, save_vis: bool = True,
num_classes: int = 150, num_classes: int = 150,
) -> None: ) -> None:
get_io_pool().submit( get_io_pool().submit(
save_segmentation, seg_ids.clone().cpu(), output_dir, stem, save_segmentation, seg_ids.clone().cpu(), output_root, rel_parent,
save_npy, save_vis, num_classes, 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 # 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]] = { _MODALITY_SPEC: dict[str, tuple[torch.dtype, str]] = {
"depth": (torch.float16, "depth"), "depth": (torch.float16, "depth"),
"edge": (torch.float16, "edge"), "edge": (torch.float16, "edge"),
@@ -237,41 +265,38 @@ _MODALITY_SPEC: dict[str, tuple[torch.dtype, str]] = {
def _load_modality_tensor( 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: ) -> torch.Tensor | None:
"""Load a single modality from .npy or .png, return [1, H, W] tensor or None.""" """Load a single modality from .npy or .png, return [1, H, W] tensor or None."""
npy_path = output_dir / f"{stem}_{suffix}.npy" np_p = npy_path(output_root, modality, rel_parent, stem)
png_path = output_dir / f"{stem}_{suffix}.png" vis_p = vis_path(output_root, modality, rel_parent, stem)
if npy_path.exists(): if np_p.exists():
arr = np.load(npy_path) arr = np.load(np_p)
t = torch.from_numpy(arr.astype(np.float32 if dtype != torch.uint8 else np.uint8)) t = torch.from_numpy(arr.astype(np.float32 if dtype != torch.uint8 else np.uint8))
if t.ndim == 2: if t.ndim == 2:
t = t.unsqueeze(0) t = t.unsqueeze(0)
return t.to(dtype) return t.to(dtype)
if png_path.exists(): if vis_p.exists():
img = np.array(Image.open(png_path)) if modality == "segm":
if suffix == "segm": pil = Image.open(vis_p)
# 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 pil.mode == "P": if pil.mode == "P":
img = np.array(pil) img = np.array(pil)
else: else:
# RGB palette render — can't recover class IDs reliably, skip. logger.debug("Skipping %s segm.png (RGB, no class IDs).", stem)
logger.debug("Skipping %s_%s.png (RGB palette, no class IDs).", stem, suffix)
return None return None
t = torch.from_numpy(img.astype(np.uint8)) t = torch.from_numpy(img.astype(np.uint8))
if t.ndim == 2: if t.ndim == 2:
t = t.unsqueeze(0) t = t.unsqueeze(0)
return t return t
else: else:
img = np.array(Image.open(vis_p))
arr = img.astype(np.float32) / 255.0 arr = img.astype(np.float32) / 255.0
if arr.ndim == 2: if arr.ndim == 2:
arr = arr[np.newaxis] arr = arr[np.newaxis]
elif arr.ndim == 3: elif arr.ndim == 3:
# Grayscale saved as RGB — take first channel.
arr = arr[:, :, 0:1].transpose(2, 0, 1) arr = arr[:, :, 0:1].transpose(2, 0, 1)
return torch.from_numpy(arr).to(dtype) return torch.from_numpy(arr).to(dtype)
@@ -279,57 +304,63 @@ def _load_modality_tensor(
def consolidate_safetensors( def consolidate_safetensors(
output_dir: Path, output_root: Path,
rel_parent: str,
stem: str, stem: str,
cleanup_npy: bool = False, cleanup_npy: bool = False,
) -> bool: ) -> 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. Returns True if the file was written, False if no modalities found.
""" """
tensors: dict[str, torch.Tensor] = {} tensors: dict[str, torch.Tensor] = {}
npy_paths: list[Path] = [] npy_paths_to_clean: list[Path] = []
for suffix, (dtype, _) in _MODALITY_SPEC.items(): for modality, (dtype, _) in _MODALITY_SPEC.items():
t = _load_modality_tensor(output_dir, stem, suffix, dtype) t = _load_modality_tensor(output_root, rel_parent, stem, modality, dtype)
if t is not None: if t is not None:
tensors[suffix] = t tensors[modality] = t
npy_path = output_dir / f"{stem}_{suffix}.npy" np_p = npy_path(output_root, modality, rel_parent, stem)
if npy_path.exists(): if np_p.exists():
npy_paths.append(npy_path) npy_paths_to_clean.append(np_p)
if not tensors: if not tensors:
return False return False
st_path = output_dir / f"{stem}.safetensors" st_p = safetensors_path(output_root, rel_parent, stem)
output_dir.mkdir(parents=True, exist_ok=True) st_p.parent.mkdir(parents=True, exist_ok=True)
# Atomic write via temp file. fd, tmp = tempfile.mkstemp(suffix=".safetensors", dir=st_p.parent)
fd, tmp = tempfile.mkstemp(suffix=".safetensors", dir=output_dir)
os.close(fd) os.close(fd)
try: try:
_st_save_file(tensors, tmp) _st_save_file(tensors, tmp)
os.replace(tmp, st_path) os.replace(tmp, st_p)
except BaseException: except BaseException:
if os.path.exists(tmp): if os.path.exists(tmp):
os.remove(tmp) os.remove(tmp)
raise raise
if cleanup_npy: if cleanup_npy:
for p in npy_paths: for p in npy_paths_to_clean:
p.unlink(missing_ok=True) p.unlink(missing_ok=True)
return True return True
def consolidate_safetensors_async( def consolidate_safetensors_async(
output_dir: Path, output_root: Path,
rel_parent: str,
stem: str, stem: str,
cleanup_npy: bool = False, cleanup_npy: bool = False,
) -> None: ) -> 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: def setup_logging(log_level: str = "INFO", log_file: Path | None = None) -> None:
"""Configure root logger with coloredlogs for console + optional file handler.""" """Configure root logger with coloredlogs for console + optional file handler."""
import coloredlogs import coloredlogs

65
src/main.py
View File

@@ -39,6 +39,7 @@ from src.augmentor.inference import (
infer_segmentation_batch, infer_segmentation_batch,
) )
from src.augmentor.io_utils import ( from src.augmentor.io_utils import (
npy_path, vis_path,
save_depth_async, save_chmv2_async, save_edges_async, save_depth_async, save_chmv2_async, save_edges_async,
save_segmentation_async, consolidate_safetensors, save_segmentation_async, consolidate_safetensors,
setup_logging, shutdown_io_pool, setup_logging, shutdown_io_pool,
@@ -57,7 +58,6 @@ _STAGE_EMOJI = {
"edges": "🔪", "edges": "🔪",
"segmentation": "🗺️", "segmentation": "🗺️",
"chmv2": "🦕", "chmv2": "🦕",
"concat": "🧩",
"consolidate": "📦", "consolidate": "📦",
} }
@@ -97,11 +97,7 @@ def _resolve_image_sizes(
records: list[ImageRecord], records: list[ImageRecord],
input_conf: InputConfig, input_conf: InputConfig,
) -> list[tuple[list[ImageRecord], int, str]]: ) -> list[tuple[list[ImageRecord], int, str]]:
"""Split records into groups by target resolution. """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).
"""
if input_conf.query_image_size == input_conf.image_size: if input_conf.query_image_size == input_conf.image_size:
return [(records, input_conf.image_size, "all")] return [(records, input_conf.image_size, "all")]
db_recs, query_recs = split_by_view(records) db_recs, query_recs = split_by_view(records)
@@ -149,7 +145,9 @@ def run_depth_stage(
for batch in pbar: for batch in pbar:
depths = infer_depth_batch(model, batch["image_raw"], device) depths = infer_depth_batch(model, batch["image_raw"], device)
for i in range(depths.shape[0]): 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], stem=batch["stem"][i],
save_npy=pipeline_conf.save_npy, save_npy=pipeline_conf.save_npy,
save_vis=pipeline_conf.save_vis) save_vis=pipeline_conf.save_vis)
@@ -169,9 +167,9 @@ def run_edges_stage(
"""🔪 Compute Sobel edges from saved depth (CPU, batched).""" """🔪 Compute Sobel edges from saved depth (CPU, batched)."""
valid: list[ImageRecord] = [] valid: list[ImageRecord] = []
for r in records: for r in records:
depth_png = r.output_dir / f"{r.stem}_depth.png" np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
depth_npy = r.output_dir / f"{r.stem}_depth.npy" vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
if depth_png.exists() or depth_npy.exists(): if np_p.exists() or vis_p.exists():
valid.append(r) valid.append(r)
else: else:
logger.warning("⚠️ No depth for %s, skipping edges.", r.rel_path) 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] chunk = valid[start : start + batch_size]
depth_tensors = [] depth_tensors = []
for r in chunk: for r in chunk:
npy_path = r.output_dir / f"{r.stem}_depth.npy" np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
png_path = r.output_dir / f"{r.stem}_depth.png" vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
if npy_path.exists(): if np_p.exists():
d = np.load(npy_path).astype(np.float32) d = np.load(np_p).astype(np.float32)
else: else:
from PIL import Image 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: if d.ndim == 2:
d = d[np.newaxis] d = d[np.newaxis]
depth_tensors.append(torch.from_numpy(d)) depth_tensors.append(torch.from_numpy(d))
@@ -200,7 +198,8 @@ def run_edges_stage(
depths = depths.unsqueeze(1) depths = depths.unsqueeze(1)
edges_batch = compute_edges_from_depth(depths) edges_batch = compute_edges_from_depth(depths)
for j, r in enumerate(chunk): 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_npy=pipeline_conf.save_npy,
save_vis=pipeline_conf.save_vis) save_vis=pipeline_conf.save_vis)
processed += len(chunk) processed += len(chunk)
@@ -245,7 +244,9 @@ def run_chmv2_stage(
for batch in pbar: for batch in pbar:
depths = infer_chmv2_batch(model, processor, batch["image_raw"], device) depths = infer_chmv2_batch(model, processor, batch["image_raw"], device)
for i in range(depths.shape[0]): 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], stem=batch["stem"][i],
save_npy=pipeline_conf.save_npy, save_npy=pipeline_conf.save_npy,
save_vis=pipeline_conf.save_vis) save_vis=pipeline_conf.save_vis)
@@ -303,7 +304,9 @@ def run_segmentation_stage(
) )
for j in range(segs.shape[0]): for j in range(segs.shape[0]):
save_segmentation_async( 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, save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis,
num_classes=num_classes, num_classes=num_classes,
) )
@@ -326,7 +329,8 @@ def run_consolidate_stage(
colour="magenta", colour="magenta",
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]") bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]")
for r in pbar: 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: if ok:
written += 1 written += 1
pbar.set_postfix(written=f"{written}/{total}") 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.") logger.error("❌ No images found. Check input_root in pipeline.gin.")
return 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 # When save_safetensors is on, always save intermediate .npy
# (consolidation reads .npy for lossless float16; PNG is lossy uint8). # (consolidation reads .npy for lossless float16; PNG is lossy uint8).
if pipeline_conf.save_safetensors and not pipeline_conf.save_npy: if pipeline_conf.save_safetensors and not pipeline_conf.save_npy:
@@ -457,7 +452,7 @@ def run_pipeline(
# Manifest. # Manifest.
manifest = { manifest = {
"pipeline_version": "3.3.0-safetensors", "pipeline_version": "4.0.0-dir-layout",
"image_size_db": input_conf.image_size, "image_size_db": input_conf.image_size,
"image_size_query": input_conf.query_image_size, "image_size_query": input_conf.query_image_size,
"profile": hw_conf.profile_name, "profile": hw_conf.profile_name,
@@ -501,16 +496,7 @@ def run_pipeline(
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def main() -> None: def main() -> None:
"""Load all gin configs and run the augmentation pipeline. """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'"
"""
import argparse import argparse
parser = argparse.ArgumentParser(description="Augmentation pipeline") parser = argparse.ArgumentParser(description="Augmentation pipeline")
@@ -523,7 +509,6 @@ def main() -> None:
proj_dir = get_proj_dir() proj_dir = get_proj_dir()
path2cfg = f"{proj_dir}in/config_files/" path2cfg = f"{proj_dir}in/config_files/"
# Load configs with optional CLI overrides.
if args.gin: if args.gin:
import gin as _gin import gin as _gin
cfg_dir = Path(path2cfg) cfg_dir = Path(path2cfg)

17
src/tests/test_dataset.py
View File

@@ -16,6 +16,7 @@ from src.augmentor.dataset import (
filter_completed, filter_completed,
INCOMPLETE_SCENES, INCOMPLETE_SCENES,
) )
from src.augmentor.io_utils import npy_path
class TestDiscoverImages: class TestDiscoverImages:
@@ -42,9 +43,7 @@ class TestDiscoverImages:
(tmp_path / "good.png").write_bytes( (tmp_path / "good.png").write_bytes(
Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).tobytes() 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) 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) assert all(r.abs_path.suffix in {".png", ".jpg", ".jpeg", ".bmp"} for r in records)
def test_subset_filter(self, tmp_path: Path) -> None: def test_subset_filter(self, tmp_path: Path) -> None:
@@ -65,14 +64,13 @@ class TestDiscoverImages:
class TestAttachOutputDirs: 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) records = discover_images(fake_image_dir)
out_root = tmp_path / "output" out_root = tmp_path / "output"
attached = attach_output_dirs(records, out_root) attached = attach_output_dirs(records, out_root)
for r in attached: for r in attached:
assert str(r.output_dir).startswith(str(out_root)) assert r.output_root == out_root
# output_dir is the parent directory (same structure, no per-image subfolder) assert r.rel_parent == str(Path(r.rel_path).parent)
assert r.output_dir == out_root / Path(r.rel_path).parent
class TestFilterCompleted: class TestFilterCompleted:
@@ -83,8 +81,9 @@ class TestFilterCompleted:
def test_skips_completed(self, sample_records: list[ImageRecord]) -> None: def test_skips_completed(self, sample_records: list[ImageRecord]) -> None:
r = sample_records[0] r = sample_records[0]
r.output_dir.mkdir(parents=True, exist_ok=True) p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
np.save(r.output_dir / f"{r.stem}_depth.npy", np.zeros((1, 8, 8))) p.parent.mkdir(parents=True, exist_ok=True)
np.save(p, np.zeros((1, 8, 8)))
pending, skipped = filter_completed(sample_records, "depth") pending, skipped = filter_completed(sample_records, "depth")
assert skipped == 1 assert skipped == 1
assert len(pending) == len(sample_records) - 1 assert len(pending) == len(sample_records) - 1
@@ -111,4 +110,4 @@ class TestAugmentDataset:
def test_getitem_keys(self, sample_records: list[ImageRecord]) -> None: def test_getitem_keys(self, sample_records: list[ImageRecord]) -> None:
ds = AugmentDataset(sample_records, image_size=32) ds = AugmentDataset(sample_records, image_size=32)
item = ds[0] 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"}

212
src/tests/test_io_utils.py
View File

@@ -11,6 +11,7 @@ from safetensors.torch import load_file as st_load_file
from src.augmentor.io_utils import ( from src.augmentor.io_utils import (
make_palette, make_palette,
npy_path, vis_path, safetensors_path,
save_chmv2, save_chmv2,
save_chmv2_async, save_chmv2_async,
save_depth, save_depth,
@@ -19,7 +20,6 @@ from src.augmentor.io_utils import (
save_edges_async, save_edges_async,
save_segmentation, save_segmentation,
save_segmentation_async, save_segmentation_async,
save_concat_6ch,
consolidate_safetensors, consolidate_safetensors,
consolidate_safetensors_async, consolidate_safetensors_async,
shutdown_io_pool, shutdown_io_pool,
@@ -53,126 +53,143 @@ class TestAtomicSaveNpy:
assert len(tmp_files) == 0 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 # save_depth
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
_RP = "sub" # rel_parent for tests
class TestSaveDepth: class TestSaveDepth:
def test_saves_float16_npy(self, tmp_path: Path) -> None: 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) save_depth(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
arr = np.load(tmp_path / "img01_depth.npy") 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.dtype == np.float16
assert arr.shape == (1, 32, 32) assert arr.shape == (1, 32, 32)
def test_saves_vis_png(self, tmp_path: Path) -> None: 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) save_depth(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
assert (tmp_path / "img01_depth.png").exists() save_npy=False, save_vis=True)
assert not (tmp_path / "img01_depth.npy").exists() 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: def test_npy_values_in_range(self, tmp_path: Path) -> None:
save_depth(torch.rand(1, 16, 16), tmp_path, "img01") save_depth(torch.rand(1, 16, 16), tmp_path, _RP, "img01")
arr = np.load(tmp_path / "img01_depth.npy").astype(np.float32) p = npy_path(tmp_path, "depth", _RP, "img01")
arr = np.load(p).astype(np.float32)
assert arr.min() >= 0.0 assert arr.min() >= 0.0
assert arr.max() <= 1.0 assert arr.max() <= 1.0
class TestSaveChmv2: class TestSaveChmv2:
def test_saves_float16_npy(self, tmp_path: Path) -> None: 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) save_chmv2(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
arr = np.load(tmp_path / "img01_chm.npy") save_npy=True, save_vis=False)
arr = np.load(npy_path(tmp_path, "chm", _RP, "img01"))
assert arr.dtype == np.float16 assert arr.dtype == np.float16
assert arr.shape == (1, 32, 32) assert arr.shape == (1, 32, 32)
def test_saves_vis_png(self, tmp_path: Path) -> None: 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) save_chmv2(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
assert (tmp_path / "img01_chm.png").exists() save_npy=False, save_vis=True)
assert not (tmp_path / "img01_chm.npy").exists() 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: def test_npy_values_in_range(self, tmp_path: Path) -> None:
save_chmv2(torch.rand(1, 16, 16), tmp_path, "img01") save_chmv2(torch.rand(1, 16, 16), tmp_path, _RP, "img01")
arr = np.load(tmp_path / "img01_chm.npy").astype(np.float32) arr = np.load(npy_path(tmp_path, "chm", _RP, "img01")).astype(np.float32)
assert arr.min() >= 0.0 assert arr.min() >= 0.0
assert arr.max() <= 1.0 assert arr.max() <= 1.0
class TestSaveEdges: class TestSaveEdges:
def test_saves_float16_npy(self, tmp_path: Path) -> None: 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) save_edges(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
arr = np.load(tmp_path / "img01_edge.npy") save_npy=True, save_vis=False)
arr = np.load(npy_path(tmp_path, "edge", _RP, "img01"))
assert arr.dtype == np.float16 assert arr.dtype == np.float16
def test_saves_vis(self, tmp_path: Path) -> None: 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) save_edges(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
assert (tmp_path / "img01_edge.png").exists() save_npy=False, save_vis=True)
assert vis_path(tmp_path, "edge", _RP, "img01").exists()
class TestSaveSegmentation: class TestSaveSegmentation:
def test_saves_uint8_npy(self, tmp_path: Path) -> None: def test_saves_uint8_npy(self, tmp_path: Path) -> None:
seg = torch.randint(0, 5, (1, 32, 32), dtype=torch.uint8) 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) save_segmentation(seg, tmp_path, _RP, "img01",
arr = np.load(tmp_path / "img01_segm.npy") 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.dtype == np.uint8
assert arr.shape == (1, 32, 32) assert arr.shape == (1, 32, 32)
def test_saves_vis(self, tmp_path: Path) -> None: def test_saves_vis(self, tmp_path: Path) -> None:
seg = torch.randint(0, 3, (1, 32, 32), dtype=torch.uint8) 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) save_segmentation(seg, tmp_path, _RP, "img01",
assert (tmp_path / "img01_segm.png").exists() 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: def test_int_tensor_also_works(self, tmp_path: Path) -> None:
seg = torch.randint(0, 5, (1, 16, 16), dtype=torch.int64) seg = torch.randint(0, 5, (1, 16, 16), dtype=torch.int64)
save_segmentation(seg, tmp_path, "img01", num_classes=5) save_segmentation(seg, tmp_path, _RP, "img01", num_classes=5)
arr = np.load(tmp_path / "img01_segm.npy") arr = np.load(npy_path(tmp_path, "segm", _RP, "img01"))
assert arr.dtype == np.uint8 assert arr.dtype == np.uint8
class TestAsyncSaves: class TestAsyncSaves:
def test_depth_async(self, tmp_path: Path) -> None: 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() 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: 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() 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: 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() 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: def test_segmentation_async(self, tmp_path: Path) -> None:
seg = torch.randint(0, 3, (1, 16, 16), dtype=torch.uint8) 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() 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: def test_multiple_async_writes(self, tmp_path: Path) -> None:
for i in range(8): 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) save_npy=True, save_vis=False)
shutdown_io_pool() shutdown_io_pool()
for i in range(8): for i in range(8):
assert (tmp_path / f"img_{i}_depth.npy").exists() assert npy_path(tmp_path, "depth", _RP, f"img_{i}").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()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -190,9 +207,9 @@ class TestMakePalette:
np.testing.assert_array_equal(pal[0], [0, 0, 0]) np.testing.assert_array_equal(pal[0], [0, 0, 0])
def test_cache_returns_same(self) -> None: def test_cache_returns_same(self) -> None:
p1 = make_palette(7) pal = make_palette(7)
p2 = make_palette(7) pal2 = make_palette(7)
assert p1 is p2 assert pal is pal2
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -201,32 +218,31 @@ class TestMakePalette:
class TestConsolidateSafetensors: class TestConsolidateSafetensors:
def _save_all_modalities(self, tmp_path: Path, stem: str = "img01") -> None: def _save_all_modalities(self, tmp_path: Path, stem: str = "img01") -> None:
"""Helper: save all 4 modalities as .npy files."""
H, W = 32, 32 H, W = 32, 32
save_depth(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, 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, 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( save_segmentation(
torch.randint(0, 11, (1, H, W), dtype=torch.uint8), 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: def test_creates_safetensors_file(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
ok = consolidate_safetensors(tmp_path, "img01") ok = consolidate_safetensors(tmp_path, _RP, "img01")
assert ok 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: def test_contains_all_modalities(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
consolidate_safetensors(tmp_path, "img01") consolidate_safetensors(tmp_path, _RP, "img01")
data = st_load_file(tmp_path / "img01.safetensors") data = st_load_file(safetensors_path(tmp_path, _RP, "img01"))
assert set(data.keys()) == {"depth", "edge", "chm", "segm"} assert set(data.keys()) == {"depth", "edge", "chm", "segm"}
def test_dtypes_correct(self, tmp_path: Path) -> None: def test_dtypes_correct(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
consolidate_safetensors(tmp_path, "img01") consolidate_safetensors(tmp_path, _RP, "img01")
data = st_load_file(tmp_path / "img01.safetensors") data = st_load_file(safetensors_path(tmp_path, _RP, "img01"))
assert data["depth"].dtype == torch.float16 assert data["depth"].dtype == torch.float16
assert data["edge"].dtype == torch.float16 assert data["edge"].dtype == torch.float16
assert data["chm"].dtype == torch.float16 assert data["chm"].dtype == torch.float16
@@ -234,70 +250,64 @@ class TestConsolidateSafetensors:
def test_shapes_correct(self, tmp_path: Path) -> None: def test_shapes_correct(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
consolidate_safetensors(tmp_path, "img01") consolidate_safetensors(tmp_path, _RP, "img01")
data = st_load_file(tmp_path / "img01.safetensors") data = st_load_file(safetensors_path(tmp_path, _RP, "img01"))
for key in ("depth", "edge", "chm", "segm"): for key in ("depth", "edge", "chm", "segm"):
assert data[key].shape == (1, 32, 32), f"{key} shape mismatch" assert data[key].shape == (1, 32, 32), f"{key} shape mismatch"
def test_partial_modalities(self, tmp_path: Path) -> None: 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, _RP, "img02", save_npy=True, save_vis=False)
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, _RP, "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, _RP, "img02")
ok = consolidate_safetensors(tmp_path, "img02")
assert ok 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"} assert set(data.keys()) == {"depth", "edge"}
def test_no_modalities_returns_false(self, tmp_path: Path) -> None: 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 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: def test_cleanup_npy(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
consolidate_safetensors(tmp_path, "img01", cleanup_npy=True) consolidate_safetensors(tmp_path, _RP, "img01", cleanup_npy=True)
assert (tmp_path / "img01.safetensors").exists() assert safetensors_path(tmp_path, _RP, "img01").exists()
assert not (tmp_path / "img01_depth.npy").exists() assert not npy_path(tmp_path, "depth", _RP, "img01").exists()
assert not (tmp_path / "img01_edge.npy").exists() assert not npy_path(tmp_path, "edge", _RP, "img01").exists()
assert not (tmp_path / "img01_chm.npy").exists() assert not npy_path(tmp_path, "chm", _RP, "img01").exists()
assert not (tmp_path / "img01_segm.npy").exists() assert not npy_path(tmp_path, "segm", _RP, "img01").exists()
def test_no_cleanup_keeps_npy(self, tmp_path: Path) -> None: def test_no_cleanup_keeps_npy(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
consolidate_safetensors(tmp_path, "img01", cleanup_npy=False) consolidate_safetensors(tmp_path, _RP, "img01", cleanup_npy=False)
assert (tmp_path / "img01.safetensors").exists() assert safetensors_path(tmp_path, _RP, "img01").exists()
assert (tmp_path / "img01_depth.npy").exists() assert npy_path(tmp_path, "depth", _RP, "img01").exists()
def test_async_consolidation(self, tmp_path: Path) -> None: def test_async_consolidation(self, tmp_path: Path) -> None:
self._save_all_modalities(tmp_path) self._save_all_modalities(tmp_path)
consolidate_safetensors_async(tmp_path, "img01") consolidate_safetensors_async(tmp_path, _RP, "img01")
shutdown_io_pool() 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: def test_from_png_only(self, tmp_path: Path) -> None:
"""Consolidation works when only .png exist (no .npy)."""
H, W = 32, 32 H, W = 32, 32
# Save depth/edge/chm as vis-only PNG (no npy). save_depth(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True)
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, _RP, "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, _RP, "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( save_segmentation(
torch.randint(0, 11, (1, H, W), dtype=torch.uint8), 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 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 set(data.keys()) == {"depth", "edge", "chm", "segm"}
assert data["segm"].dtype == torch.uint8 assert data["segm"].dtype == torch.uint8
def test_values_preserved(self, tmp_path: Path) -> None: def test_values_preserved(self, tmp_path: Path) -> None:
"""Verify tensor values survive round-trip."""
depth = torch.rand(1, 16, 16) depth = torch.rand(1, 16, 16)
save_depth(depth, tmp_path, "img03", save_npy=True, save_vis=False) save_depth(depth, tmp_path, _RP, "img03", save_npy=True, save_vis=False)
consolidate_safetensors(tmp_path, "img03") consolidate_safetensors(tmp_path, _RP, "img03")
data = st_load_file(tmp_path / "img03.safetensors") data = st_load_file(safetensors_path(tmp_path, _RP, "img03"))
# float16 round-trip: compare at fp16 precision
expected = depth.half() expected = depth.half()
torch.testing.assert_close(data["depth"], expected, atol=0, rtol=0) torch.testing.assert_close(data["depth"], expected, atol=0, rtol=0)

65
src/tests/test_pipeline_integration.py
View File

@@ -16,7 +16,7 @@ from src.augmentor.dataset import (
discover_images, discover_images,
) )
from src.augmentor.inference import compute_edges_from_depth 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.hardware_conf import HardwareConfig
from src.conf.input_conf import InputConfig from src.conf.input_conf import InputConfig
from src.conf.models_conf import ModelsConfig from src.conf.models_conf import ModelsConfig
@@ -103,9 +103,6 @@ class TestDepthStageIntegration:
) -> None: ) -> None:
from src.main import run_depth_stage 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) 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), \ with patch("src.main.load_depth_model", return_value=mock_model), \
@@ -116,9 +113,9 @@ class TestDepthStageIntegration:
) )
for r in sample_records: for r in sample_records:
npy_path = r.output_dir / f"{r.stem}_depth.npy" p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
assert npy_path.exists(), f"Missing {r.stem}_depth.npy in {r.output_dir}" assert p.exists(), f"Missing depth npy for {r.stem}"
arr = np.load(npy_path) arr = np.load(p)
assert arr.dtype == np.float16 assert arr.dtype == np.float16
assert arr.shape[0] == 1 assert arr.shape[0] == 1
@@ -133,9 +130,6 @@ class TestChmv2StageIntegration:
) -> None: ) -> None:
from src.main import run_chmv2_stage 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( mock_model, mock_processor = _make_mock_chmv2(
input_conf.image_size, input_conf.image_size, input_conf.image_size, input_conf.image_size,
) )
@@ -148,9 +142,9 @@ class TestChmv2StageIntegration:
) )
for r in sample_records: for r in sample_records:
npy_path = r.output_dir / f"{r.stem}_chm.npy" p = npy_path(r.output_root, "chm", r.rel_parent, r.stem)
assert npy_path.exists(), f"Missing {r.stem}_chm.npy in {r.output_dir}" assert p.exists(), f"Missing chm npy for {r.stem}"
arr = np.load(npy_path) arr = np.load(p)
assert arr.dtype == np.float16 assert arr.dtype == np.float16
assert arr.shape[0] == 1 assert arr.shape[0] == 1
@@ -164,18 +158,16 @@ class TestEdgesStageIntegration:
from src.main import run_edges_stage from src.main import run_edges_stage
for r in sample_records: for r in sample_records:
r.output_dir.mkdir(parents=True, exist_ok=True) p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
np.save( p.parent.mkdir(parents=True, exist_ok=True)
r.output_dir / f"{r.stem}_depth.npy", np.save(p, np.random.rand(1, 64, 64).astype(np.float16))
np.random.rand(1, 64, 64).astype(np.float16),
)
run_edges_stage(sample_records, pipeline_conf) run_edges_stage(sample_records, pipeline_conf)
for r in sample_records: for r in sample_records:
npy_path = r.output_dir / f"{r.stem}_edge.npy" p = npy_path(r.output_root, "edge", r.rel_parent, r.stem)
assert npy_path.exists(), f"Missing {r.stem}_edge.npy in {r.output_dir}" assert p.exists(), f"Missing edge npy for {r.stem}"
arr = np.load(npy_path) arr = np.load(p)
assert arr.dtype == np.float16 assert arr.dtype == np.float16
def test_skips_missing_depth( def test_skips_missing_depth(
@@ -185,11 +177,9 @@ class TestEdgesStageIntegration:
) -> None: ) -> None:
from src.main import run_edges_stage 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) run_edges_stage(sample_records, pipeline_conf)
for r in sample_records: 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: class TestSegmentationStageIntegration:
@@ -203,9 +193,6 @@ class TestSegmentationStageIntegration:
) -> None: ) -> None:
from src.main import run_segmentation_stage 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 num_classes = seg_conf.num_classes
mock_model, seg_config_dict = _make_mock_segformer( mock_model, seg_config_dict = _make_mock_segformer(
num_classes, input_conf.image_size, input_conf.image_size, num_classes, input_conf.image_size, input_conf.image_size,
@@ -221,9 +208,9 @@ class TestSegmentationStageIntegration:
) )
for r in sample_records: for r in sample_records:
npy_path = r.output_dir / f"{r.stem}_segm.npy" p = npy_path(r.output_root, "segm", r.rel_parent, r.stem)
assert npy_path.exists() assert p.exists()
arr = np.load(npy_path) arr = np.load(p)
assert arr.dtype == np.uint8 assert arr.dtype == np.uint8
@@ -275,10 +262,11 @@ class TestFullPipelineSmoke:
output_root = Path(pipeline_conf.output_root) output_root = Path(pipeline_conf.output_root)
assert (output_root / "manifest.json").exists() assert (output_root / "manifest.json").exists()
found_depth = list(output_root.rglob("*_depth.npy")) # New dir layout: npy/depth/..., npy/edge/..., etc.
found_edges = list(output_root.rglob("*_edge.npy")) found_depth = list((output_root / "npy" / "depth").rglob("*.npy"))
found_seg = list(output_root.rglob("*_segm.npy")) found_edges = list((output_root / "npy" / "edge").rglob("*.npy"))
found_chmv2 = list(output_root.rglob("*_chm.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_depth) > 0
assert len(found_edges) > 0 assert len(found_edges) > 0
assert len(found_seg) > 0 assert len(found_seg) > 0
@@ -297,6 +285,7 @@ class TestFullPipelineVisOnly:
output_root=str(tmp_path / "output"), output_root=str(tmp_path / "output"),
stages=["depth", "edges", "segmentation", "chmv2"], stages=["depth", "edges", "segmentation", "chmv2"],
save_npy=False, save_vis=True, save_npy=False, save_vis=True,
save_safetensors=False,
save_concat=False, resume=False, log_level="WARNING", save_concat=False, resume=False, log_level="WARNING",
) )
hw_conf = HardwareConfig( hw_conf = HardwareConfig(
@@ -321,8 +310,8 @@ class TestFullPipelineVisOnly:
run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf) run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf)
output_root = Path(pipeline_conf.output_root) output_root = Path(pipeline_conf.output_root)
assert len(list(output_root.rglob("*_depth.png"))) > 0 assert len(list((output_root / "depth").rglob("*.png"))) > 0
assert len(list(output_root.rglob("*_edge.png"))) > 0 assert len(list((output_root / "edge").rglob("*.png"))) > 0
assert len(list(output_root.rglob("*_segm.png"))) > 0 assert len(list((output_root / "segm").rglob("*.png"))) > 0
assert len(list(output_root.rglob("*_chm.png"))) > 0 assert len(list((output_root / "chm").rglob("*.png"))) > 0
assert len(list(output_root.rglob("*.npy"))) == 0 assert len(list(output_root.rglob("*.npy"))) == 0