From 467e5fc97660bb3be135fb07ffb136edebc58403 Mon Sep 17 00:00:00 2001 From: Pikaliov Date: Sat, 11 Jul 2026 18:36:53 +0300 Subject: [PATCH] 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. --- README.md | 21 ++- scripts/run_folder.py | 235 ++++++++++++++++++++++++++ src/augmentor/dataset.py | 25 ++- src/main.py | 32 ++-- src/tests/test_run_folder.py | 309 +++++++++++++++++++++++++++++++++++ src/utils/profiler.py | 47 ++++-- 6 files changed, 641 insertions(+), 28 deletions(-) create mode 100644 scripts/run_folder.py create mode 100644 src/tests/test_run_folder.py diff --git a/README.md b/README.md index aab051f..94bdb5e 100644 --- a/README.md +++ b/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 +# Папка: выход по умолчанию — сиблинг -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 классов) diff --git a/scripts/run_folder.py b/scripts/run_folder.py new file mode 100644 index 0000000..0629af8 --- /dev/null +++ b/scripts/run_folder.py @@ -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 '-aug'; " + "for a single file: '-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() diff --git a/src/augmentor/dataset.py b/src/augmentor/dataset.py index 793f3a6..ce9b689 100644 --- a/src/augmentor/dataset.py +++ b/src/augmentor/dataset.py @@ -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 diff --git a/src/main.py b/src/main.py index 2cd981d..f43f4be 100644 --- a/src/main.py +++ b/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.") diff --git a/src/tests/test_run_folder.py b/src/tests/test_run_folder.py new file mode 100644 index 0000000..28dae39 --- /dev/null +++ b/src/tests/test_run_folder.py @@ -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 /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 diff --git a/src/utils/profiler.py b/src/utils/profiler.py index 428723c..651185d 100644 --- a/src/utils/profiler.py +++ b/src/utils/profiler.py @@ -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%%)",