from __future__ import annotations """Offline caption generation for UAV-VisLoc images. Supports three strategies: - template: rule-based from SegEarth-OV3 masks (fastest, generic). - vlm: Qwen2.5-VL / InternVL2 VLM (slowest, most diverse). - hybrid: template first, VLM refinement on 10% sample (balanced). Writes a manifest JSON that is directly consumable by VisLocCaptionDataset. Usage: python -m scripts.generate_captions \ --image_root data/visloc/images \ --pairs_csv data/visloc/pairs.csv \ --output data/visloc_train.json \ --strategy hybrid \ --vlm_refine_ratio 0.1 """ import argparse import json import logging import random from pathlib import Path LOGGER = logging.getLogger("caption_test.generate") _TEMPLATE_SAT_PATTERNS = [ "aerial satellite view of {area_type} with {feature_1} and {feature_2}", "orthogonal satellite image showing {area_type}, visible {feature_1}", "top-down satellite photo of {area_type}", ] _TEMPLATE_DRONE_PATTERNS = [ "low-altitude drone photo of {area_type} with {feature_1} and {feature_2}", "oblique UAV view of {area_type}, showing {feature_1}", "aerial drone image of {area_type}", ] def _template_caption( view: str, area_type: str, features: list[str], rng: random.Random, ) -> str: """Generate rule-based caption from semantic masks.""" patterns = _TEMPLATE_SAT_PATTERNS if view == "sat" else _TEMPLATE_DRONE_PATTERNS pattern = rng.choice(patterns) feat_1 = features[0] if len(features) > 0 else "varied terrain" feat_2 = features[1] if len(features) > 1 else "natural features" return pattern.format(area_type=area_type, feature_1=feat_1, feature_2=feat_2) def _placeholder_vlm_caption(image_path: Path, view: str) -> str: """Placeholder for VLM caption. Replace with real Qwen2.5-VL inference. Returns a short deterministic caption for smoke-testing pipelines. """ # TODO: integrate Qwen2.5-VL / InternVL2 inference here. if view == "sat": return f"satellite aerial view (placeholder for {image_path.name})" return f"drone low-altitude view (placeholder for {image_path.name})" def _parse_pairs_csv(pairs_csv: Path) -> list[dict]: """Load pair metadata (drone_path, sat_path, gps, optional masks).""" import csv entries: list[dict] = [] with pairs_csv.open("r", encoding="utf-8", newline="") as f: reader = csv.DictReader(f) for row in reader: entries.append( { "pair_id": row.get("pair_id", f"pair_{len(entries)}"), "drone_path": row["drone_path"], "sat_path": row["sat_path"], "gps": [float(row.get("lat", 0.0)), float(row.get("lon", 0.0))], "area_type": row.get("area_type", "mixed terrain"), "features": [ s.strip() for s in row.get("features", "").split(";") if s.strip() ], } ) return entries def build_manifest( image_root: Path, pairs_csv: Path, output_path: Path, strategy: str, vlm_refine_ratio: float, seed: int, ) -> None: """Build a manifest JSON with captions for all pairs. Args: image_root: Directory prefix for images. pairs_csv: CSV with pair metadata (drone_path, sat_path, ...). output_path: Output JSON path. strategy: 'template' | 'vlm' | 'hybrid'. vlm_refine_ratio: Fraction to refine with VLM when strategy='hybrid'. seed: Random seed. """ rng = random.Random(seed) entries = _parse_pairs_csv(pairs_csv) LOGGER.info("loaded %d pairs from %s", len(entries), pairs_csv) manifest: list[dict] = [] n_vlm_refined = 0 for i, entry in enumerate(entries): area_type = entry["area_type"] features = entry["features"] # Template captions cap_sat_tpl = _template_caption("sat", area_type, features, rng) cap_drone_tpl = _template_caption("drone", area_type, features, rng) # VLM captions (optional, based on strategy) use_vlm = False if strategy == "vlm": use_vlm = True elif strategy == "hybrid" and rng.random() < vlm_refine_ratio: use_vlm = True if use_vlm: cap_sat_vlm = _placeholder_vlm_caption( image_root / entry["sat_path"], "sat" ) cap_drone_vlm = _placeholder_vlm_caption( image_root / entry["drone_path"], "drone" ) n_vlm_refined += 1 else: cap_sat_vlm = cap_sat_tpl cap_drone_vlm = cap_drone_tpl # Final hybrid caption prefers VLM when present. final_sat = cap_sat_vlm if use_vlm else cap_sat_tpl final_drone = cap_drone_vlm if use_vlm else cap_drone_tpl manifest.append( { "pair_id": entry["pair_id"], "drone_path": entry["drone_path"], "sat_path": entry["sat_path"], "gps": entry["gps"], # Strategy-specific captions (for ablations). "caption_sat_template": cap_sat_tpl, "caption_drone_template": cap_drone_tpl, "caption_sat_vlm": cap_sat_vlm, "caption_drone_vlm": cap_drone_vlm, # Generic 'hybrid' fields used by default dataset. "caption_sat": final_sat, "caption_drone": final_drone, } ) if (i + 1) % 1000 == 0: LOGGER.info("processed %d / %d pairs", i + 1, len(entries)) output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8") as f: json.dump(manifest, f, indent=2, ensure_ascii=False) LOGGER.info( "wrote %d entries to %s (%d VLM-refined, strategy=%s)", len(manifest), output_path, n_vlm_refined, strategy, ) def main() -> None: parser = argparse.ArgumentParser(description="Generate captions for UAV-VisLoc.") parser.add_argument("--image_root", type=Path, required=True) parser.add_argument("--pairs_csv", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument( "--strategy", choices=["template", "vlm", "hybrid"], default="hybrid", ) parser.add_argument("--vlm_refine_ratio", type=float, default=0.1) parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() logging.basicConfig( level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s", ) build_manifest( image_root=args.image_root, pairs_csv=args.pairs_csv, output_path=args.output, strategy=args.strategy, vlm_refine_ratio=args.vlm_refine_ratio, seed=args.seed, ) if __name__ == "__main__": main()