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

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)