Rewrite: GatedFusion architecture + UAV-GeoLoc dataset
Architecture v2: - Query branch: drone + text -> GatedFusion -> proj -> query_emb - Gallery branch: satellite -> proj -> gallery_emb - Single InfoNCE loss (asymmetric 0.6/0.4) - GatedFusion: learnable gated addition (sigma(alpha)*img + (1-sigma(alpha))*text) - Baseline mode: gate=1.0 (text ignored) Dataset: - UAV-GeoLoc loader with template captions from path metadata - 27 terrain types with predefined features - Random positive crop sampling per epoch Configs: balanced (gate=0.7), baseline (no text), text_heavy (gate=0.3) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
@@ -1,26 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""UAV-VisLoc dataset loader augmented with generated captions.
|
||||
"""UAV-GeoLoc dataset loader with template captions for CVGL caption test.
|
||||
|
||||
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]
|
||||
},
|
||||
...
|
||||
]
|
||||
Reads UAV-GeoLoc Index format (train_query.txt / train_db.txt) and generates
|
||||
template captions from path metadata (terrain type, scene, altitude, heading).
|
||||
|
||||
Captions are produced offline by scripts/generate_captions.py using one of
|
||||
three strategies: template, VLM, or hybrid (see АНАЛИЗ_caption_quality_test).
|
||||
train_query.txt format:
|
||||
Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg 0 .../DB/img/crop_X_Y.png ...
|
||||
|
||||
Template caption (P3-style fingerprint for cross-view matching):
|
||||
"Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano.
|
||||
Plan-view features: dark lava flows, crater edges, sparse vegetation patches."
|
||||
"""
|
||||
|
||||
import json
|
||||
import random
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
@@ -29,130 +23,191 @@ import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
# Compass lookup: heading_deg -> direction name.
|
||||
_COMPASS = ["north", "northeast", "east", "southeast",
|
||||
"south", "southwest", "west", "northwest"]
|
||||
|
||||
# Simple terrain-to-features mapping for template captions.
|
||||
_TERRAIN_FEATURES: dict[str, str] = {
|
||||
"Volcano": "lava flows, crater edges, volcanic rock",
|
||||
"Mountain": "ridgelines, steep slopes, rocky terrain",
|
||||
"Hill": "rolling hills, gentle slopes, scattered trees",
|
||||
"Desert": "sand dunes, arid ground, sparse scrub",
|
||||
"Plain": "flat open fields, agricultural plots, dirt roads",
|
||||
"Plateau": "flat elevated terrain, cliffs, mesa edges",
|
||||
"Basin": "lowland depression, dry lake bed, sediment patterns",
|
||||
"Delta": "river channels, sediment fans, wetland patches",
|
||||
"Gorge": "deep canyon, exposed rock layers, narrow valley",
|
||||
"Island": "shoreline, coastal vegetation, water boundary",
|
||||
"Wetland": "marshes, water channels, aquatic vegetation",
|
||||
"Glacier": "ice formations, crevasses, glacial moraine",
|
||||
"Forest": "dense canopy, tree shadows, clearings",
|
||||
"Farm": "crop fields, irrigation lines, farm buildings",
|
||||
"Prairie": "grassland, open meadow, fence lines",
|
||||
"Finca": "rural estate, orchards, scattered structures",
|
||||
"Calcification": "mineral deposits, white crusts, barren soil",
|
||||
"StoneForest": "karst pillars, eroded limestone, sparse vegetation",
|
||||
"Oasis": "palm clusters, water pool, surrounding desert",
|
||||
"Flowers": "colorful ground cover, floral patterns, open field",
|
||||
"Terrace": "terraced hillside, stepped fields, retaining walls",
|
||||
"Snow": "snow cover, tracks, exposed rock patches",
|
||||
"Pasture": "grazing land, fenced paddocks, grass",
|
||||
"Danxia": "red sandstone, layered cliffs, erosion patterns",
|
||||
"Hylare": "sparse woodland, rocky outcrops",
|
||||
"Karst": "sinkholes, limestone towers, caves",
|
||||
"Fall": "autumn foliage, colored canopy, leaf litter",
|
||||
}
|
||||
|
||||
# Country/city features.
|
||||
_COUNTRY_FEATURES: str = "buildings, roads, urban blocks, rooftops, intersections"
|
||||
|
||||
|
||||
def _parse_query_path(query_path: str) -> dict[str, str]:
|
||||
"""Extract metadata from UAV-GeoLoc query path.
|
||||
|
||||
Example: Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg
|
||||
Returns: {category, terrain_type, scene, height_m, heading_deg, heading_dir}
|
||||
"""
|
||||
parts = query_path.split("/")
|
||||
|
||||
category = parts[0] if parts else "Unknown"
|
||||
|
||||
if category == "Terrain":
|
||||
terrain_type = parts[1] if len(parts) > 1 else "Unknown"
|
||||
scene = parts[2] if len(parts) > 2 else "Unknown"
|
||||
elif category == "Country":
|
||||
terrain_type = "Urban"
|
||||
scene = "/".join(parts[1:3]) if len(parts) > 2 else "Unknown"
|
||||
else:
|
||||
terrain_type = "Unknown"
|
||||
scene = parts[1] if len(parts) > 1 else "Unknown"
|
||||
|
||||
# Parse height and rotation from trajectory folder name.
|
||||
height_m = "100"
|
||||
heading_deg = "0"
|
||||
for part in parts:
|
||||
m = re.match(r"height(\d+)_rot(\d+)", part)
|
||||
if m:
|
||||
height_m = m.group(1)
|
||||
heading_deg = m.group(2)
|
||||
break
|
||||
|
||||
heading_idx = round(int(heading_deg) / 45) % 8
|
||||
heading_dir = _COMPASS[heading_idx]
|
||||
|
||||
return {
|
||||
"category": category,
|
||||
"terrain_type": terrain_type,
|
||||
"scene": scene,
|
||||
"height_m": height_m,
|
||||
"heading_deg": heading_deg,
|
||||
"heading_dir": heading_dir,
|
||||
}
|
||||
|
||||
|
||||
def _make_template_caption(meta: dict[str, str]) -> str:
|
||||
"""Generate a template caption from parsed metadata."""
|
||||
terrain = meta["terrain_type"]
|
||||
features = _TERRAIN_FEATURES.get(terrain, _COUNTRY_FEATURES)
|
||||
|
||||
return (
|
||||
f"Aerial view at {meta['height_m']}m facing {meta['heading_dir']} "
|
||||
f"over {terrain.lower()} terrain near {meta['scene']}. "
|
||||
f"Plan-view features: {features}."
|
||||
)
|
||||
|
||||
|
||||
@gin.configurable
|
||||
class VisLocCaptionDataset(Dataset):
|
||||
"""UAV-VisLoc pairs with generated captions.
|
||||
class GeoLocCaptionDataset(Dataset):
|
||||
"""UAV-GeoLoc pairs with template captions.
|
||||
|
||||
Reads train_query.txt, randomly samples one positive crop per query.
|
||||
|
||||
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.
|
||||
query_file: Path to train_query.txt (or test_query.txt).
|
||||
data_root: Root directory of UAV-GeoLoc dataset.
|
||||
image_transform: Callable applied to PIL images.
|
||||
drop_caption_prob: Probability of dropping caption (for ablation).
|
||||
seed: Random seed.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: str,
|
||||
image_root: str,
|
||||
query_file: str,
|
||||
data_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.data_root = Path(data_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.entries: list[dict[str, Any]] = []
|
||||
self._load_query_file(Path(query_file))
|
||||
|
||||
self._validate_entries()
|
||||
def _load_query_file(self, query_file: Path) -> None:
|
||||
"""Parse train_query.txt into list of entries."""
|
||||
with open(query_file) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split()
|
||||
query_path = parts[0]
|
||||
# parts[1] is label (always 0), parts[2:] are positive crop paths.
|
||||
positive_crops = parts[2:]
|
||||
if not positive_crops:
|
||||
continue
|
||||
|
||||
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}
|
||||
meta = _parse_query_path(query_path)
|
||||
caption = _make_template_caption(meta)
|
||||
|
||||
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}"
|
||||
self.entries.append({
|
||||
"query_path": query_path,
|
||||
"positive_crops": positive_crops,
|
||||
"caption": caption,
|
||||
"meta": meta,
|
||||
})
|
||||
|
||||
def _load_image(self, relative_path: str) -> torch.Tensor:
|
||||
"""Load image and apply preprocessing."""
|
||||
path = self.image_root / relative_path
|
||||
path = self.data_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"])
|
||||
drone_img = self._load_image(entry["query_path"])
|
||||
|
||||
caption_drone = self._maybe_drop(entry[self._caption_key("drone")])
|
||||
caption_sat = self._maybe_drop(entry[self._caption_key("sat")])
|
||||
# Randomly sample one positive crop.
|
||||
crop_path = self._rng.choice(entry["positive_crops"])
|
||||
sat_img = self._load_image(crop_path)
|
||||
|
||||
caption = entry["caption"]
|
||||
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
|
||||
caption = ""
|
||||
|
||||
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}"),
|
||||
"caption_drone": caption,
|
||||
"pair_id": entry["query_path"],
|
||||
}
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
"""Collate into batched dict. Captions stay as string 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],
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user