Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
769 lines
28 KiB
Python
769 lines
28 KiB
Python
"""
|
|
PyTorch DataLoader for UAV-GeoLoc dataset (Cross-View Geo-Localization).
|
|
|
|
Dataset structure (discovered empirically):
|
|
- 372 scenes: 171 Country + 200 Terrain + 1 Rot
|
|
- 927K images: 652K query (drone, 512x512 JPEG) + 275K DB (satellite, PNG)
|
|
- Satellite crops: stride = crop_size / 2 (50% overlap), 11 unique sizes (100-1000 px)
|
|
- Query variants: 3 heights (100/125/150m) x 8 azimuths (0-315, step 45) = 24 per scene
|
|
- Camera: 30 deg vertical FOV, top-down, 76 frames per trajectory
|
|
|
|
Supports:
|
|
- Triplet training with random / semi-positive-aware negative mining
|
|
- Pair training for contrastive learning
|
|
- Evaluation with separate query / DB sets
|
|
- Camera metadata (height, azimuth, GPS) per query
|
|
- GPS-based localization error computation
|
|
- Satellite tiling utility for new data
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import re
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset, DataLoader
|
|
import torchvision.transforms as T
|
|
|
|
|
|
# ── Index file parsing ──────────────────────────────────────────────────────
|
|
|
|
def _parse_query_line(line: str):
|
|
"""Parse a query index line. Handles paths with spaces.
|
|
|
|
Format: <query_path> <label_int> <pos_db_1> [pos_db_2 ...]
|
|
DB paths always contain /DB/img/crop_.
|
|
"""
|
|
line = line.strip()
|
|
if not line:
|
|
return None
|
|
db_pattern = re.compile(r'\S*DB/img/crop_\S+')
|
|
db_matches = list(db_pattern.finditer(line))
|
|
if not db_matches:
|
|
return None
|
|
before_db = line[:db_matches[0].start()].rstrip()
|
|
label_match = re.search(r'\s(\d+)\s*$', before_db)
|
|
if not label_match:
|
|
return None
|
|
label = int(label_match.group(1))
|
|
query_path = before_db[:label_match.start()].strip()
|
|
positives = [m.group() for m in db_matches]
|
|
return query_path, label, positives
|
|
|
|
|
|
def _parse_height_rot(query_path: str):
|
|
"""Extract height and rotation from query path.
|
|
|
|
e.g. '.../height125_rot270/footage/...' -> (125, 270)
|
|
Handles typos like 'eight150', '125_rot315', 'ght100'.
|
|
"""
|
|
m = re.search(r'(?:h(?:eigh?t?)?)?(\d{3})_rot(\d+)', query_path)
|
|
if m:
|
|
return int(m.group(1)), int(m.group(2))
|
|
return None, None
|
|
|
|
|
|
# ── GPS utilities ───────────────────────────────────────────────────────────
|
|
|
|
def load_db_positions(db_pos_path: str) -> dict:
|
|
"""Load db_postion.txt -> {crop_filename: (lon, lat, res_x, res_y)}."""
|
|
positions = {}
|
|
if not os.path.isfile(db_pos_path):
|
|
return positions
|
|
with open(db_pos_path) as f:
|
|
for line in f:
|
|
parts = line.strip().split()
|
|
if len(parts) >= 3:
|
|
name = parts[0]
|
|
lon, lat = float(parts[1]), float(parts[2])
|
|
res_x = float(parts[3]) if len(parts) > 3 else 0.0
|
|
res_y = float(parts[4]) if len(parts) > 4 else 0.0
|
|
positions[name] = (lon, lat, res_x, res_y)
|
|
return positions
|
|
|
|
|
|
def haversine_m(lon1, lat1, lon2, lat2):
|
|
"""Haversine distance in meters between two GPS points."""
|
|
R = 6_371_000
|
|
phi1, phi2 = math.radians(lat1), math.radians(lat2)
|
|
dphi = math.radians(lat2 - lat1)
|
|
dlam = math.radians(lon2 - lon1)
|
|
a = math.sin(dphi / 2) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlam / 2) ** 2
|
|
return 2 * R * math.atan2(math.sqrt(a), math.sqrt(1 - a))
|
|
|
|
|
|
# ── Satellite tiling utility ────────────────────────────────────────────────
|
|
|
|
def tile_satellite_image(
|
|
image: np.ndarray,
|
|
crop_size: int = 200,
|
|
stride: Optional[int] = None,
|
|
) -> list:
|
|
"""Tile a satellite image following the UAV-GeoLoc convention.
|
|
|
|
Args:
|
|
image: HxWx3 numpy array (the merge.tif).
|
|
crop_size: Size of each square crop in pixels.
|
|
stride: Step between crops. Default = crop_size // 2 (50% overlap).
|
|
|
|
Returns:
|
|
List of (crop_array, x_idx, y_idx, pixel_x, pixel_y) tuples.
|
|
crop_X_Y.png naming: X = col index, Y = row index.
|
|
"""
|
|
if stride is None:
|
|
stride = crop_size // 2
|
|
h, w = image.shape[:2]
|
|
crops = []
|
|
for x_idx, px in enumerate(range(0, w - crop_size + 1, stride)):
|
|
for y_idx, py in enumerate(range(0, h - crop_size + 1, stride)):
|
|
crop = image[py:py + crop_size, px:px + crop_size]
|
|
crops.append((crop, x_idx, y_idx, px, py))
|
|
return crops
|
|
|
|
|
|
# ── Default transforms ──────────────────────────────────────────────────────
|
|
|
|
def _default_train_query_transform(img_size=224):
|
|
return T.Compose([
|
|
T.Resize((img_size, img_size)),
|
|
T.RandomHorizontalFlip(),
|
|
T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
|
|
def _default_train_db_transform(img_size=224):
|
|
return T.Compose([
|
|
T.Resize((img_size, img_size)),
|
|
T.RandomHorizontalFlip(),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
|
|
def _default_eval_transform(img_size=224):
|
|
return T.Compose([
|
|
T.Resize((img_size, img_size)),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
|
|
# ── Training dataset: triplets ──────────────────────────────────────────────
|
|
|
|
class UAVGeoLocTrain(Dataset):
|
|
"""Training dataset returning (query, positive, negative) triplets.
|
|
|
|
Supports semi-positive-aware negative mining: negatives are guaranteed
|
|
to NOT be in the positive or semi-positive set for the query's scene.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
root: str,
|
|
query_file: str = "Index/train_query.txt",
|
|
db_file: str = "Index/train_db.txt",
|
|
img_size: int = 224,
|
|
transform_query=None,
|
|
transform_db=None,
|
|
):
|
|
self.root = root
|
|
|
|
self.queries = []
|
|
with open(os.path.join(root, query_file)) as f:
|
|
for line in f:
|
|
parsed = _parse_query_line(line)
|
|
if parsed is None:
|
|
continue
|
|
query_path, label, positives = parsed
|
|
height, rot = _parse_height_rot(query_path)
|
|
self.queries.append({
|
|
"path": query_path,
|
|
"label": label,
|
|
"positives": positives,
|
|
"height": height,
|
|
"rotation": rot,
|
|
})
|
|
|
|
self.db_paths = []
|
|
with open(os.path.join(root, db_file)) as f:
|
|
for line in f:
|
|
p = line.strip()
|
|
if p:
|
|
self.db_paths.append(p)
|
|
|
|
# Build label -> all positive+semi-positive DB paths for clean negative mining
|
|
self._label_positives = {}
|
|
for q in self.queries:
|
|
lbl = q["label"]
|
|
if lbl not in self._label_positives:
|
|
self._label_positives[lbl] = set()
|
|
self._label_positives[lbl].update(q["positives"])
|
|
|
|
self.transform_query = transform_query or _default_train_query_transform(img_size)
|
|
self.transform_db = transform_db or _default_train_db_transform(img_size)
|
|
|
|
def __len__(self):
|
|
return len(self.queries)
|
|
|
|
def _load(self, rel_path):
|
|
return Image.open(os.path.join(self.root, rel_path)).convert("RGB")
|
|
|
|
def __getitem__(self, idx):
|
|
q = self.queries[idx]
|
|
|
|
query_img = self.transform_query(self._load(q["path"]))
|
|
pos_img = self.transform_db(self._load(random.choice(q["positives"])))
|
|
|
|
# Negative: random DB image not in this scene's positive set
|
|
pos_set = self._label_positives.get(q["label"], set())
|
|
while True:
|
|
neg_path = random.choice(self.db_paths)
|
|
if neg_path not in pos_set:
|
|
break
|
|
neg_img = self.transform_db(self._load(neg_path))
|
|
|
|
return {
|
|
"query": query_img,
|
|
"positive": pos_img,
|
|
"negative": neg_img,
|
|
"label": q["label"],
|
|
"height": q["height"] or 0,
|
|
"rotation": q["rotation"] or 0,
|
|
}
|
|
|
|
|
|
# ── Training dataset: pairs ─────────────────────────────────────────────────
|
|
|
|
class UAVGeoLocPair(Dataset):
|
|
"""Training dataset returning (query, positive) pairs for contrastive learning."""
|
|
|
|
def __init__(
|
|
self,
|
|
root: str,
|
|
query_file: str = "Index/train_query.txt",
|
|
img_size: int = 224,
|
|
transform_query=None,
|
|
transform_db=None,
|
|
):
|
|
self.root = root
|
|
|
|
self.queries = []
|
|
with open(os.path.join(root, query_file)) as f:
|
|
for line in f:
|
|
parsed = _parse_query_line(line)
|
|
if parsed is None:
|
|
continue
|
|
query_path, label, positives = parsed
|
|
height, rot = _parse_height_rot(query_path)
|
|
self.queries.append({
|
|
"path": query_path,
|
|
"label": label,
|
|
"positives": positives,
|
|
"height": height,
|
|
"rotation": rot,
|
|
})
|
|
|
|
self.transform_query = transform_query or _default_train_query_transform(img_size)
|
|
self.transform_db = transform_db or _default_train_db_transform(img_size)
|
|
|
|
def __len__(self):
|
|
return len(self.queries)
|
|
|
|
def __getitem__(self, idx):
|
|
q = self.queries[idx]
|
|
query_img = self.transform_query(
|
|
Image.open(os.path.join(self.root, q["path"])).convert("RGB")
|
|
)
|
|
pos_img = self.transform_db(
|
|
Image.open(os.path.join(self.root, random.choice(q["positives"]))).convert("RGB")
|
|
)
|
|
return {
|
|
"query": query_img,
|
|
"positive": pos_img,
|
|
"label": q["label"],
|
|
"height": q["height"] or 0,
|
|
"rotation": q["rotation"] or 0,
|
|
}
|
|
|
|
|
|
# ── Evaluation dataset ──────────────────────────────────────────────────────
|
|
|
|
class UAVGeoLocEval(Dataset):
|
|
"""Evaluation dataset for retrieval. Returns single images with metadata.
|
|
|
|
mode="query": loads UAV query images with labels, positives, height, rotation.
|
|
mode="db": loads satellite DB images.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
root: str,
|
|
index_file: str,
|
|
mode: str = "query",
|
|
img_size: int = 224,
|
|
transform=None,
|
|
):
|
|
self.root = root
|
|
self.mode = mode
|
|
|
|
self.images = []
|
|
self.labels = []
|
|
self.positives = []
|
|
self.heights = []
|
|
self.rotations = []
|
|
|
|
with open(os.path.join(root, index_file)) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
if mode == "db":
|
|
self.images.append(line)
|
|
else:
|
|
parsed = _parse_query_line(line)
|
|
if parsed is None:
|
|
continue
|
|
query_path, label, positives = parsed
|
|
self.images.append(query_path)
|
|
self.labels.append(label)
|
|
self.positives.append(positives)
|
|
h, r = _parse_height_rot(query_path)
|
|
self.heights.append(h or 0)
|
|
self.rotations.append(r or 0)
|
|
|
|
self.transform = transform or _default_eval_transform(img_size)
|
|
|
|
if mode == "db":
|
|
self.path_to_idx = {p: i for i, p in enumerate(self.images)}
|
|
|
|
def __len__(self):
|
|
return len(self.images)
|
|
|
|
def __getitem__(self, idx):
|
|
path = self.images[idx]
|
|
img = Image.open(os.path.join(self.root, path)).convert("RGB")
|
|
img = self.transform(img)
|
|
|
|
out = {"image": img, "path": path, "index": idx}
|
|
if self.mode == "query":
|
|
out["label"] = self.labels[idx]
|
|
out["positives"] = self.positives[idx]
|
|
out["height"] = self.heights[idx]
|
|
out["rotation"] = self.rotations[idx]
|
|
return out
|
|
|
|
|
|
# ── Scene-based dataset (direct from directory, no index files) ─────────────
|
|
|
|
class UAVGeoLocScene(Dataset):
|
|
"""Load data directly from a scene directory. Useful for custom splits
|
|
or scenes not covered by Index files (e.g., Rot subset).
|
|
|
|
Returns (query_img, db_positive_img, metadata_dict) pairs.
|
|
|
|
Args:
|
|
scene_dir: Path to scene (e.g., .../Country/Australia/Adelaide/AdelaideCBD)
|
|
heights: List of heights to include. Default: [100, 125, 150].
|
|
rotations: List of rotations to include. Default: all 8.
|
|
frames: List of frame indices or None for all.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
scene_dir: str,
|
|
heights: Optional[list] = None,
|
|
rotations: Optional[list] = None,
|
|
frames: Optional[list] = None,
|
|
img_size: int = 224,
|
|
transform_query=None,
|
|
transform_db=None,
|
|
):
|
|
self.scene_dir = scene_dir
|
|
heights = heights or [100, 125, 150]
|
|
rotations = rotations or [0, 45, 90, 135, 180, 225, 270, 315]
|
|
|
|
# Check for incomplete scene (missing DB crops or annotations)
|
|
pos_path = os.path.join(scene_dir, "positive.json")
|
|
db_img_dir = os.path.join(scene_dir, "DB", "img")
|
|
missing = []
|
|
if not os.path.isfile(pos_path):
|
|
missing.append("positive.json")
|
|
if not os.path.isdir(db_img_dir):
|
|
missing.append("DB/img/")
|
|
if missing:
|
|
scene_name = os.path.basename(scene_dir)
|
|
raise FileNotFoundError(
|
|
f"Incomplete scene '{scene_name}': missing {', '.join(missing)}. "
|
|
f"17 known incomplete scenes (Edinburgh, London, Manchester, "
|
|
f"Birmingham/JewelleryQuarter, Chicago/__MACOSX) lack DB crops "
|
|
f"and annotations — they cannot be used for training or evaluation."
|
|
)
|
|
|
|
# Load positive.json
|
|
with open(pos_path) as f:
|
|
self.positive_map = json.load(f)
|
|
|
|
# Load semi_positive.json if available
|
|
semi_path = os.path.join(scene_dir, "semi_positive.json")
|
|
self.semi_positive_map = {}
|
|
if os.path.isfile(semi_path):
|
|
with open(semi_path) as f:
|
|
self.semi_positive_map = json.load(f)
|
|
|
|
# Load DB GPS positions
|
|
db_dir = os.path.join(scene_dir, "DB")
|
|
self.db_positions = load_db_positions(os.path.join(db_dir, "db_postion.txt"))
|
|
|
|
# Enumerate valid (variant, frame) pairs
|
|
self.samples = []
|
|
query_dir = os.path.join(scene_dir, "query")
|
|
for h in heights:
|
|
for r in rotations:
|
|
variant_name = f"height{h}_rot{r}"
|
|
footage_dir = os.path.join(query_dir, variant_name, "footage")
|
|
if not os.path.isdir(footage_dir):
|
|
continue
|
|
|
|
available_frames = sorted([
|
|
f for f in os.listdir(footage_dir)
|
|
if f.endswith((".jpeg", ".jpg"))
|
|
])
|
|
|
|
for frame_file in available_frames:
|
|
# Extract frame index (e.g. "height100_rot0_38.jpeg" -> "38")
|
|
m = re.search(r'_(\d+)\.jpe?g$', frame_file)
|
|
if m is None:
|
|
continue
|
|
frame_idx = m.group(1)
|
|
|
|
if frames is not None and int(frame_idx) not in frames:
|
|
continue
|
|
|
|
# Get positive DB crop(s)
|
|
pos_crops = self.positive_map.get(frame_idx, [])
|
|
if not pos_crops:
|
|
continue
|
|
|
|
self.samples.append({
|
|
"query_path": os.path.join(footage_dir, frame_file),
|
|
"frame_idx": frame_idx,
|
|
"height": h,
|
|
"rotation": r,
|
|
"positives": pos_crops,
|
|
"semi_positives": self.semi_positive_map.get(frame_idx, []),
|
|
})
|
|
|
|
self.db_img_dir = os.path.join(db_dir, "img")
|
|
|
|
self.transform_query = transform_query or _default_eval_transform(img_size)
|
|
self.transform_db = transform_db or _default_eval_transform(img_size)
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx):
|
|
s = self.samples[idx]
|
|
|
|
query_img = self.transform_query(
|
|
Image.open(s["query_path"]).convert("RGB")
|
|
)
|
|
pos_crop_name = random.choice(s["positives"])
|
|
pos_img = self.transform_db(
|
|
Image.open(os.path.join(self.db_img_dir, pos_crop_name)).convert("RGB")
|
|
)
|
|
|
|
# GPS of positive crop centroid
|
|
gps = self.db_positions.get(pos_crop_name, (0.0, 0.0, 0.0, 0.0))
|
|
|
|
return {
|
|
"query": query_img,
|
|
"positive": pos_img,
|
|
"height": s["height"],
|
|
"rotation": s["rotation"],
|
|
"frame_idx": int(s["frame_idx"]),
|
|
"positive_name": pos_crop_name,
|
|
"positive_lon": gps[0],
|
|
"positive_lat": gps[1],
|
|
"semi_positives": s["semi_positives"],
|
|
}
|
|
|
|
|
|
# ── Collate functions ───────────────────────────────────────────────────────
|
|
|
|
def eval_collate_fn(batch):
|
|
"""Collate for eval datasets (handles variable-length positives)."""
|
|
images = torch.stack([item["image"] for item in batch])
|
|
indices = torch.tensor([item["index"] for item in batch])
|
|
paths = [item["path"] for item in batch]
|
|
out = {"image": images, "index": indices, "path": paths}
|
|
if "label" in batch[0]:
|
|
out["label"] = torch.tensor([item["label"] for item in batch])
|
|
out["positives"] = [item["positives"] for item in batch]
|
|
out["height"] = torch.tensor([item["height"] for item in batch])
|
|
out["rotation"] = torch.tensor([item["rotation"] for item in batch])
|
|
return out
|
|
|
|
|
|
def scene_collate_fn(batch):
|
|
"""Collate for UAVGeoLocScene (handles variable-length semi_positives)."""
|
|
return {
|
|
"query": torch.stack([b["query"] for b in batch]),
|
|
"positive": torch.stack([b["positive"] for b in batch]),
|
|
"height": torch.tensor([b["height"] for b in batch]),
|
|
"rotation": torch.tensor([b["rotation"] for b in batch]),
|
|
"frame_idx": torch.tensor([b["frame_idx"] for b in batch]),
|
|
"positive_name": [b["positive_name"] for b in batch],
|
|
"positive_lon": torch.tensor([b["positive_lon"] for b in batch], dtype=torch.float64),
|
|
"positive_lat": torch.tensor([b["positive_lat"] for b in batch], dtype=torch.float64),
|
|
"semi_positives": [b["semi_positives"] for b in batch],
|
|
}
|
|
|
|
|
|
# ── Localization error evaluation ───────────────────────────────────────────
|
|
|
|
def compute_localization_error(
|
|
query_dataset: UAVGeoLocEval,
|
|
db_dataset: UAVGeoLocEval,
|
|
predictions: np.ndarray,
|
|
db_positions_cache: Optional[dict] = None,
|
|
) -> dict:
|
|
"""Compute localization error in meters given retrieval predictions.
|
|
|
|
Args:
|
|
query_dataset: UAVGeoLocEval in "query" mode.
|
|
db_dataset: UAVGeoLocEval in "db" mode.
|
|
predictions: (N_query,) array of predicted DB indices.
|
|
db_positions_cache: Pre-loaded {db_path: (lon, lat, ...)} dict.
|
|
If None, loads from db_postion.txt files on the fly.
|
|
|
|
Returns:
|
|
dict with 'mean_error_m', 'median_error_m', 'errors' (per-query list).
|
|
"""
|
|
if db_positions_cache is None:
|
|
db_positions_cache = {}
|
|
# Discover all unique scene DB dirs from db paths
|
|
scene_dirs = set()
|
|
for p in db_dataset.images:
|
|
# e.g. "Terrain/Mountain/Andes/DB/img/crop_0_0.png" -> "Terrain/Mountain/Andes/DB"
|
|
db_dir = str(Path(p).parent.parent)
|
|
scene_dirs.add(db_dir)
|
|
for sd in scene_dirs:
|
|
pos_file = os.path.join(db_dataset.root, sd, "db_postion.txt")
|
|
positions = load_db_positions(pos_file)
|
|
for crop_name, coords in positions.items():
|
|
full_path = os.path.join(sd, "img", crop_name)
|
|
db_positions_cache[full_path] = coords
|
|
|
|
errors = []
|
|
for i, pred_idx in enumerate(predictions):
|
|
# Get GT positive crops for this query
|
|
gt_positives = query_dataset.positives[i]
|
|
# Get predicted DB path
|
|
pred_path = db_dataset.images[int(pred_idx)]
|
|
|
|
# GPS of prediction
|
|
pred_gps = db_positions_cache.get(pred_path)
|
|
if pred_gps is None:
|
|
continue
|
|
|
|
# GPS of nearest GT positive
|
|
min_dist = float("inf")
|
|
for gt_path in gt_positives:
|
|
gt_gps = db_positions_cache.get(gt_path)
|
|
if gt_gps is None:
|
|
continue
|
|
dist = haversine_m(pred_gps[0], pred_gps[1], gt_gps[0], gt_gps[1])
|
|
min_dist = min(min_dist, dist)
|
|
|
|
if min_dist < float("inf"):
|
|
errors.append(min_dist)
|
|
|
|
errors_arr = np.array(errors)
|
|
return {
|
|
"mean_error_m": float(np.mean(errors_arr)) if len(errors_arr) else 0.0,
|
|
"median_error_m": float(np.median(errors_arr)) if len(errors_arr) else 0.0,
|
|
"errors": errors,
|
|
"num_evaluated": len(errors),
|
|
}
|
|
|
|
|
|
# ── Convenience builder ─────────────────────────────────────────────────────
|
|
|
|
def build_dataloaders(
|
|
root: str,
|
|
split: str = "terrain",
|
|
batch_size: int = 32,
|
|
img_size: int = 224,
|
|
num_workers: int = 4,
|
|
mode: str = "triplet",
|
|
):
|
|
"""Build train/val/test dataloaders.
|
|
|
|
Args:
|
|
root: Path to UAV-GeoLoc dataset root.
|
|
split: "terrain" (default), "country", or "all".
|
|
batch_size: Batch size for all loaders.
|
|
img_size: Resize images to (img_size, img_size).
|
|
num_workers: DataLoader workers.
|
|
mode: "triplet" for (query, pos, neg) or "pair" for (query, pos).
|
|
|
|
Returns:
|
|
dict with keys: "train", "val_query", "val_db", "test_query", "test_db"
|
|
"""
|
|
suffix = {"terrain": "", "country": "_country", "all": "_all"}[split]
|
|
|
|
if mode == "triplet":
|
|
train_ds = UAVGeoLocTrain(
|
|
root,
|
|
query_file=f"Index/train_query{suffix}.txt",
|
|
db_file=f"Index/train_db{suffix}.txt",
|
|
img_size=img_size,
|
|
)
|
|
else:
|
|
train_ds = UAVGeoLocPair(
|
|
root,
|
|
query_file=f"Index/train_query{suffix}.txt",
|
|
img_size=img_size,
|
|
)
|
|
|
|
train_loader = DataLoader(
|
|
train_ds,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=num_workers,
|
|
pin_memory=True,
|
|
drop_last=True,
|
|
)
|
|
|
|
loaders = {"train": train_loader}
|
|
|
|
for phase in ["val", "test"]:
|
|
q_file = f"Index/{phase}_query{suffix}.txt"
|
|
d_file = f"Index/{phase}_db{suffix}.txt"
|
|
# Fall back to unsuffixed if suffixed file doesn't exist
|
|
if not os.path.isfile(os.path.join(root, q_file)):
|
|
q_file = f"Index/{phase}_query.txt"
|
|
if not os.path.isfile(os.path.join(root, d_file)):
|
|
d_file = f"Index/{phase}_db.txt"
|
|
|
|
q_ds = UAVGeoLocEval(root, q_file, mode="query", img_size=img_size)
|
|
d_ds = UAVGeoLocEval(root, d_file, mode="db", img_size=img_size)
|
|
|
|
loaders[f"{phase}_query"] = DataLoader(
|
|
q_ds, batch_size=batch_size, shuffle=False,
|
|
num_workers=num_workers, pin_memory=True, collate_fn=eval_collate_fn,
|
|
)
|
|
loaders[f"{phase}_db"] = DataLoader(
|
|
d_ds, batch_size=batch_size, shuffle=False,
|
|
num_workers=num_workers, pin_memory=True, collate_fn=eval_collate_fn,
|
|
)
|
|
|
|
return loaders
|
|
|
|
|
|
def build_rot_loader(
|
|
root: str,
|
|
batch_size: int = 32,
|
|
img_size: int = 224,
|
|
num_workers: int = 4,
|
|
heights: Optional[list] = None,
|
|
rotations: Optional[list] = None,
|
|
):
|
|
"""Build a DataLoader for the Rot evaluation subset.
|
|
|
|
The Rot subset has 88 query variants (72 at h100 every 5deg + 16 at h125/h150)
|
|
over a single scene (SouthernSuburbs), useful for rotation robustness evaluation.
|
|
|
|
Returns:
|
|
DataLoader yielding scene_collate_fn batches.
|
|
"""
|
|
scene_dir = os.path.join(root, "Rot", "SouthernSuburbs")
|
|
# Rot has fine-grained rotations at height100
|
|
if rotations is None and heights is None:
|
|
# Include all: h100 has 72 rotations (0-355 step 5), h125/h150 have 8 each
|
|
heights_rot = [(100, r) for r in range(0, 360, 5)]
|
|
heights_rot += [(h, r) for h in [125, 150] for r in range(0, 360, 45)]
|
|
all_heights = list(set(h for h, _ in heights_rot))
|
|
all_rotations = list(set(r for _, r in heights_rot))
|
|
else:
|
|
all_heights = heights or [100, 125, 150]
|
|
all_rotations = rotations or list(range(0, 360, 5))
|
|
|
|
ds = UAVGeoLocScene(
|
|
scene_dir,
|
|
heights=all_heights,
|
|
rotations=all_rotations,
|
|
img_size=img_size,
|
|
)
|
|
return DataLoader(
|
|
ds, batch_size=batch_size, shuffle=False,
|
|
num_workers=num_workers, pin_memory=True, collate_fn=scene_collate_fn,
|
|
)
|
|
|
|
|
|
# ── Quick test ──────────────────────────────────────────────────────────────
|
|
|
|
if __name__ == "__main__":
|
|
ROOT = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"
|
|
|
|
print("=" * 60)
|
|
print("Building dataloaders (terrain split, img_size=224)...")
|
|
loaders = build_dataloaders(ROOT, split="terrain", batch_size=4, num_workers=0)
|
|
|
|
for name, loader in loaders.items():
|
|
print(f" {name}: {len(loader.dataset)} samples")
|
|
|
|
batch = next(iter(loaders["train"]))
|
|
print(f"\nTrain batch:")
|
|
print(f" query: {batch['query'].shape}")
|
|
print(f" positive: {batch['positive'].shape}")
|
|
print(f" negative: {batch['negative'].shape}")
|
|
print(f" labels: {batch['label']}")
|
|
print(f" heights: {batch['height']}")
|
|
print(f" rotations:{batch['rotation']}")
|
|
|
|
batch = next(iter(loaders["val_query"]))
|
|
print(f"\nVal query batch:")
|
|
print(f" image: {batch['image'].shape}")
|
|
print(f" heights: {batch['height']}")
|
|
print(f" rotations:{batch['rotation']}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Scene-based loader (AdelaideCBD, h100 only)...")
|
|
scene_ds = UAVGeoLocScene(
|
|
os.path.join(ROOT, "Country/Australia/Adelaide/AdelaideCBD"),
|
|
heights=[100],
|
|
rotations=[0, 90, 180, 270],
|
|
)
|
|
print(f" Samples: {len(scene_ds)}")
|
|
s = scene_ds[0]
|
|
print(f" query: {s['query'].shape}, height={s['height']}, rot={s['rotation']}")
|
|
print(f" positive: {s['positive_name']}, GPS=({s['positive_lon']:.4f}, {s['positive_lat']:.4f})")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Rot evaluation loader...")
|
|
rot_loader = build_rot_loader(ROOT, batch_size=4, num_workers=0, heights=[100], rotations=[0, 45, 90])
|
|
print(f" Samples: {len(rot_loader.dataset)}")
|
|
batch = next(iter(rot_loader))
|
|
print(f" query: {batch['query'].shape}")
|
|
print(f" rotations: {batch['rotation']}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Tiling utility test...")
|
|
dummy = np.random.randint(0, 255, (500, 600, 3), dtype=np.uint8)
|
|
crops = tile_satellite_image(dummy, crop_size=200, stride=100)
|
|
print(f" Input: 600x500, crop=200, stride=100")
|
|
print(f" Crops generated: {len(crops)}")
|
|
print(f" Grid: {max(c[1] for c in crops)+1} x {max(c[2] for c in crops)+1}")
|
|
|
|
print("\nAll tests passed.")
|