Improve training: learnable temperature, per-group LR, warmup, augmentations

Loss:
- Learnable temperature (CLIP-style logit_scale) with clamp [0.01, 0.5]
- Replaces fixed cosine schedule (still available via --no-learnable-temp)
- Default tau_init=0.07

Optimizer:
- Per-group LR: projections 1e-4, text encoder 1e-5 (10x lower)
- Learnable temperature included in projection param group

Scheduler:
- Linear warmup (2 epochs default) + cosine annealing
- Per-step scheduling (not per-epoch)

Augmentations (separate drone/satellite):
- Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15), ColorJitter,
  RandomGrayscale(0.05), GaussianBlur
- Satellite: RandomResizedCrop(0.7-1.0), HFlip, ColorJitter, RandomGrayscale
- Eval: clean Resize+CenterCrop (no augmentation)

Dataset: supports separate drone_transform/sat_transform args

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

src/datasets/gtauav_dataset.py

@@ -99,7 +99,9 @@ class GTAUAVDataset(Dataset):
         pair_json: Path to cross-area-drone2sate-{train,test}.json.
         rgb_root: Root of GTA-UAV-LR RGB images.
         caption_root: Root of GTA-UAV-LR-captions.
-        image_transform: Callable applied to PIL images.
+        drone_transform: Transform for drone images (can include augmentations).
+        sat_transform: Transform for satellite images (can include augmentations).
+        image_transform: Fallback single transform for both (used if drone/sat not set).
         filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
         drop_caption_prob: Probability of dropping captions (ablation).
         seed: Random seed.
@@ -110,6 +112,8 @@ class GTAUAVDataset(Dataset):
         pair_json: str,
         rgb_root: str = str(_RGB_ROOT),
         caption_root: str = str(_CAPTION_ROOT),
+        drone_transform: Callable[[Image.Image], torch.Tensor] | None = None,
+        sat_transform: Callable[[Image.Image], torch.Tensor] | None = None,
         image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
         filter_meta: str | None = None,
         drop_caption_prob: float = 0.0,
@@ -117,7 +121,8 @@ class GTAUAVDataset(Dataset):
     ) -> None:
         self.rgb_root = Path(rgb_root)
         self.caption_root = Path(caption_root)
-        self.image_transform = image_transform
+        self.drone_transform = drone_transform or image_transform
+        self.sat_transform = sat_transform or image_transform
         self.drop_caption_prob = drop_caption_prob
         self._rng = random.Random(seed)
 
@@ -189,12 +194,12 @@ class GTAUAVDataset(Dataset):
                 "caption_l3": l3,
             })
 
-    def _load_image(self, directory: str, filename: str) -> torch.Tensor:
+    def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
         path = self.rgb_root / directory / filename
         with Image.open(path) as img:
             rgb = img.convert("RGB")
-        if self.image_transform is not None:
-            return self.image_transform(rgb)
+        if transform is not None:
+            return transform(rgb)
         return torch.tensor(0)  # placeholder if no transform
 
     def __len__(self) -> int:
@@ -203,7 +208,7 @@ class GTAUAVDataset(Dataset):
     def __getitem__(self, idx: int) -> dict[str, Any]:
         entry = self.entries[idx]
 
-        drone_img = self._load_image(entry["drone_dir"], entry["drone_name"])
+        drone_img = self._load_image(entry["drone_dir"], entry["drone_name"], self.drone_transform)
 
         # Sample satellite match (weighted if semi-positive).
         if entry["sat_weights"] is not None:
@@ -215,7 +220,7 @@ class GTAUAVDataset(Dataset):
         else:
             sat_name = self._rng.choice(entry["sat_candidates"])
 
-        sat_img = self._load_image(entry["sat_dir"], sat_name)
+        sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform)
 
         # Captions with optional dropout.
         if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob: