Initial commit: caption quality test on UAV-VisLoc
Self-contained experimental track validating generated text captions
via retrieval R@1 lift on UAV-VisLoc.
Architecture: GeoRSCLIP ViT-B/32 dual encoder, 512-dim shared space.
Loss: 4-term InfoNCE (img-img + sat-cap + drone-cap + cap-cap)
with cosine temperature decay, PALW-like curriculum.
Metric: delta R@1 (with text - without text) >= +3% => PASS.
Gin-configured (balanced / baseline_no_text / text_heavy variants).
Follows NADEZHDA code style.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
210
scripts/generate_captions.py
Normal file
210
scripts/generate_captions.py
Normal file
@@ -0,0 +1,210 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user