Fix GTA-UAV eval + training pipeline: full gallery, mutex sampler, per-sample mask

Six critical fixes to the caption-test training/eval stack:

1. **IndentationError blocker** (train_gtauav.py:765-766)
   Unparseable file — train-recall LOGGER.info block was orphaned outside
   its `if eval_every` guard. Wrapped in `if train_recall:` so val eval
   and Grad-CAM only run on eval epochs.

2. **Full satellite gallery in `_evaluate`**
   Old code assembled gallery from DataLoader batches (one random sat per
   drone), producing an incomplete gallery of size ≈ N_query instead of
   N_unique_sat. Metrics were inflated because retrieval was against a
   subset that always contained the target.
   New `GTAUAVSatGallery` / `GTAUAVDroneQuery` iterate all unique tiles
   and queries independently; full-gallery multi-match R@K + MRR.

3. **Per-sample caption mask** (`AsymmetricEncoder._fuse_with_mask`)
   Mixed batches (some samples have captions, some don't) previously
   encoded empty strings through DGTRS and mixed the noise output into
   every sample via scalar gate. New `encode_query`/`encode_gallery` use
   `torch.where` to fall back to pure image features for empty-caption
   samples. Training `forward()` routes through the same helper so
   training and eval share code.

4. **Symmetric InfoNCE as primary loss** (multi_infonce.InfoNCELoss)
   Switched gin default from `WeightedInfoNCELoss` (adaptive label
   smoothing — not the Game4Loc soft-IoU target it claimed) to the
   existing symmetric InfoNCE with q2g=0.6/g2q=0.4 weighting. Loss type
   now selectable via `cfg.loss_type ∈ {"symmetric", "weighted"}`.

5. **MutuallyExclusiveSampler** (new file)
   BatchSampler that greedily packs drones whose `sat_candidates` sets
   are pairwise disjoint within a batch. Eliminates false negatives from
   the semi-positive graph without needing soft-label losses.
   At bs=8 keeps 100% of 24,891 train entries; at bs=64 keeps 92.6%.
   `set_epoch()` for reproducibility + different batches per epoch.

6. **Temperature clamp [0.01, 0.1]** (both loss modules)
   Old tau_max=0.5 allowed the logit distribution to collapse into a
   near-uniform softmax. Tightened to the CLIP-standard range.

Also:
- Added `scripts/smoke_eval.py` / `scripts/smoke_train.py` for fast
  regression checks (eval in ~2 min, 2 train steps in ~1 min on RTX 4090).
- CLAUDE.md updated to reflect the new pipeline.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/datasets/gtauav_dataset.py

@@ -287,3 +287,109 @@ def collate_gtauav_batch(
         "sat_names": [b["sat_name"] for b in batch],
         "positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32),
     }
+
+
+def _load_rgb_image(
+    rgb_root: Path,
+    directory: str,
+    filename: str,
+    transform: Callable | None,
+) -> torch.Tensor:
+    path = rgb_root / directory / filename
+    with Image.open(path) as img:
+        rgb = img.convert("RGB")
+    if transform is not None:
+        return transform(rgb)
+    return torch.tensor(0)
+
+
+class GTAUAVSatGallery(Dataset):
+    """Unique satellite gallery for retrieval evaluation.
+
+    Takes a GTAUAVDataset and extracts the set of unique satellite names
+    appearing in any entry's sat_candidates. Yields one (sat_img, captions)
+    per unique name — suitable as the gallery side of a retrieval benchmark.
+
+    Used by `_evaluate` in train_gtauav.py to forward the full gallery once.
+    """
+
+    def __init__(self, source: "GTAUAVDataset") -> None:
+        self.rgb_root = source.rgb_root
+        self.sat_transform = source.sat_transform
+        # Collect unique sats (preserve first-seen order for determinism).
+        unique: dict[str, tuple[str, tuple[str, str, str] | None]] = {}
+        for entry in source.entries:
+            sat_dir = entry["sat_dir"]
+            for sat_name in entry["sat_candidates"]:
+                if sat_name in unique:
+                    continue
+                caps = entry["sat_captions"].get(sat_name)
+                unique[sat_name] = (sat_dir, caps)
+        self.sat_names: list[str] = list(unique.keys())
+        self._sat_info: dict[str, tuple[str, tuple[str, str, str] | None]] = unique
+
+    def __len__(self) -> int:
+        return len(self.sat_names)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        sat_name = self.sat_names[idx]
+        sat_dir, caps = self._sat_info[sat_name]
+        sat_img = _load_rgb_image(self.rgb_root, sat_dir, sat_name, self.sat_transform)
+        if caps is not None:
+            l1, l2, l3 = caps
+        else:
+            l1 = l2 = l3 = _EMPTY_CAPTION
+        return {
+            "sat_img": sat_img,
+            "sat_name": sat_name,
+            "sat_caption_l1": l1,
+            "sat_caption_l2": l2,
+            "sat_caption_l3": l3,
+        }
+
+
+class GTAUAVDroneQuery(Dataset):
+    """Drone queries with valid satellite names for multi-match evaluation."""
+
+    def __init__(self, source: "GTAUAVDataset") -> None:
+        self.rgb_root = source.rgb_root
+        self.drone_transform = source.drone_transform
+        self.entries = source.entries
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        entry = self.entries[idx]
+        drone_img = _load_rgb_image(
+            self.rgb_root, entry["drone_dir"], entry["drone_name"], self.drone_transform,
+        )
+        return {
+            "drone_img": drone_img,
+            "drone_name": entry["drone_name"],
+            "caption_l1": entry["caption_l1"],
+            "caption_l2": entry["caption_l2"],
+            "caption_l3": entry["caption_l3"],
+            "valid_sat_names": list(entry["sat_candidates"]),
+        }
+
+
+def collate_sat_gallery(batch: list[dict[str, Any]]) -> dict[str, Any]:
+    return {
+        "sat_img": torch.stack([b["sat_img"] for b in batch], dim=0),
+        "sat_names": [b["sat_name"] for b in batch],
+        "sat_caption_l1": [b["sat_caption_l1"] for b in batch],
+        "sat_caption_l2": [b["sat_caption_l2"] for b in batch],
+        "sat_caption_l3": [b["sat_caption_l3"] for b in batch],
+    }
+
+
+def collate_drone_query(batch: list[dict[str, Any]]) -> dict[str, Any]:
+    return {
+        "drone_img": torch.stack([b["drone_img"] for b in batch], dim=0),
+        "drone_names": [b["drone_name"] for b in batch],
+        "caption_l1": [b["caption_l1"] for b in batch],
+        "caption_l2": [b["caption_l2"] for b in batch],
+        "caption_l3": [b["caption_l3"] for b in batch],
+        "valid_sat_names": [b["valid_sat_names"] for b in batch],
+    }

src/datasets/mutually_exclusive_sampler.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+"""Batch sampler that prevents false negatives from GTA-UAV's semi-positive graph.
+
+In GTA-UAV a single satellite tile can be a valid (semi-)positive for multiple
+drone frames. When those frames land in the same InfoNCE batch, the shared
+satellite becomes a negative for every drone except one — which trains the
+model to push apart embeddings that should actually be close.
+
+MutuallyExclusiveSampler resolves this by greedily building batches where no
+two drone indices share ANY entry in their `sat_candidates` set. This keeps
+the diagonal InfoNCE formulation valid: every off-diagonal satellite is a
+genuine negative for every query in the row/column.
+
+References: Zhu et al., Game4Loc (arXiv:2409.16925), §3.2 "Sample ID".
+"""
+
+import logging
+import random
+from typing import Iterator, Sequence
+
+from torch.utils.data.sampler import Sampler
+
+LOGGER = logging.getLogger("caption_test.mutex_sampler")
+
+
+class MutuallyExclusiveSampler(Sampler[list[int]]):
+    """Batch sampler yielding drone-index lists with disjoint sat_candidates.
+
+    Args:
+        sat_candidates_per_item: For each dataset index i, the list of
+            satellite names considered valid matches (positive + semi-positive).
+        batch_size: Target batch size. Partial batches are dropped.
+        shuffle: Shuffle the index pool each epoch (use set_epoch for reproducibility).
+        seed: Base RNG seed — the effective seed is `seed + epoch`.
+        allow_partial: If True, yield the trailing partial batch. Default False.
+    """
+
+    def __init__(
+        self,
+        sat_candidates_per_item: Sequence[Sequence[str]],
+        batch_size: int,
+        shuffle: bool = True,
+        seed: int = 0,
+        allow_partial: bool = False,
+    ) -> None:
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {batch_size}")
+        self._item_sats: list[frozenset[str]] = [
+            frozenset(s) for s in sat_candidates_per_item
+        ]
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.seed = seed
+        self.allow_partial = allow_partial
+        self.epoch = 0
+        self._cached_len: int | None = None
+
+    def set_epoch(self, epoch: int) -> None:
+        """Advance epoch (invalidates cached length) — call from training loop."""
+        self.epoch = epoch
+        self._cached_len = None
+
+    def _generate_batches(self) -> list[list[int]]:
+        rng = random.Random(self.seed + self.epoch) if self.shuffle else None
+        remaining = list(range(len(self._item_sats)))
+        if rng is not None:
+            rng.shuffle(remaining)
+
+        batches: list[list[int]] = []
+
+        # Each outer iteration produces at most one batch. Items that conflict
+        # with the batch's claimed-sat set roll over to the next iteration.
+        while remaining:
+            batch: list[int] = []
+            claimed: set[str] = set()
+            next_remaining: list[int] = []
+
+            for idx in remaining:
+                sats = self._item_sats[idx]
+                if len(batch) < self.batch_size and not (sats & claimed):
+                    batch.append(idx)
+                    claimed |= sats
+                else:
+                    next_remaining.append(idx)
+
+            if len(batch) == self.batch_size:
+                batches.append(batch)
+            elif self.allow_partial and batch:
+                batches.append(batch)
+                break  # leftover rollover produces no more full batches
+            else:
+                break  # drop partial trailing batch
+
+            remaining = next_remaining
+            if rng is not None:
+                rng.shuffle(remaining)
+
+        return batches
+
+    def __iter__(self) -> Iterator[list[int]]:
+        batches = self._generate_batches()
+        self._cached_len = len(batches)
+        for batch in batches:
+            yield batch
+
+    def __len__(self) -> int:
+        if self._cached_len is None:
+            # Estimate by actually generating — correct count needed by
+            # DataLoader/tqdm. Cached until next set_epoch.
+            self._cached_len = len(self._generate_batches())
+        return self._cached_len