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>

src/datasets/__init__.py

@@ -0,0 +1,7 @@
+"""Dataset loaders for caption quality test."""
+from src.datasets.visloc_with_captions import (
+    VisLocCaptionDataset,
+    collate_caption_batch,
+)
+
+__all__ = ["VisLocCaptionDataset", "collate_caption_batch"]

src/datasets/visloc_with_captions.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+"""UAV-VisLoc dataset loader augmented with generated captions.
+
+Expects a manifest JSON of the form:
+    [
+        {
+            "pair_id": "v001_0042",
+            "drone_path": "drone/v001_0042.jpg",
+            "sat_path": "satellite/v001_0042.png",
+            "caption_drone": "low-altitude photo of residential ...",
+            "caption_sat": "aerial view of urban area ...",
+            "gps": [lat, lon]
+        },
+        ...
+    ]
+
+Captions are produced offline by scripts/generate_captions.py using one of
+three strategies: template, VLM, or hybrid (see АНАЛИЗ_caption_quality_test).
+"""
+
+import json
+import random
+from pathlib import Path
+from typing import Any, Callable
+
+import gin
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+@gin.configurable
+class VisLocCaptionDataset(Dataset):
+    """UAV-VisLoc pairs with generated captions.
+
+    Args:
+        manifest_path: Path to JSON manifest with pair entries.
+        image_root: Directory prefix joined with manifest relative paths.
+        image_transform: Callable applied to PIL images (e.g., GeoRSCLIP preprocess).
+        caption_strategy: Which caption field to use ('template', 'vlm', 'hybrid').
+            The corresponding field must exist in the manifest
+            (e.g., 'caption_sat_vlm', or the generic 'caption_sat').
+        drop_caption_prob: Random probability of replacing a caption with ''.
+            Useful for dropout ablations during training.
+        seed: Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        manifest_path: str,
+        image_root: str,
+        image_transform: Callable[[Image.Image], torch.Tensor],
+        caption_strategy: str = "hybrid",
+        drop_caption_prob: float = 0.0,
+        seed: int = 0,
+    ) -> None:
+        self.manifest_path = Path(manifest_path)
+        self.image_root = Path(image_root)
+        self.image_transform = image_transform
+        self.caption_strategy = caption_strategy
+        self.drop_caption_prob = drop_caption_prob
+        self._rng = random.Random(seed)
+
+        with self.manifest_path.open("r", encoding="utf-8") as f:
+            self.entries: list[dict[str, Any]] = json.load(f)
+
+        self._validate_entries()
+
+    def _validate_entries(self) -> None:
+        """Ensure all entries have required fields for the chosen strategy."""
+        required = {"drone_path", "sat_path"}
+        caption_sat_key = self._caption_key("sat")
+        caption_drone_key = self._caption_key("drone")
+        required |= {caption_sat_key, caption_drone_key}
+
+        for i, entry in enumerate(self.entries):
+            missing = required - entry.keys()
+            if missing:
+                raise KeyError(
+                    f"Entry {i} (pair_id={entry.get('pair_id', '?')}) missing fields: "
+                    f"{sorted(missing)}"
+                )
+
+    def _caption_key(self, view: str) -> str:
+        """Resolve caption field name from strategy + view."""
+        if self.caption_strategy == "hybrid":
+            return f"caption_{view}"
+        return f"caption_{view}_{self.caption_strategy}"
+
+    def _load_image(self, relative_path: str) -> torch.Tensor:
+        """Load image and apply preprocessing."""
+        path = self.image_root / relative_path
+        with Image.open(path) as img:
+            rgb = img.convert("RGB")
+        return self.image_transform(rgb)
+
+    def _maybe_drop(self, caption: str) -> str:
+        """Stochastically drop caption to empty string for robustness training."""
+        if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
+            return ""
+        return caption
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Return one pair with images and captions.
+
+        Args:
+            idx: Index into the manifest.
+
+        Returns:
+            Dict with:
+                - 'drone_img': [3, H, W] tensor
+                - 'sat_img': [3, H, W] tensor
+                - 'caption_drone': str (possibly empty)
+                - 'caption_sat': str (possibly empty)
+                - 'pair_id': str for logging
+        """
+        entry = self.entries[idx]
+
+        drone_img = self._load_image(entry["drone_path"])
+        sat_img = self._load_image(entry["sat_path"])
+
+        caption_drone = self._maybe_drop(entry[self._caption_key("drone")])
+        caption_sat = self._maybe_drop(entry[self._caption_key("sat")])
+
+        return {
+            "drone_img": drone_img,
+            "sat_img": sat_img,
+            "caption_drone": caption_drone,
+            "caption_sat": caption_sat,
+            "pair_id": entry.get("pair_id", f"idx_{idx}"),
+        }
+
+
+def collate_caption_batch(
+    batch: list[dict[str, Any]],
+) -> dict[str, Any]:
+    """Collate VisLocCaptionDataset items into a batched dict.
+
+    Images are stacked; captions remain Python lists so the tokenizer can
+    process them inside the model.forward().
+
+    Args:
+        batch: List of samples from VisLocCaptionDataset.__getitem__.
+
+    Returns:
+        Batched dict with stacked image tensors and caption lists.
+    """
+    return {
+        "drone_img": torch.stack([b["drone_img"] for b in batch], dim=0),
+        "sat_img": torch.stack([b["sat_img"] for b in batch], dim=0),
+        "caption_drone": [b["caption_drone"] for b in batch],
+        "caption_sat": [b["caption_sat"] for b in batch],
+        "pair_ids": [b["pair_id"] for b in batch],
+    }