b6dccbba7b
PROBLEM: GTA-UAV has overlapping satellite crops (partial IoU). Standard InfoNCE with diagonal targets treated valid matches as negatives. R@K checked only diagonal — missed valid matches, artificially low recall. FIXES: 1. WeightedInfoNCE loss (src/losses/weighted_infonce.py): - Per-sample adaptive label smoothing from positive_weights (IoU) - Higher weight → sharper target, lower → softer (semi-positive tolerance) - Based on Game4Loc reference implementation 2. Multi-match R@K evaluation: - Uses dataset.get_all_valid_sat_names() to get ALL valid matches per query - R@K counts hit if ANY valid satellite is in top-K (not just diagonal) - AP computed as MRR over first valid match 3. Dataset returns positive_weight per sample: - Sampled satellite weight passed to loss for adaptive smoothing - All valid satellite candidates exposed for evaluation Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
290 lines
11 KiB
Python
290 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
"""GTA-UAV-LR dataset loader with L1/L2/L3 hierarchical captions for CVGL.
|
|
|
|
Reads cross-area pair JSONs and VLM-generated caption JSONs.
|
|
Produces (drone_img, sat_img, caption_l1, caption_l2, caption_l3) tuples.
|
|
|
|
Caption levels:
|
|
L1 (overview): First sentence of P1 (land cover summary).
|
|
L2 (full): Complete P1 + P2 text.
|
|
L3 (fingerprint): P3 unique signature section.
|
|
|
|
For short captions (pure water, no P markers): all levels get the same short text.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import random
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
import coloredlogs
|
|
import torch
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
from tqdm import tqdm
|
|
|
|
LOGGER = logging.getLogger("caption_test.gtauav_dataset")
|
|
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
|
|
|
|
# Default paths.
|
|
_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
|
|
_CAPTION_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-captions")
|
|
_EMPTY_CAPTION = ""
|
|
|
|
# Regex to split P1/P2/P3 sections.
|
|
_P_SPLIT = re.compile(r"\*\*P[123][^*]*\*\*\s*:?\s*")
|
|
|
|
|
|
def _parse_caption_levels(output: str) -> tuple[str, str, str]:
|
|
"""Split VLM caption output into L1, L2, L3 levels.
|
|
|
|
Returns:
|
|
(l1_overview, l2_full, l3_fingerprint)
|
|
"""
|
|
sections = _P_SPLIT.split(output)
|
|
# sections[0] is empty (before **P1**), sections[1]=P1, [2]=P2, [3]=P3
|
|
sections = [s.strip() for s in sections if s.strip()]
|
|
|
|
if len(sections) >= 3:
|
|
p1, p2, p3 = sections[0], sections[1], sections[2]
|
|
elif len(sections) == 2:
|
|
p1, p2, p3 = sections[0], sections[1], sections[0]
|
|
elif len(sections) == 1:
|
|
p1 = p2 = p3 = sections[0]
|
|
else:
|
|
p1 = p2 = p3 = output.strip()
|
|
|
|
# L1: first sentence of P1.
|
|
first_dot = p1.find(". ")
|
|
l1 = p1[:first_dot + 1] if first_dot > 0 else p1
|
|
|
|
# L2: full P1 + P2.
|
|
l2 = p1 + " " + p2
|
|
|
|
# L3: P3 (fingerprint / unique signature).
|
|
l3 = p3
|
|
|
|
return l1, l2, l3
|
|
|
|
|
|
def _load_caption_index(caption_root: Path) -> dict[str, dict]:
|
|
"""Build index: image_name -> caption JSON data.
|
|
|
|
Scans drone/images/ and satellite/ directories.
|
|
"""
|
|
index: dict[str, dict] = {}
|
|
|
|
for subdir in ["drone/images", "satellite"]:
|
|
cap_dir = caption_root / subdir
|
|
if not cap_dir.exists():
|
|
continue
|
|
cap_files = sorted(cap_dir.glob("*_caption.json"))
|
|
for cap_file in tqdm(cap_files, desc=f" 📄 Loading {subdir} captions", unit="file", leave=False):
|
|
with open(cap_file) as f:
|
|
data = json.load(f)
|
|
# Key by the image name (without _caption suffix).
|
|
img_name = cap_file.name.replace("_caption.json", ".png")
|
|
index[img_name] = data
|
|
|
|
return index
|
|
|
|
|
|
class GTAUAVDataset(Dataset):
|
|
"""GTA-UAV-LR dataset with hierarchical L1/L2/L3 captions.
|
|
|
|
Args:
|
|
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.
|
|
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.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
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,
|
|
seed: int = 0,
|
|
) -> None:
|
|
self.rgb_root = Path(rgb_root)
|
|
self.caption_root = Path(caption_root)
|
|
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)
|
|
|
|
# Load exclusion set from segmentation filter.
|
|
self.excluded: set[str] = set()
|
|
if filter_meta is not None:
|
|
self._load_filter(Path(filter_meta))
|
|
|
|
# Load caption index.
|
|
LOGGER.info("📚 Loading caption index from %s", caption_root)
|
|
self.caption_index = _load_caption_index(self.caption_root)
|
|
LOGGER.info("📚 Caption index: %d entries", len(self.caption_index))
|
|
|
|
# Load pairs.
|
|
self.entries: list[dict[str, Any]] = []
|
|
self._load_pairs(Path(pair_json))
|
|
LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json)
|
|
|
|
def _load_filter(self, path: Path) -> None:
|
|
with open(path) as f:
|
|
meta = json.load(f)
|
|
# excluded list contains paths like "drone/images/xxx.png"
|
|
for exc in meta.get("excluded", []):
|
|
# Extract image name from segm path.
|
|
self.excluded.add(Path(exc).name)
|
|
LOGGER.info("🔻 Filter loaded: %d excluded images", len(self.excluded))
|
|
|
|
def _load_pairs(self, pair_json: Path) -> None:
|
|
with open(pair_json) as f:
|
|
raw_pairs = json.load(f)
|
|
|
|
for pair in raw_pairs:
|
|
drone_name = pair["drone_img_name"]
|
|
|
|
# Skip excluded images.
|
|
if drone_name in self.excluded:
|
|
continue
|
|
|
|
# Get positive/semi-positive satellite images.
|
|
pos_list = pair.get("pair_pos_sate_img_list", [])
|
|
semipos_list = pair.get("pair_pos_semipos_sate_img_list", [])
|
|
semipos_weights = pair.get("pair_pos_semipos_sate_weight_list", [])
|
|
|
|
# Use positives if available, else semi-positives.
|
|
if pos_list:
|
|
sat_candidates = pos_list
|
|
sat_weights = None
|
|
elif semipos_list:
|
|
sat_candidates = semipos_list
|
|
sat_weights = semipos_weights if semipos_weights else None
|
|
else:
|
|
continue # No match, skip.
|
|
|
|
# Get drone captions.
|
|
cap_data = self.caption_index.get(drone_name)
|
|
if cap_data is not None:
|
|
l1, l2, l3 = _parse_caption_levels(cap_data["output"])
|
|
else:
|
|
l1 = l2 = l3 = _EMPTY_CAPTION
|
|
|
|
# Pre-parse satellite captions for all candidates.
|
|
sat_captions: dict[str, tuple[str, str, str]] = {}
|
|
for sat_name in sat_candidates:
|
|
sat_cap = self.caption_index.get(sat_name)
|
|
if sat_cap is not None:
|
|
sat_captions[sat_name] = _parse_caption_levels(sat_cap["output"])
|
|
|
|
self.entries.append({
|
|
"drone_name": drone_name,
|
|
"drone_dir": pair["drone_img_dir"],
|
|
"sat_dir": pair["sate_img_dir"],
|
|
"sat_candidates": sat_candidates,
|
|
"sat_weights": sat_weights,
|
|
"caption_l1": l1,
|
|
"caption_l2": l2,
|
|
"caption_l3": l3,
|
|
"sat_captions": sat_captions,
|
|
})
|
|
|
|
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 transform is not None:
|
|
return transform(rgb)
|
|
return torch.tensor(0) # placeholder if no transform
|
|
|
|
def get_all_valid_sat_names(self) -> list[list[str]]:
|
|
"""Return all valid satellite matches per drone query (for evaluation).
|
|
|
|
In GTA-UAV, each drone has multiple valid satellite crops (partial IoU).
|
|
Standard diagonal R@K is wrong — must check if ANY valid match is in top-K.
|
|
"""
|
|
return [entry["sat_candidates"] for entry in self.entries]
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.entries)
|
|
|
|
def __getitem__(self, idx: int) -> dict[str, Any]:
|
|
entry = self.entries[idx]
|
|
|
|
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:
|
|
sat_idx = self._rng.choices(
|
|
range(len(entry["sat_candidates"])),
|
|
weights=entry["sat_weights"],
|
|
k=1,
|
|
)[0]
|
|
sat_name = entry["sat_candidates"][sat_idx]
|
|
pos_weight = entry["sat_weights"][sat_idx]
|
|
else:
|
|
sat_idx = self._rng.randrange(len(entry["sat_candidates"]))
|
|
sat_name = entry["sat_candidates"][sat_idx]
|
|
pos_weight = 1.0 # strict positive
|
|
|
|
sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform)
|
|
|
|
# Drone captions with optional dropout.
|
|
if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
|
|
l1 = l2 = l3 = _EMPTY_CAPTION
|
|
else:
|
|
l1, l2, l3 = entry["caption_l1"], entry["caption_l2"], entry["caption_l3"]
|
|
|
|
# Satellite captions (empty string if not available).
|
|
sat_caps = entry["sat_captions"].get(sat_name)
|
|
if sat_caps is not None:
|
|
sat_l1, sat_l2, sat_l3 = sat_caps
|
|
else:
|
|
sat_l1 = sat_l2 = sat_l3 = _EMPTY_CAPTION
|
|
|
|
return {
|
|
"drone_img": drone_img,
|
|
"sat_img": sat_img,
|
|
"caption_l1": l1,
|
|
"caption_l2": l2,
|
|
"caption_l3": l3,
|
|
"sat_caption_l1": sat_l1,
|
|
"sat_caption_l2": sat_l2,
|
|
"sat_caption_l3": sat_l3,
|
|
"pair_id": entry["drone_name"],
|
|
"sat_name": sat_name,
|
|
"positive_weight": pos_weight,
|
|
}
|
|
|
|
|
|
def collate_gtauav_batch(
|
|
batch: list[dict[str, Any]],
|
|
) -> dict[str, Any]:
|
|
"""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_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],
|
|
"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],
|
|
"pair_ids": [b["pair_id"] for b in batch],
|
|
"sat_names": [b["sat_name"] for b in batch],
|
|
"positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32),
|
|
}
|