From 6f04856b6b00dd3295f3502ae3e0f7aabd05e636 Mon Sep 17 00:00:00 2001
From: pikaliov <pikaliov@gmail.com>
Date: Fri, 17 Apr 2026 17:18:38 +0300
Subject: [PATCH] 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>
---
 in/config_files/pipeline_uav_visloc.gin | 13 ++++
 scripts/run_uav_visloc.py               | 91 +++++++++++++++++++++++++
 src/augmentor/dataset.py                |  5 +-
 3 files changed, 107 insertions(+), 2 deletions(-)
 create mode 100644 in/config_files/pipeline_uav_visloc.gin
 create mode 100644 scripts/run_uav_visloc.py

diff --git a/in/config_files/pipeline_uav_visloc.gin b/in/config_files/pipeline_uav_visloc.gin
new file mode 100644
index 0000000..99b7c6a
--- /dev/null
+++ b/in/config_files/pipeline_uav_visloc.gin
@@ -0,0 +1,13 @@
+# Pipeline configuration for UAV_VisLoc_processed dataset (74K images)
+PipelineConfig.input_root = '/home/servml/Документы/datasets/UAV_VisLoc_processed'
+PipelineConfig.output_root = '/home/servml/Документы/datasets/UAV_VisLoc_processed-aug'
+PipelineConfig.stages = ['depth', 'edges', 'segmentation', 'chmv2']
+PipelineConfig.save_npy = False
+PipelineConfig.save_vis = True
+PipelineConfig.save_concat = False
+PipelineConfig.save_safetensors = True
+PipelineConfig.cleanup_npy = True
+PipelineConfig.resume = True
+PipelineConfig.subset = None
+PipelineConfig.source = None
+PipelineConfig.log_level = 'INFO'
diff --git a/scripts/run_uav_visloc.py b/scripts/run_uav_visloc.py
new file mode 100644
index 0000000..e7329ca
--- /dev/null
+++ b/scripts/run_uav_visloc.py
@@ -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()
diff --git a/src/augmentor/dataset.py b/src/augmentor/dataset.py
index 8823aa0..7f08825 100644
--- a/src/augmentor/dataset.py
+++ b/src/augmentor/dataset.py
@@ -41,7 +41,8 @@ class ImageRecord(NamedTuple):
 
 def is_query_record(record: ImageRecord) -> bool:
     """Return True if the record belongs to a query (drone) image."""
-    return "query" in Path(record.rel_path).parts
+    parts = Path(record.rel_path).parts
+    return "query" in parts or "drone" in parts
 
 
 def split_by_view(
@@ -102,7 +103,7 @@ def discover_images(
             rel_parts = p.relative_to(root).parts
             if source == "query" and "DB" in rel_parts:
                 continue
-            if source == "db" and "query" in rel_parts:
+            if source == "db" and ("query" in rel_parts or "drone" in rel_parts):
                 continue
 
         rel = p.relative_to(root)