Files
depth_edges_annotate_worlduav/src/main.py
pikaliov f3cb18ac4d Add SafeTensors consolidation stage for zero-copy tensor loading
Bundle all per-image modalities (depth, edge, chm, segm) into a single
.safetensors file for fast training DataLoader reads (~0.1ms zero-copy mmap
vs ~5ms for 4x PNG). Adds consolidate stage after main pipeline stages,
save_safetensors/cleanup_npy config flags, resume support, and 10 new tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-17 16:53:29 +03:00

553 lines
20 KiB
Python

"""Entry point for the depth/edges/segmentation/chmv2 augmentation pipeline.
All parameters loaded from gin config files — no argparse.
Sequential stage processing: one model at a time, load → process all → unload.
Usage:
python -m src.main
"""
from __future__ import annotations
import gc
import json
import logging
import time
from datetime import datetime
from pathlib import Path
from typing import Any
import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from src.conf.config_loader import load_all_configs
from src.conf.pipeline_conf import PipelineConfig
from src.conf.hardware_conf import HardwareConfig
from src.conf.models_conf import ModelsConfig
from src.conf.input_conf import InputConfig
from src.conf.seg_conf import SegConfig
from src.utils.utils_file_dir import get_proj_dir
from src.augmentor.dataset import (
AugmentDataset, ImageRecord, attach_output_dirs,
discover_images, filter_completed, split_by_view,
)
from src.augmentor.inference import (
compute_edges_from_depth, infer_depth_batch, infer_chmv2_batch,
infer_segmentation_batch,
)
from src.augmentor.io_utils import (
save_depth_async, save_chmv2_async, save_edges_async,
save_segmentation_async, consolidate_safetensors,
setup_logging, shutdown_io_pool,
)
from src.augmentor.models import (
load_depth_model, load_chmv2_model, load_segmentation_model, unload_model,
)
from src.utils.profiler import (
log_system_info, log_disk_info, log_vram_snapshot, log_ram_snapshot,
)
logger = logging.getLogger(__name__)
_STAGE_EMOJI = {
"depth": "🌊",
"edges": "🔪",
"segmentation": "🗺️",
"chmv2": "🦕",
"concat": "🧩",
"consolidate": "📦",
}
def _silence_model_loggers() -> None:
"""Suppress verbose inference logs from all models."""
import os
import warnings
os.environ["DA3_LOG_LEVEL"] = "ERROR"
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
for name in (
"depth_anything_3", "depth_anything_3.api",
"depth_anything_3.utils.logger",
"transformers", "transformers.modeling_utils",
"transformers.configuration_utils", "transformers.image_processing_utils",
"transformers.image_processing_base",
"sam3", "segearthov3_segmentor",
"huggingface_hub", "httpx", "filelock",
"numexpr", "numexpr.utils",
"torch", "py.warnings",
):
logging.getLogger(name).setLevel(logging.ERROR)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", message=".*not sharded.*")
# ---------------------------------------------------------------------------
# Stage runners
# ---------------------------------------------------------------------------
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).
"""
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)
groups: list[tuple[list[ImageRecord], int, str]] = []
if db_recs:
groups.append((db_recs, input_conf.image_size, f"db {input_conf.image_size}"))
if query_recs:
groups.append((query_recs, input_conf.query_image_size, f"query {input_conf.query_image_size}"))
return groups
def run_depth_stage(
records: list[ImageRecord],
pipeline_conf: PipelineConfig,
hw_conf: HardwareConfig,
models_conf: ModelsConfig,
input_conf: InputConfig,
device: torch.device,
) -> None:
"""🌊 Load DA3, process all images, unload."""
model = load_depth_model(models_conf, hw_conf, device)
for group_records, sz, label in _resolve_image_sizes(records, input_conf):
bs = hw_conf.find_batch_size(
inference_fn=lambda x: infer_depth_batch(model, x, device),
input_shape=(3, sz, sz),
device=device,
)
ds = AugmentDataset(group_records, image_size=sz)
logger.info("🌊 [depth/%s] DataLoader: batch_size=%d, %d images, %d batches",
label, bs, len(ds), (len(ds) + bs - 1) // bs)
loader = DataLoader(
ds, batch_size=bs, shuffle=False,
num_workers=hw_conf.num_workers, pin_memory=True,
persistent_workers=hw_conf.num_workers > 0,
prefetch_factor=4 if hw_conf.num_workers > 0 else None,
)
total_images = len(ds)
pbar = tqdm(loader, desc=f"🌊 depth/{label} (bs={bs})", unit="batch",
total=len(loader), colour="green",
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
processed = 0
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]),
stem=batch["stem"][i],
save_npy=pipeline_conf.save_npy,
save_vis=pipeline_conf.save_vis)
processed += depths.shape[0]
pbar.set_postfix(images=f"{processed}/{total_images}")
shutdown_io_pool()
unload_model(model)
def run_edges_stage(
records: list[ImageRecord],
pipeline_conf: PipelineConfig,
batch_size: int = 32,
) -> None:
"""🔪 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():
valid.append(r)
else:
logger.warning("⚠️ No depth for %s, skipping edges.", r.rel_path)
total_images = len(valid)
pbar = tqdm(range(0, len(valid), batch_size), desc="🔪 edges (sobel)", unit="batch",
colour="cyan",
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
processed = 0
for start in pbar:
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)
else:
from PIL import Image
d = np.array(Image.open(png_path)).astype(np.float32) / 255.0
if d.ndim == 2:
d = d[np.newaxis]
depth_tensors.append(torch.from_numpy(d))
depths = torch.stack(depth_tensors)
if depths.ndim == 3:
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_npy=pipeline_conf.save_npy,
save_vis=pipeline_conf.save_vis)
processed += len(chunk)
pbar.set_postfix(images=f"{processed}/{total_images}")
shutdown_io_pool()
def run_chmv2_stage(
records: list[ImageRecord],
pipeline_conf: PipelineConfig,
hw_conf: HardwareConfig,
models_conf: ModelsConfig,
input_conf: InputConfig,
device: torch.device,
) -> None:
"""🦕 Load CHMv2 (DINOv3), process all images, unload."""
model, processor = load_chmv2_model(models_conf, hw_conf, device)
for group_records, sz, label in _resolve_image_sizes(records, input_conf):
bs = hw_conf.find_batch_size(
inference_fn=lambda x: infer_chmv2_batch(model, processor, x, device),
input_shape=(3, sz, sz),
device=device,
)
ds = AugmentDataset(group_records, image_size=sz)
logger.info("🦕 [chmv2/%s] DataLoader: batch_size=%d, %d images, %d batches",
label, bs, len(ds), (len(ds) + bs - 1) // bs)
loader = DataLoader(
ds, batch_size=bs, shuffle=False,
num_workers=hw_conf.num_workers, pin_memory=True,
persistent_workers=hw_conf.num_workers > 0,
prefetch_factor=4 if hw_conf.num_workers > 0 else None,
)
total_images = len(ds)
pbar = tqdm(loader, desc=f"🦕 chmv2/{label} (bs={bs})", unit="batch",
total=len(loader), colour="blue",
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
processed = 0
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]),
stem=batch["stem"][i],
save_npy=pipeline_conf.save_npy,
save_vis=pipeline_conf.save_vis)
processed += depths.shape[0]
pbar.set_postfix(images=f"{processed}/{total_images}")
shutdown_io_pool()
unload_model(model)
def run_segmentation_stage(
records: list[ImageRecord],
pipeline_conf: PipelineConfig,
hw_conf: HardwareConfig,
models_conf: ModelsConfig,
input_conf: InputConfig,
seg_conf: SegConfig,
device: torch.device,
) -> None:
"""🗺️ Load segmentation model, process all images, unload."""
model, seg_config = load_segmentation_model(models_conf, hw_conf, seg_conf, device)
num_classes = seg_config.get("num_classes", 150)
is_segearth = seg_config.get("type") == "segearth-ov3"
for group_records, sz, label in _resolve_image_sizes(records, input_conf):
if is_segearth:
bs = 16
else:
bs = hw_conf.find_batch_size(
inference_fn=lambda x: infer_segmentation_batch(model, seg_config, x, device),
input_shape=(3, sz, sz),
device=device,
)
ds = AugmentDataset(group_records, image_size=sz)
total_images = len(ds)
logger.info("🗺️ [segmentation/%s] DataLoader: batch_size=%d, %d images, %d batches",
label, bs, total_images, (total_images + bs - 1) // bs)
loader = DataLoader(
ds, batch_size=bs, shuffle=False,
num_workers=hw_conf.num_workers, pin_memory=True,
persistent_workers=hw_conf.num_workers > 0,
prefetch_factor=4 if hw_conf.num_workers > 0 else None,
)
seg_type = "SegEarth-OV3" if is_segearth else "SegFormer"
pbar = tqdm(loader, desc=f"🗺️ seg/{label} {seg_type} (bs={bs})", unit="batch",
colour="yellow",
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
processed = 0
for batch in pbar:
segs = infer_segmentation_batch(
model, seg_config, batch["image_raw"], device,
)
for j in range(segs.shape[0]):
save_segmentation_async(
segs[j], Path(batch["output_dir"][j]), stem=batch["stem"][j],
save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis,
num_classes=num_classes,
)
processed += segs.shape[0]
pbar.set_postfix(images=f"{processed}/{total_images}")
shutdown_io_pool()
unload_model(model)
def run_consolidate_stage(
records: list[ImageRecord],
cleanup_npy: bool = False,
) -> None:
"""📦 Bundle per-image .npy/.png modalities into .safetensors files."""
total = len(records)
written = 0
pbar = tqdm(records, desc="📦 consolidate → safetensors", unit="img",
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)
if ok:
written += 1
pbar.set_postfix(written=f"{written}/{total}")
logger.info("📦 Consolidated %d / %d images to .safetensors.", written, total)
# ---------------------------------------------------------------------------
# Pipeline orchestration
# ---------------------------------------------------------------------------
def run_pipeline(
pipeline_conf: PipelineConfig,
hw_conf: HardwareConfig,
models_conf: ModelsConfig,
input_conf: InputConfig,
seg_conf: SegConfig,
) -> None:
"""Execute the full augmentation pipeline: one stage at a time."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type != "cuda":
logger.warning("⚠️ CUDA not available, running on CPU (very slow).")
_silence_model_loggers()
# System profiling at startup.
log_system_info()
log_disk_info(Path(pipeline_conf.input_root), Path(pipeline_conf.output_root))
# Discover images.
input_root = Path(pipeline_conf.input_root)
output_root = Path(pipeline_conf.output_root)
logger.info(
"🔍 Discovering images in %s (subset=%s, source=%s) ...",
input_root, pipeline_conf.subset or "all", pipeline_conf.source or "all",
)
all_records = discover_images(input_root, subset=pipeline_conf.subset,
source=pipeline_conf.source)
all_records = attach_output_dirs(all_records, output_root)
logger.info("📸 Found %d images.", len(all_records))
if not all_records:
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)
# Process each stage sequentially.
stage_times: dict[str, float] = {}
stage_counts: dict[str, int] = {}
failed_stages: set[str] = set()
for stage in pipeline_conf.stages:
emoji = _STAGE_EMOJI.get(stage, "⚙️")
if stage == "edges" and "depth" in failed_stages:
logger.error("❌ [edges] Skipped — depth stage failed.")
failed_stages.add(stage)
continue
pending, skipped = filter_completed(all_records, stage)
logger.info("%s [%s] %d pending, %d skipped.", emoji, stage, len(pending), skipped)
stage_counts[stage] = len(pending)
if not pending:
stage_times[stage] = 0.0
continue
logger.info("%s [%s] Starting stage...", emoji, stage)
log_vram_snapshot(f"before {stage}")
log_ram_snapshot(f"before {stage}")
t0 = time.perf_counter()
try:
if stage == "depth":
run_depth_stage(pending, pipeline_conf, hw_conf, models_conf,
input_conf, device)
elif stage == "edges":
run_edges_stage(pending, pipeline_conf)
elif stage == "segmentation":
run_segmentation_stage(pending, pipeline_conf, hw_conf, models_conf,
input_conf, seg_conf, device)
elif stage == "chmv2":
run_chmv2_stage(pending, pipeline_conf, hw_conf, models_conf,
input_conf, device)
except Exception:
logger.exception("💥 Stage '%s' failed.", stage)
failed_stages.add(stage)
elapsed = time.perf_counter() - t0
stage_times[stage] = elapsed
if stage not in failed_stages:
logger.info("✅ [%s] Completed in %.1f s (%d images).", stage, elapsed, len(pending))
log_vram_snapshot(f"after {stage}")
log_ram_snapshot(f"after {stage}")
# SafeTensors consolidation: bundle all modalities per image.
if pipeline_conf.save_safetensors:
pending_st, skipped_st = filter_completed(all_records, "consolidate")
logger.info("📦 [consolidate] %d pending, %d skipped.", len(pending_st), skipped_st)
if pending_st:
t0 = time.perf_counter()
run_consolidate_stage(pending_st, cleanup_npy=pipeline_conf.cleanup_npy)
elapsed_st = time.perf_counter() - t0
stage_times["consolidate"] = elapsed_st
stage_counts["consolidate"] = len(pending_st)
logger.info("✅ [consolidate] Completed in %.1f s (%d images).",
elapsed_st, len(pending_st))
else:
stage_times["consolidate"] = 0.0
stage_counts["consolidate"] = 0
# Manifest.
manifest = {
"pipeline_version": "3.3.0-safetensors",
"image_size_db": input_conf.image_size,
"image_size_query": input_conf.query_image_size,
"profile": hw_conf.profile_name,
"models": {
"depth": models_conf.depth_model_id,
"edges": "Sobel from depth (CPU)",
"segmentation": models_conf.seg_model_type,
"chmv2": models_conf.chmv2_model_id,
},
"save_safetensors": pipeline_conf.save_safetensors,
"seg_prompts": seg_conf.prompts,
"total_images": len(all_records),
"stages": {
s: {"processed": stage_counts.get(s, 0),
"time_sec": round(stage_times.get(s, 0), 1)}
for s in pipeline_conf.stages
},
"timestamp": datetime.now().isoformat(),
}
manifest_path = output_root / "manifest.json"
manifest_path.parent.mkdir(parents=True, exist_ok=True)
manifest_path.write_text(
json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8",
)
# Summary.
total = sum(stage_times.values())
logger.info("=" * 60)
logger.info("🏁 DONE: %d images, %.1f s total", len(all_records), total)
for s, t in stage_times.items():
cnt = stage_counts.get(s, len(all_records))
fps = cnt / t if t > 0 else 0
emoji = _STAGE_EMOJI.get(s, "⚙️")
logger.info(" %s %-13s %6.1f s (%d images, %.0f FPS)", emoji, s, t, cnt, fps)
logger.info("📂 Output: %s", output_root)
logger.info("=" * 60)
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
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'"
"""
import argparse
parser = argparse.ArgumentParser(description="Augmentation pipeline")
parser.add_argument("--gin", action="append", default=[],
help="Gin parameter overrides (repeatable)")
args = parser.parse_args()
_silence_model_loggers()
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)
gin_files = sorted(cfg_dir.glob("*.gin"))
_gin.clear_config()
_gin.parse_config_files_and_bindings(
config_files=[str(f) for f in gin_files],
bindings=args.gin,
)
configs = {
"pipeline": PipelineConfig(),
"hardware": HardwareConfig(),
"models": ModelsConfig(),
"input": InputConfig(),
"seg": SegConfig(),
}
else:
configs = load_all_configs(path2cfg)
pipeline_conf: PipelineConfig = configs["pipeline"]
setup_logging(pipeline_conf.log_level,
log_file=Path(pipeline_conf.output_root) / "pipeline.log")
torch.manual_seed(42)
np.random.seed(42)
run_pipeline(
configs["pipeline"],
configs["hardware"],
configs["models"],
configs["input"],
configs["seg"],
)
if __name__ == "__main__":
main()