Add generic entry point for arbitrary RGB folders/single images
scripts/run_folder.py annotates any folder or single image with the same depth/edges/segmentation/chmv2 -> safetensors pipeline used by the dataset scripts, without World-UAV-specific scene/dir filters. run_pipeline() gains an optional records= parameter to bypass discovery for explicit inputs. Resume now also recognizes modalities already present in a consolidated .safetensors file, so a save_vis=False run can be resumed without redoing GPU stages. --no-vis + --no-safetensors together is rejected instead of silently running inference with no output. psutil made optional in profiler.py (CPU-core fallback via os.cpu_count()) since it was missing from the local test venv, unblocking 7 pre-existing tests unrelated to this change.
This commit is contained in:
21
README.md
21
README.md
@@ -30,10 +30,28 @@ python scripts/run_uav_visloc.py
|
||||
# GTA-UAV-LR
|
||||
python scripts/run_gta_uav.py
|
||||
|
||||
# Тесты (149 шт, без GPU)
|
||||
# Тесты (без GPU)
|
||||
python -m pytest src/tests/ -v
|
||||
```
|
||||
|
||||
### Произвольная папка / одно изображение
|
||||
|
||||
Универсальная точка входа без датасет-специфичных путей — `scripts/run_folder.py`. Принимает папку с RGB-изображениями (рекурсивный обход, относительные пути сохраняются в выходном layout) или одиночный файл (`.png/.jpg/.jpeg/.bmp`):
|
||||
|
||||
```bash
|
||||
# Папка: выход по умолчанию — сиблинг <input>-aug
|
||||
python scripts/run_folder.py /path/to/images
|
||||
|
||||
# Одно изображение: выход по умолчанию — <родитель>-aug
|
||||
python scripts/run_folder.py /path/to/photo.jpg
|
||||
|
||||
# Явный выход, только depth+edges, без PNG-визуализаций
|
||||
python scripts/run_folder.py /path/to/images --output /path/to/out \
|
||||
--stages depth edges --no-vis
|
||||
```
|
||||
|
||||
Дефолты: все 4 стадии, `image_size=256`, unified 17 классов (`threshold=0.15`), dark-water fix включён, `.safetensors` консолидация включена (`--no-safetensors` для отключения), resume включён. GTA-специфичная переклассификация wetland выключена (включается флагом `--wetland-reclassify`). `--query-image-size` актуален только если в папке есть `drone/`/`query/`-поддиректории.
|
||||
|
||||
## Структура проекта
|
||||
|
||||
```
|
||||
@@ -77,6 +95,7 @@ python -m pytest src/tests/ -v
|
||||
│ ├── seg_classes.py # UNIFIED_PROMPTS — 17 классов (единый источник)
|
||||
│ ├── run_uav_visloc.py # Запуск для UAV_VisLoc
|
||||
│ ├── run_gta_uav.py # Запуск для GTA-UAV-LR
|
||||
│ ├── run_folder.py # Произвольная папка / одно изображение
|
||||
│ └── migrate_layout.py # Миграция со старого prefix-формата
|
||||
└── docs/
|
||||
├── segmentation_class_analysis.md # Анализ классов сегментации (11 классов)
|
||||
|
||||
235
scripts/run_folder.py
Normal file
235
scripts/run_folder.py
Normal file
@@ -0,0 +1,235 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run annotation pipeline on an arbitrary folder of RGB images or a single image.
|
||||
|
||||
Generic entry point: no dataset-specific assumptions. Images are discovered
|
||||
recursively with a local walker (relative paths preserved in the output
|
||||
layout); the World-UAV incomplete-scene and service-dir filters of
|
||||
``discover_images`` are deliberately NOT applied — an arbitrary folder may
|
||||
legitimately contain directories named ``SoHo`` or ``Index``. A single image
|
||||
file is also accepted — its modalities land directly under the output root.
|
||||
|
||||
Usage:
|
||||
python scripts/run_folder.py /path/to/images
|
||||
python scripts/run_folder.py /path/to/photo.jpg
|
||||
python scripts/run_folder.py /path/to/images --output /path/to/out
|
||||
python scripts/run_folder.py /path/to/images --stages depth edges --no-vis
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Ensure project root is importable.
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from src.conf.hardware_conf import HardwareConfig
|
||||
from src.conf.input_conf import InputConfig
|
||||
from src.conf.models_conf import ModelsConfig
|
||||
from src.conf.pipeline_conf import PipelineConfig
|
||||
from src.conf.seg_conf import SegConfig
|
||||
from src.augmentor.dataset import EXTENSIONS, ImageRecord
|
||||
from src.augmentor.io_utils import setup_logging
|
||||
from src.main import run_pipeline
|
||||
from scripts.seg_classes import UNIFIED_PROMPTS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Only truly generic service dirs are skipped; the dataset-specific
|
||||
#: EXCLUDE_DIRS / INCOMPLETE_SCENES of ``discover_images`` do not apply here.
|
||||
GENERIC_EXCLUDE_DIRS = {"__pycache__", "__MACOSX"}
|
||||
|
||||
|
||||
def discover_folder_images(
|
||||
root: Path,
|
||||
source: str | None = None,
|
||||
) -> list[ImageRecord]:
|
||||
"""Recursively find images under *root* without dataset-specific filters.
|
||||
|
||||
Args:
|
||||
root: Folder to scan.
|
||||
source: Optional filter — 'query' keeps drone/query images, 'db' keeps
|
||||
satellite/DB images (same path-part convention as discover_images).
|
||||
|
||||
Returns:
|
||||
Sorted list of ImageRecord with output_root left unset.
|
||||
"""
|
||||
records: list[ImageRecord] = []
|
||||
for p in sorted(root.rglob("*")):
|
||||
if not p.is_file() or p.suffix.lower() not in EXTENSIONS:
|
||||
continue
|
||||
if any(d in p.parts for d in GENERIC_EXCLUDE_DIRS):
|
||||
continue
|
||||
rel = p.relative_to(root)
|
||||
if source is not None:
|
||||
is_db = "DB" in rel.parts or "satellite" in rel.parts
|
||||
is_query = "query" in rel.parts or "drone" in rel.parts
|
||||
if source == "query" and is_db:
|
||||
continue
|
||||
if source == "db" and is_query:
|
||||
continue
|
||||
records.append(ImageRecord(
|
||||
abs_path=p,
|
||||
rel_path=str(rel),
|
||||
stem=p.stem,
|
||||
output_root=Path(),
|
||||
rel_parent=str(rel.parent),
|
||||
))
|
||||
return records
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
"""Build the argparse parser (separate function for testability)."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Annotate an arbitrary folder of RGB images or a single image",
|
||||
)
|
||||
parser.add_argument("input",
|
||||
help="Path to a folder with RGB images or a single "
|
||||
"image file (.png/.jpg/.jpeg/.bmp)")
|
||||
parser.add_argument("--output", default=None,
|
||||
help="Output root (default: sibling '<input>-aug'; "
|
||||
"for a single file: '<parent>-aug')")
|
||||
parser.add_argument("--stages", nargs="+",
|
||||
default=["depth", "edges", "segmentation", "chmv2"],
|
||||
help="Stages to run")
|
||||
parser.add_argument("--image-size", type=int, default=256,
|
||||
help="Output resolution for db/satellite images")
|
||||
parser.add_argument("--query-image-size", type=int, default=None,
|
||||
help="Output resolution for query/drone images "
|
||||
"(default: same as --image-size; only relevant "
|
||||
"when the folder contains drone/query subdirs)")
|
||||
parser.add_argument("--source", choices=["db", "drone", "all"], default="all",
|
||||
help="Process only db (satellite), drone, or all "
|
||||
"(default; ignored for a single-file input)")
|
||||
parser.add_argument("--no-vis", action="store_true",
|
||||
help="Do not save PNG visualizations")
|
||||
parser.add_argument("--no-safetensors", action="store_true",
|
||||
help="Do not consolidate modalities into .safetensors")
|
||||
parser.add_argument("--wetland-reclassify", action="store_true",
|
||||
help="Reclassify wetland into vegetation/bare soil "
|
||||
"(GTA-specific post-processing, off by default)")
|
||||
parser.add_argument("--weights-dir",
|
||||
default=str(_PROJECT_ROOT / "in" / "weights"),
|
||||
help="Directory with model weights")
|
||||
parser.add_argument("--num-workers", type=int, default=4,
|
||||
help="DataLoader workers")
|
||||
parser.add_argument("--profile", default="rtx4090",
|
||||
help="Hardware profile name")
|
||||
return parser
|
||||
|
||||
|
||||
def resolve_output_root(input_path: Path, output: str | None) -> Path:
|
||||
"""Return the output root: explicit --output or a '-aug' sibling."""
|
||||
if output is not None:
|
||||
return Path(output)
|
||||
base = input_path if input_path.is_dir() else input_path.parent
|
||||
return base.parent / f"{base.name}-aug"
|
||||
|
||||
|
||||
def build_single_file_record(image_path: Path) -> ImageRecord:
|
||||
"""Build one ImageRecord for a standalone image file."""
|
||||
return ImageRecord(
|
||||
abs_path=image_path,
|
||||
rel_path=image_path.name,
|
||||
stem=image_path.stem,
|
||||
output_root=Path(),
|
||||
rel_parent=".",
|
||||
)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.no_vis and args.no_safetensors:
|
||||
parser.error(
|
||||
"--no-vis together with --no-safetensors would run inference "
|
||||
"without writing any output; drop one of the flags"
|
||||
)
|
||||
|
||||
import gin
|
||||
gin.clear_config()
|
||||
|
||||
input_path = Path(args.input).resolve()
|
||||
if not input_path.exists():
|
||||
raise SystemExit(f"Input path does not exist: {input_path}")
|
||||
|
||||
source = None if args.source == "all" else args.source
|
||||
if source == "drone":
|
||||
source = "query"
|
||||
|
||||
if input_path.is_file():
|
||||
if input_path.suffix.lower() not in EXTENSIONS:
|
||||
raise SystemExit(
|
||||
f"Unsupported image extension '{input_path.suffix}' "
|
||||
f"(expected one of {sorted(EXTENSIONS)})"
|
||||
)
|
||||
records = [build_single_file_record(input_path)]
|
||||
input_root = input_path.parent
|
||||
else:
|
||||
input_root = input_path
|
||||
records = discover_folder_images(input_root, source=source)
|
||||
if not records:
|
||||
raise SystemExit(f"No RGB images found under: {input_root}")
|
||||
|
||||
output_root = resolve_output_root(input_path, args.output)
|
||||
|
||||
pipeline_conf = PipelineConfig(
|
||||
input_root=str(input_root),
|
||||
output_root=str(output_root),
|
||||
stages=args.stages,
|
||||
save_npy=False,
|
||||
save_vis=not args.no_vis,
|
||||
save_safetensors=not args.no_safetensors,
|
||||
cleanup_npy=True,
|
||||
seg_fix_dark_water=True,
|
||||
seg_reclassify_wetland=args.wetland_reclassify,
|
||||
resume=True,
|
||||
source=source,
|
||||
log_level="INFO",
|
||||
)
|
||||
|
||||
hw_conf = HardwareConfig(
|
||||
profile_name=args.profile,
|
||||
total_ram_gb=24.0,
|
||||
reserve_gb=2.0,
|
||||
use_fp16=True,
|
||||
batch_size=None, # auto
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
|
||||
input_conf = InputConfig(
|
||||
image_size=args.image_size,
|
||||
query_image_size=args.query_image_size or args.image_size,
|
||||
)
|
||||
|
||||
# Unified 17-class prompt list — shared across all datasets.
|
||||
seg_conf = SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS)
|
||||
|
||||
models_conf = ModelsConfig(weights_dir=args.weights_dir)
|
||||
|
||||
setup_logging(
|
||||
pipeline_conf.log_level,
|
||||
log_file=output_root / "pipeline.log",
|
||||
)
|
||||
|
||||
if input_path.is_file() and source is not None:
|
||||
logger.warning("--source is ignored for a single-file input.")
|
||||
logger.info("Input: %s (%d image(s)) -> %s",
|
||||
input_path, len(records), output_root)
|
||||
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf,
|
||||
records=records)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -142,6 +142,19 @@ STAGE_MODALITY: dict[str, str] = {
|
||||
}
|
||||
|
||||
|
||||
def _consolidated_keys(st_path: Path) -> set[str]:
|
||||
"""Return tensor keys of a consolidated .safetensors file (empty if unreadable)."""
|
||||
try:
|
||||
from safetensors import safe_open
|
||||
except ImportError:
|
||||
return set()
|
||||
try:
|
||||
with safe_open(str(st_path), framework="pt") as f:
|
||||
return set(f.keys())
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
|
||||
def filter_completed(
|
||||
records: list[ImageRecord],
|
||||
stage: str,
|
||||
@@ -159,8 +172,16 @@ def filter_completed(
|
||||
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)
|
||||
continue
|
||||
# A consolidated .safetensors already holding this modality also counts
|
||||
# as completed: a save_vis=False run leaves neither .npy nor .png behind
|
||||
# (cleanup_npy), and without this check every resume would redo the GPU
|
||||
# stages. safe_open reads only the file header — cheap per record.
|
||||
st_p = safetensors_path(r.output_root, r.rel_parent, r.stem)
|
||||
if st_p.exists() and modality in _consolidated_keys(st_p):
|
||||
skipped += 1
|
||||
continue
|
||||
pending.append(r)
|
||||
return pending, skipped
|
||||
|
||||
|
||||
|
||||
32
src/main.py
32
src/main.py
@@ -370,8 +370,15 @@ def run_pipeline(
|
||||
models_conf: ModelsConfig,
|
||||
input_conf: InputConfig,
|
||||
seg_conf: SegConfig,
|
||||
records: list[ImageRecord] | None = None,
|
||||
) -> None:
|
||||
"""Execute the full augmentation pipeline: one stage at a time."""
|
||||
"""Execute the full augmentation pipeline: one stage at a time.
|
||||
|
||||
Args:
|
||||
records: Optional pre-built list of ImageRecord. When provided,
|
||||
image discovery is skipped and exactly these records are
|
||||
processed (output_root is attached automatically).
|
||||
"""
|
||||
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).")
|
||||
@@ -383,17 +390,22 @@ def run_pipeline(
|
||||
print(Path(pipeline_conf.input_root))
|
||||
log_disk_info(Path(pipeline_conf.input_root), Path(pipeline_conf.output_root))
|
||||
|
||||
# Discover images.
|
||||
# Discover images (or use externally provided records).
|
||||
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 records is not None:
|
||||
all_records = attach_output_dirs(records, output_root)
|
||||
logger.info("📸 Using %d provided records (discovery skipped).",
|
||||
len(all_records))
|
||||
else:
|
||||
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.")
|
||||
|
||||
309
src/tests/test_run_folder.py
Normal file
309
src/tests/test_run_folder.py
Normal file
@@ -0,0 +1,309 @@
|
||||
"""Tests for scripts/run_folder.py and run_pipeline(records=...)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from src.augmentor.dataset import ImageRecord
|
||||
from src.augmentor.io_utils import npy_path, safetensors_path
|
||||
from src.conf.hardware_conf import HardwareConfig
|
||||
from src.conf.input_conf import InputConfig
|
||||
from src.conf.models_conf import ModelsConfig
|
||||
from src.conf.pipeline_conf import PipelineConfig
|
||||
from src.conf.seg_conf import SegConfig
|
||||
from src.tests.test_pipeline_integration import (
|
||||
_make_mock_chmv2,
|
||||
_make_mock_depth_model_da3,
|
||||
_make_mock_segformer,
|
||||
)
|
||||
|
||||
from scripts.run_folder import (
|
||||
build_parser,
|
||||
build_single_file_record,
|
||||
resolve_output_root,
|
||||
)
|
||||
from scripts.seg_classes import UNIFIED_PROMPTS
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (a) run_pipeline with explicit records: no discovery, exact set processed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestRunPipelineWithRecords:
|
||||
def test_processes_exactly_given_records_and_skips_discovery(
|
||||
self, fake_image_dir: Path, tmp_path: Path,
|
||||
) -> None:
|
||||
import gin
|
||||
gin.clear_config()
|
||||
|
||||
img_dir = fake_image_dir / "scene01" / "query"
|
||||
all_files = sorted(img_dir.glob("*.png"))
|
||||
chosen = all_files[:2]
|
||||
records = [
|
||||
ImageRecord(
|
||||
abs_path=f,
|
||||
rel_path=str(Path("scene01") / "query" / f.name),
|
||||
stem=f.stem,
|
||||
output_root=Path(),
|
||||
rel_parent=str(Path("scene01") / "query"),
|
||||
)
|
||||
for f in chosen
|
||||
]
|
||||
|
||||
pipeline_conf = PipelineConfig(
|
||||
input_root=str(fake_image_dir),
|
||||
output_root=str(tmp_path / "output_records"),
|
||||
stages=["depth"],
|
||||
save_npy=True, save_vis=False,
|
||||
save_safetensors=False,
|
||||
resume=False, log_level="WARNING",
|
||||
)
|
||||
hw_conf = HardwareConfig(
|
||||
profile_name="test", total_ram_gb=8.0, reserve_gb=1.0,
|
||||
batch_size=2, num_workers=0,
|
||||
)
|
||||
input_conf = InputConfig(image_size=32)
|
||||
seg_conf = SegConfig(prompts=["bg", "building", "road"])
|
||||
|
||||
mock_depth = _make_mock_depth_model_da3(32, 32)
|
||||
|
||||
with patch("src.main.load_depth_model", return_value=mock_depth), \
|
||||
patch("src.main.unload_model"), \
|
||||
patch("src.main.discover_images") as mock_discover:
|
||||
from src.main import run_pipeline
|
||||
run_pipeline(pipeline_conf, hw_conf, ModelsConfig(),
|
||||
input_conf, seg_conf, records=records)
|
||||
|
||||
mock_discover.assert_not_called()
|
||||
|
||||
output_root = Path(pipeline_conf.output_root)
|
||||
for r in records:
|
||||
p = npy_path(output_root, "depth", r.rel_parent, r.stem)
|
||||
assert p.exists(), f"Missing depth npy for {r.stem}"
|
||||
# Only the 2 given records were processed, not all 4 discovered ones.
|
||||
produced = list((output_root / "npy" / "depth").rglob("*.npy"))
|
||||
assert len(produced) == len(records)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (b) Single-file record construction and path normalization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSingleFileRecord:
|
||||
def test_record_fields(self, tmp_path: Path) -> None:
|
||||
img = tmp_path / "photo.jpg"
|
||||
img.touch()
|
||||
r = build_single_file_record(img)
|
||||
assert r.abs_path == img
|
||||
assert r.rel_path == "photo.jpg"
|
||||
assert r.stem == "photo"
|
||||
assert r.rel_parent == "."
|
||||
|
||||
def test_dot_rel_parent_collapses_in_paths(self, tmp_path: Path) -> None:
|
||||
# pathlib collapses '.' components: the safetensors file lands
|
||||
# directly under <out>/safetensors/ with no intermediate dir.
|
||||
out = tmp_path / "out"
|
||||
st = safetensors_path(out, ".", "photo")
|
||||
assert st == out / "safetensors" / "photo.safetensors"
|
||||
np_p = npy_path(out, "depth", ".", "photo")
|
||||
assert np_p == out / "npy" / "depth" / "photo.npy"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (c) argparse defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestBuildParser:
|
||||
def test_defaults(self) -> None:
|
||||
args = build_parser().parse_args(["some/input"])
|
||||
assert args.input == "some/input"
|
||||
assert args.output is None
|
||||
assert args.stages == ["depth", "edges", "segmentation", "chmv2"]
|
||||
assert args.image_size == 256
|
||||
assert args.query_image_size is None
|
||||
assert args.source == "all"
|
||||
assert args.no_vis is False
|
||||
assert args.no_safetensors is False
|
||||
assert args.wetland_reclassify is False
|
||||
assert args.num_workers == 4
|
||||
assert args.profile == "rtx4090"
|
||||
assert Path(args.weights_dir).name == "weights"
|
||||
|
||||
def test_overrides(self) -> None:
|
||||
args = build_parser().parse_args([
|
||||
"in_dir", "--output", "out_dir", "--stages", "depth", "edges",
|
||||
"--image-size", "128", "--query-image-size", "512",
|
||||
"--source", "drone", "--no-vis", "--no-safetensors",
|
||||
"--wetland-reclassify", "--num-workers", "0",
|
||||
])
|
||||
assert args.output == "out_dir"
|
||||
assert args.stages == ["depth", "edges"]
|
||||
assert args.image_size == 128
|
||||
assert args.query_image_size == 512
|
||||
assert args.source == "drone"
|
||||
assert args.no_vis is True
|
||||
assert args.no_safetensors is True
|
||||
assert args.wetland_reclassify is True
|
||||
assert args.num_workers == 0
|
||||
|
||||
def test_resolve_output_root_folder(self, tmp_path: Path) -> None:
|
||||
folder = tmp_path / "myphotos"
|
||||
folder.mkdir()
|
||||
assert resolve_output_root(folder, None) == tmp_path / "myphotos-aug"
|
||||
assert resolve_output_root(folder, "explicit") == Path("explicit")
|
||||
|
||||
def test_resolve_output_root_single_file(self, tmp_path: Path) -> None:
|
||||
folder = tmp_path / "myphotos"
|
||||
folder.mkdir()
|
||||
img = folder / "a.png"
|
||||
img.touch()
|
||||
assert resolve_output_root(img, None) == tmp_path / "myphotos-aug"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (d) End-to-end single image: canonical safetensors keys/dtypes/shapes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSingleImageEndToEnd:
|
||||
def test_safetensors_canon(self, tmp_path: Path) -> None:
|
||||
img_path = tmp_path / "solo" / "photo.png"
|
||||
img_path.parent.mkdir(parents=True)
|
||||
Image.fromarray(
|
||||
np.random.randint(0, 255, (48, 48, 3), dtype=np.uint8),
|
||||
).save(img_path)
|
||||
|
||||
out = tmp_path / "solo-out"
|
||||
H = W = 32
|
||||
num_classes = len(UNIFIED_PROMPTS)
|
||||
|
||||
mock_depth = _make_mock_depth_model_da3(H, W)
|
||||
mock_chmv2_model, mock_chmv2_proc = _make_mock_chmv2(H, W)
|
||||
mock_seg, seg_config_dict = _make_mock_segformer(num_classes, H, W)
|
||||
|
||||
with patch("src.main.load_depth_model", return_value=mock_depth), \
|
||||
patch("src.main.load_chmv2_model",
|
||||
return_value=(mock_chmv2_model, mock_chmv2_proc)), \
|
||||
patch("src.main.load_segmentation_model",
|
||||
return_value=(mock_seg, seg_config_dict)), \
|
||||
patch("src.main.unload_model"):
|
||||
from scripts.run_folder import main
|
||||
main([
|
||||
str(img_path), "--output", str(out),
|
||||
"--image-size", "32", "--num-workers", "0", "--no-vis",
|
||||
])
|
||||
|
||||
st_path = out / "safetensors" / "photo.safetensors"
|
||||
assert st_path.exists(), "Missing consolidated safetensors file"
|
||||
|
||||
data = load_file(str(st_path))
|
||||
assert set(data.keys()) == {"depth", "edge", "chm", "segm"}
|
||||
for key in ("depth", "edge", "chm"):
|
||||
assert data[key].dtype == torch.float16, key
|
||||
assert data[key].shape == (1, H, W), key
|
||||
assert data["segm"].dtype == torch.uint8
|
||||
assert data["segm"].shape == (1, H, W)
|
||||
assert int(data["segm"].max()) <= num_classes - 1
|
||||
# Intermediate npy cleaned up after consolidation.
|
||||
assert list(out.rglob("*.npy")) == []
|
||||
|
||||
def test_rejects_unsupported_extension(self, tmp_path: Path) -> None:
|
||||
bad = tmp_path / "doc.txt"
|
||||
bad.touch()
|
||||
from scripts.run_folder import main
|
||||
with pytest.raises(SystemExit):
|
||||
main([str(bad)])
|
||||
|
||||
def test_rejects_no_vis_with_no_safetensors(self, tmp_path: Path) -> None:
|
||||
# Inference without any output sink is a user error, not a silent no-op.
|
||||
img = tmp_path / "a.png"
|
||||
img.touch()
|
||||
from scripts.run_folder import main
|
||||
with pytest.raises(SystemExit):
|
||||
main([str(img), "--no-vis", "--no-safetensors"])
|
||||
|
||||
def test_rejects_empty_folder(self, tmp_path: Path) -> None:
|
||||
empty = tmp_path / "nothing"
|
||||
empty.mkdir()
|
||||
from scripts.run_folder import main
|
||||
with pytest.raises(SystemExit):
|
||||
main([str(empty)])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (e) Generic discovery: no dataset-specific scene/dir filters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestDiscoverFolderImages:
|
||||
def _touch_png(self, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.touch()
|
||||
|
||||
def test_keeps_dataset_reserved_dir_names(self, tmp_path: Path) -> None:
|
||||
# discover_images drops World-UAV INCOMPLETE_SCENES ('SoHo', ...) and
|
||||
# EXCLUDE_DIRS ('Index', 'charts'); the generic walker must not.
|
||||
from scripts.run_folder import discover_folder_images
|
||||
from src.augmentor.dataset import discover_images
|
||||
|
||||
for rel in ("SoHo/a.png", "Index/b.png", "charts/c.png", "plain/d.png"):
|
||||
self._touch_png(tmp_path / rel)
|
||||
|
||||
generic = discover_folder_images(tmp_path)
|
||||
assert sorted(r.rel_path for r in generic) == sorted(
|
||||
str(Path(p)) for p in
|
||||
("SoHo/a.png", "Index/b.png", "charts/c.png", "plain/d.png")
|
||||
)
|
||||
# Sanity: the dataset walker does filter these out.
|
||||
dataset_recs = discover_images(tmp_path)
|
||||
assert {r.rel_path for r in dataset_recs} == {str(Path("plain/d.png"))}
|
||||
|
||||
def test_skips_generic_service_dirs(self, tmp_path: Path) -> None:
|
||||
from scripts.run_folder import discover_folder_images
|
||||
self._touch_png(tmp_path / "__MACOSX" / "junk.png")
|
||||
self._touch_png(tmp_path / "ok" / "a.png")
|
||||
recs = discover_folder_images(tmp_path)
|
||||
assert [r.rel_path for r in recs] == [str(Path("ok/a.png"))]
|
||||
|
||||
def test_source_filter(self, tmp_path: Path) -> None:
|
||||
from scripts.run_folder import discover_folder_images
|
||||
self._touch_png(tmp_path / "drone" / "q.png")
|
||||
self._touch_png(tmp_path / "satellite" / "s.png")
|
||||
assert [r.rel_path for r in discover_folder_images(tmp_path, source="query")] \
|
||||
== [str(Path("drone/q.png"))]
|
||||
assert [r.rel_path for r in discover_folder_images(tmp_path, source="db")] \
|
||||
== [str(Path("satellite/s.png"))]
|
||||
assert len(discover_folder_images(tmp_path)) == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (f) Resume via consolidated safetensors (save_vis=False runs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestFilterCompletedViaSafetensors:
|
||||
def test_consolidated_modality_counts_as_done(self, tmp_path: Path) -> None:
|
||||
from safetensors.torch import save_file
|
||||
from src.augmentor.dataset import filter_completed
|
||||
|
||||
rec = ImageRecord(
|
||||
abs_path=tmp_path / "a.png",
|
||||
rel_path="a.png",
|
||||
stem="a",
|
||||
output_root=tmp_path / "out",
|
||||
rel_parent=".",
|
||||
)
|
||||
st = safetensors_path(rec.output_root, rec.rel_parent, rec.stem)
|
||||
st.parent.mkdir(parents=True, exist_ok=True)
|
||||
save_file({"depth": torch.zeros(1, 8, 8, dtype=torch.float16)}, str(st))
|
||||
|
||||
# depth present in the consolidated file -> stage skipped ...
|
||||
pending, skipped = filter_completed([rec], "depth")
|
||||
assert pending == [] and skipped == 1
|
||||
# ... but a modality absent from it stays pending.
|
||||
pending, skipped = filter_completed([rec], "chmv2")
|
||||
assert pending == [rec] and skipped == 0
|
||||
@@ -11,10 +11,14 @@ import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import psutil
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
try:
|
||||
import psutil
|
||||
except ImportError: # pragma: no cover — environment-dependent
|
||||
psutil = None # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -51,23 +55,31 @@ def log_system_info() -> dict[str, Any]:
|
||||
info: dict[str, Any] = {}
|
||||
|
||||
# CPU
|
||||
import os
|
||||
cpu_name = platform.processor() or platform.machine()
|
||||
cpu_cores_phys = psutil.cpu_count(logical=False) or 0
|
||||
cpu_cores_logic = psutil.cpu_count(logical=True) or 0
|
||||
if psutil is not None:
|
||||
cpu_cores_phys = psutil.cpu_count(logical=False) or 0
|
||||
cpu_cores_logic = psutil.cpu_count(logical=True) or 0
|
||||
else:
|
||||
cpu_cores_phys = 0
|
||||
cpu_cores_logic = os.cpu_count() or 0
|
||||
info["cpu"] = {
|
||||
"name": cpu_name,
|
||||
"cores_physical": cpu_cores_phys,
|
||||
"cores_logical": cpu_cores_logic,
|
||||
}
|
||||
|
||||
# RAM
|
||||
mem = psutil.virtual_memory()
|
||||
info["ram"] = {
|
||||
"total": mem.total,
|
||||
"available": mem.available,
|
||||
"used": mem.used,
|
||||
"percent": mem.percent,
|
||||
}
|
||||
# RAM (psutil optional — stats unavailable without it).
|
||||
mem = psutil.virtual_memory() if psutil is not None else None
|
||||
if mem is not None:
|
||||
info["ram"] = {
|
||||
"total": mem.total,
|
||||
"available": mem.available,
|
||||
"used": mem.used,
|
||||
"percent": mem.percent,
|
||||
}
|
||||
else:
|
||||
info["ram"] = None
|
||||
|
||||
# GPU
|
||||
if torch.cuda.is_available():
|
||||
@@ -91,8 +103,11 @@ def log_system_info() -> dict[str, Any]:
|
||||
logger.info("🖥️ System info:")
|
||||
logger.info(" 🧠 CPU: %s (%d physical / %d logical cores)",
|
||||
cpu_name, cpu_cores_phys, cpu_cores_logic)
|
||||
logger.info(" 💾 RAM: %s used / %s total (%.1f%% used)",
|
||||
_fmt_bytes(mem.used), _fmt_bytes(mem.total), mem.percent)
|
||||
if mem is not None:
|
||||
logger.info(" 💾 RAM: %s used / %s total (%.1f%% used)",
|
||||
_fmt_bytes(mem.used), _fmt_bytes(mem.total), mem.percent)
|
||||
else:
|
||||
logger.info(" 💾 RAM: psutil not installed — stats unavailable")
|
||||
|
||||
if info["gpu"]:
|
||||
g = info["gpu"]
|
||||
@@ -151,8 +166,10 @@ def log_vram_snapshot(label: str = "") -> dict[str, float] | None:
|
||||
return {"allocated": allocated, "reserved": reserved, "free": free, "total": total}
|
||||
|
||||
|
||||
def log_ram_snapshot(label: str = "") -> dict[str, float]:
|
||||
"""Log current RAM usage."""
|
||||
def log_ram_snapshot(label: str = "") -> dict[str, float] | None:
|
||||
"""Log current RAM usage. Returns dict or None if psutil is missing."""
|
||||
if psutil is None:
|
||||
return None
|
||||
mem = psutil.virtual_memory()
|
||||
prefix = f"[{label}] " if label else ""
|
||||
logger.info(" 💾 %sRAM: %s used / %s available / %s total (%.1f%%)",
|
||||
|
||||
Reference in New Issue
Block a user