Initial import: World-UAV prepro
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
768
dataloader_v2.py
Normal file
768
dataloader_v2.py
Normal file
@@ -0,0 +1,768 @@
|
||||
"""
|
||||
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.")
|
||||
Reference in New Issue
Block a user