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:
13
in/config_files/pipeline_uav_visloc.gin
Normal file
13
in/config_files/pipeline_uav_visloc.gin
Normal file
@@ -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'
|
||||||
91
scripts/run_uav_visloc.py
Normal file
91
scripts/run_uav_visloc.py
Normal 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()
|
||||||
@@ -41,7 +41,8 @@ class ImageRecord(NamedTuple):
|
|||||||
|
|
||||||
def is_query_record(record: ImageRecord) -> bool:
|
def is_query_record(record: ImageRecord) -> bool:
|
||||||
"""Return True if the record belongs to a query (drone) image."""
|
"""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(
|
def split_by_view(
|
||||||
@@ -102,7 +103,7 @@ def discover_images(
|
|||||||
rel_parts = p.relative_to(root).parts
|
rel_parts = p.relative_to(root).parts
|
||||||
if source == "query" and "DB" in rel_parts:
|
if source == "query" and "DB" in rel_parts:
|
||||||
continue
|
continue
|
||||||
if source == "db" and "query" in rel_parts:
|
if source == "db" and ("query" in rel_parts or "drone" in rel_parts):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
rel = p.relative_to(root)
|
rel = p.relative_to(root)
|
||||||
|
|||||||
Reference in New Issue
Block a user