Add UAV_VisLoc annotation script + support drone folder naming

- Add scripts/run_uav_visloc.py for UAV_VisLoc_processed dataset (81K images)
- Add gin config for UAV_VisLoc (pipeline_uav_visloc.gin)
- Fix source filter and is_query_record to recognize "drone" folder
  (UAV_VisLoc uses "drone" instead of "query")

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-17 17:18:38 +03:00
parent 13ff079891
commit 6f04856b6b
3 changed files with 107 additions and 2 deletions

91
scripts/run_uav_visloc.py Normal file
View File

@@ -0,0 +1,91 @@
#!/usr/bin/env python3
"""Run annotation pipeline for UAV_VisLoc_processed dataset.
Usage:
python scripts/run_uav_visloc.py
python scripts/run_uav_visloc.py --source db # only satellite
python scripts/run_uav_visloc.py --source drone # only drone
"""
from __future__ import annotations
import argparse
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.io_utils import setup_logging
from src.main import run_pipeline
INPUT_ROOT = "/home/servml/Документы/datasets/UAV_VisLoc_processed"
OUTPUT_ROOT = "/home/servml/Документы/datasets/UAV_VisLoc_processed-aug"
def main() -> None:
parser = argparse.ArgumentParser(description="Annotate UAV_VisLoc_processed")
parser.add_argument("--source", choices=["db", "drone", "all"], default="all",
help="Process only db (satellite), drone, or all (default)")
parser.add_argument("--stages", nargs="+",
default=["depth", "edges", "segmentation", "chmv2"],
help="Stages to run")
args = parser.parse_args()
import gin
gin.clear_config()
source = None if args.source == "all" else args.source
# UAV_VisLoc uses "drone" folder, not "query". Map for filter compatibility.
if source == "drone":
source = "query"
pipeline_conf = PipelineConfig(
input_root=INPUT_ROOT,
output_root=OUTPUT_ROOT,
stages=args.stages,
save_npy=False,
save_vis=True,
save_safetensors=True,
cleanup_npy=True,
resume=True,
source=source,
log_level="INFO",
)
hw_conf = HardwareConfig(
profile_name="rtx4090",
total_ram_gb=24.0,
reserve_gb=2.0,
use_fp16=True,
batch_size=None, # auto
num_workers=4,
)
input_conf = InputConfig(image_size=256)
seg_conf = SegConfig()
models_conf = ModelsConfig()
setup_logging(
pipeline_conf.log_level,
log_file=Path(OUTPUT_ROOT) / "pipeline.log",
)
torch.manual_seed(42)
np.random.seed(42)
run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf)
if __name__ == "__main__":
main()