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
│ ├── 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 суммарно):

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 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 / <parent dirs>/."""
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,
}

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.
"""
@@ -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

65
src/main.py
View File

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

17
src/tests/test_dataset.py
View File

@@ -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"}

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

65
src/tests/test_pipeline_integration.py
View File

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