Add GTA-UAV-LR annotation script + dataset documentation
- Add scripts/run_gta_uav.py for GTA-UAV-LR (48K images, GTA V synthetic) - 14 segmentation classes: 11 base + bare soil, rooftop, swimming pool - Fix source filter to recognize "satellite" folder (alongside "DB") - Document GTA-UAV characteristics in segmentation_class_analysis.md - Update README and CLAUDE.md with GTA-UAV support Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
116
scripts/run_gta_uav.py
Normal file
116
scripts/run_gta_uav.py
Normal file
@@ -0,0 +1,116 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run annotation pipeline for GTA-UAV-LR dataset.
|
||||
|
||||
GTA-UAV-LR: synthetic dataset from GTA V engine.
|
||||
- drone/images/: 33763 images, 512x384, RGB PNG
|
||||
- satellite/: 14640 images, 256x256, RGBA PNG (alpha = map boundary)
|
||||
- Total: 48403 images
|
||||
- 6 flight heights: 100, 200, 300, 400, 500, 600 meters
|
||||
|
||||
Usage:
|
||||
python scripts/run_gta_uav.py
|
||||
python scripts/run_gta_uav.py --source db # only satellite (14.6K)
|
||||
python scripts/run_gta_uav.py --source drone # only drone (33.8K)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_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/GTA-UAV-LR"
|
||||
OUTPUT_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-aug"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Annotate GTA-UAV-LR")
|
||||
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
|
||||
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,
|
||||
num_workers=4,
|
||||
)
|
||||
|
||||
# GTA-UAV: satellite 256x256, drone 512x384
|
||||
# Use 256 for satellite, 512 for drone (non-square → resize to square)
|
||||
input_conf = InputConfig(image_size=256, query_image_size=512)
|
||||
|
||||
# GTA V synthetic scenes: urban, suburban, rural, coastal, mountainous
|
||||
# 11 base classes + pool (swimming pools common in GTA suburbs)
|
||||
seg_conf = SegConfig(threshold=0.15, prompts=[
|
||||
"background", # 0
|
||||
"building", # 1
|
||||
"road", # 2
|
||||
"vegetation", # 3
|
||||
"water", # 4
|
||||
"sand and gravel ground", # 5
|
||||
"rocky terrain", # 6
|
||||
"farmland", # 7
|
||||
"railway", # 8
|
||||
"parking lot", # 9
|
||||
"sidewalk", # 10
|
||||
"bare soil and plowed field", # 11
|
||||
"roof and rooftop", # 12
|
||||
"swimming pool", # 13
|
||||
])
|
||||
|
||||
models_conf = ModelsConfig(weights_dir=str(_PROJECT_ROOT / "in" / "weights"))
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user