caption-test/scripts/generate_captions.py

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()