Add GTA-UAV experiment: asymmetric DINOv3 + LRSCLIP text encoder
V3 architecture for CVGL caption validation on GTA-UAV-LR dataset: - AsymmetricEncoder: DINOv3 ViT-L/16 (LVD drone + SAT satellite, frozen) + LRSCLIP/DGTRS-CLIP ViT-L-14 text encoder (248 tok, partial unfreeze) - L1/L2/L3 hierarchical captions from VLM-generated descriptions - TextFusionMLP (concat 3x768 -> MLP -> 512) + GatedFusion - Segmentation filter: exclude images with >=90% background+water - 10.9M trainable / 733M total params, 256x256 input - coloredlogs + tqdm + emoji for training UX - Baseline mode (--baseline): image-only, no text encoder loaded Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
247
src/datasets/gtauav_dataset.py
Normal file
247
src/datasets/gtauav_dataset.py
Normal file
@@ -0,0 +1,247 @@
|
||||
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.
|
||||
image_transform: Callable applied to PIL images.
|
||||
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),
|
||||
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.image_transform = 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 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
|
||||
|
||||
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,
|
||||
})
|
||||
|
||||
def _load_image(self, directory: str, filename: str) -> 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)
|
||||
return torch.tensor(0) # placeholder if no transform
|
||||
|
||||
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"])
|
||||
|
||||
# Sample satellite match (weighted if semi-positive).
|
||||
if entry["sat_weights"] is not None:
|
||||
sat_name = self._rng.choices(
|
||||
entry["sat_candidates"],
|
||||
weights=entry["sat_weights"],
|
||||
k=1,
|
||||
)[0]
|
||||
else:
|
||||
sat_name = self._rng.choice(entry["sat_candidates"])
|
||||
|
||||
sat_img = self._load_image(entry["sat_dir"], sat_name)
|
||||
|
||||
# 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"]
|
||||
|
||||
return {
|
||||
"drone_img": drone_img,
|
||||
"sat_img": sat_img,
|
||||
"caption_l1": l1,
|
||||
"caption_l2": l2,
|
||||
"caption_l3": l3,
|
||||
"pair_id": entry["drone_name"],
|
||||
}
|
||||
|
||||
|
||||
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],
|
||||
"pair_ids": [b["pair_id"] for b in batch],
|
||||
}
|
||||
Reference in New Issue
Block a user