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:
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()
|
||||
Reference in New Issue
Block a user