fd014a3155
- Satellite crop generation (512x512, stride 256, resize 256x256) - Route 09 tile stitching (4 tiles -> 44800x33280) - GPS matching drone->crop via vectorized haversine - Index files in UAV-GeoLoc format (train/test query + DB) - positive.json / semi_positive.json / db_postion.txt per route - Route 07 excluded (satellite too narrow) - Fixed: full gallery in DB files, db_postion.txt format, frame_id keys - Fixed: file handle leaks in image processing loops Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
560 lines
18 KiB
Python
560 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare UAV-VisLoc dataset for retrieval training.
|
|
|
|
Pipeline:
|
|
1. Resize drone images -> 256x256
|
|
2. Stitch satellite tiles for route 09
|
|
3. Crop satellite maps -> 512x512 patches with stride 256, resize -> 256x256
|
|
4. Compute GPS for each crop center
|
|
5. Match drone -> crops via GPS (positive, semi-positive)
|
|
6. Generate metadata: positive.json, semi_positive.json, db_postion.txt
|
|
7. Generate Index files (train_query.txt, train_db.txt, etc.)
|
|
|
|
Usage:
|
|
python scripts/prepare_dataset.py \
|
|
--src /path/to/UAV_VisLoc_dataset \
|
|
--dst /path/to/UAV_VisLoc_processed \
|
|
--crop-size 512 \
|
|
--stride 256 \
|
|
--target-size 256
|
|
"""
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import math
|
|
import os
|
|
import warnings
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
warnings.filterwarnings("ignore", category=Image.DecompressionBombWarning)
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
|
|
# Route 07 excluded: satellite map 3000x170 — too narrow for 512x512 crops
|
|
EXCLUDED_ROUTES = {"07"}
|
|
|
|
# Route 09 has satellite split into 4 tiles
|
|
SPLIT_TILE_ROUTE = "09"
|
|
TILE_LAYOUT_09 = {
|
|
# (col, row): filename
|
|
(0, 0): "satellite09_01-01.tif",
|
|
(1, 0): "satellite09_01-02.tif",
|
|
(0, 1): "satellite09_02-01.tif",
|
|
(1, 1): "satellite09_02-02.tif",
|
|
}
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Prepare UAV-VisLoc for retrieval.")
|
|
parser.add_argument("--src", required=True, help="Path to raw UAV_VisLoc_dataset")
|
|
parser.add_argument("--dst", required=True, help="Path to output processed dataset")
|
|
parser.add_argument("--crop-size", type=int, default=512, help="Satellite crop size in px")
|
|
parser.add_argument("--stride", type=int, default=256, help="Crop stride in px")
|
|
parser.add_argument("--target-size", type=int, default=256, help="Final resize for both drone and crops")
|
|
parser.add_argument("--routes", nargs="*", default=None, help="Process only these routes (e.g. 01 02)")
|
|
return parser.parse_args()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. Read metadata
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def read_satellite_bbox(src: Path) -> dict:
|
|
"""Read satellite GPS bounding boxes from CSV."""
|
|
bbox = {}
|
|
csv_path = src / "satellite_ coordinates_range.csv"
|
|
with open(csv_path) as f:
|
|
reader = csv.DictReader(f)
|
|
for row in reader:
|
|
name = row["mapname"]
|
|
route = name.replace("satellite", "").replace(".tif", "")
|
|
bbox[route] = {
|
|
"lt_lat": float(row["LT_lat_map"]),
|
|
"lt_lon": float(row["LT_lon_map"]),
|
|
"rb_lat": float(row["RB_lat_map"]),
|
|
"rb_lon": float(row["RB_lon_map"]),
|
|
}
|
|
return bbox
|
|
|
|
|
|
def read_drone_metadata(src: Path, route: str) -> list[dict]:
|
|
"""Read drone GPS + pose from per-route CSV."""
|
|
csv_path = src / route / f"{route}.csv"
|
|
entries = []
|
|
with open(csv_path) as f:
|
|
reader = csv.DictReader(f)
|
|
for row in reader:
|
|
entries.append({
|
|
"filename": row["filename"],
|
|
"lat": float(row["lat"]),
|
|
"lon": float(row["lon"]),
|
|
"height": float(row["height"]),
|
|
})
|
|
return entries
|
|
|
|
|
|
def read_split(src: Path, split: str) -> set[str]:
|
|
"""Read train/test split CSV, return set of drone filenames."""
|
|
csv_path = src / f"visloc_{split}.csv"
|
|
filenames = set()
|
|
with open(csv_path) as f:
|
|
reader = csv.DictReader(f, delimiter="\t")
|
|
for row in reader:
|
|
fn = row["filename"]
|
|
basename = os.path.basename(fn)
|
|
filenames.add(basename)
|
|
return filenames
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. Resize drone images
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def resize_drone_images(src: Path, dst: Path, route: str, target_size: int) -> int:
|
|
"""Resize all drone images for a route to target_size x target_size."""
|
|
drone_src = src / route / "drone"
|
|
drone_dst = dst / route / "drone"
|
|
drone_dst.mkdir(parents=True, exist_ok=True)
|
|
|
|
count = 0
|
|
for img_name in sorted(os.listdir(drone_src)):
|
|
if not img_name.upper().endswith(".JPG"):
|
|
continue
|
|
src_path = drone_src / img_name
|
|
dst_path = drone_dst / img_name
|
|
|
|
if dst_path.exists():
|
|
count += 1
|
|
continue
|
|
|
|
with Image.open(src_path) as img:
|
|
img_resized = img.resize((target_size, target_size), Image.LANCZOS)
|
|
img_resized.save(dst_path, "JPEG", quality=95)
|
|
count += 1
|
|
|
|
return count
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. Stitch satellite tiles (route 09)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def stitch_route09(src: Path) -> Image.Image:
|
|
"""Stitch 4 satellite tiles for route 09 into a single image."""
|
|
tiles = {}
|
|
for (col, row), fname in TILE_LAYOUT_09.items():
|
|
tile_path = src / SPLIT_TILE_ROUTE / fname
|
|
tiles[(col, row)] = Image.open(tile_path)
|
|
|
|
# Compute full image dimensions
|
|
# Row 0: tiles (0,0) and (1,0) side by side
|
|
# Row 1: tiles (0,1) and (1,1) side by side
|
|
w0 = tiles[(0, 0)].width + tiles[(1, 0)].width
|
|
w1 = tiles[(0, 1)].width + tiles[(1, 1)].width
|
|
full_w = max(w0, w1)
|
|
|
|
h0 = max(tiles[(0, 0)].height, tiles[(1, 0)].height)
|
|
h1 = max(tiles[(0, 1)].height, tiles[(1, 1)].height)
|
|
full_h = h0 + h1
|
|
|
|
merged = Image.new("RGB", (full_w, full_h))
|
|
merged.paste(tiles[(0, 0)], (0, 0))
|
|
merged.paste(tiles[(1, 0)], (tiles[(0, 0)].width, 0))
|
|
merged.paste(tiles[(0, 1)], (0, h0))
|
|
merged.paste(tiles[(1, 1)], (tiles[(0, 1)].width, h0))
|
|
|
|
# Close tiles to free memory (~4.3 GB).
|
|
for tile in tiles.values():
|
|
tile.close()
|
|
|
|
return merged
|
|
|
|
|
|
def load_satellite(src: Path, route: str) -> Image.Image:
|
|
"""Load satellite map for a route (handles route 09 stitching)."""
|
|
if route == SPLIT_TILE_ROUTE:
|
|
return stitch_route09(src)
|
|
sat_path = src / route / f"satellite{route}.tif"
|
|
return Image.open(sat_path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 4. Crop satellite map
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def crop_satellite(
|
|
sat_img: Image.Image,
|
|
dst: Path,
|
|
route: str,
|
|
crop_size: int,
|
|
stride: int,
|
|
target_size: int,
|
|
) -> list[dict]:
|
|
"""Crop satellite image into patches and save resized versions.
|
|
|
|
Returns list of crop metadata: [{"name": "crop_X_Y.png", "x": X, "y": Y, "px_x": ..., "px_y": ...}]
|
|
"""
|
|
crop_dir = dst / route / "DB" / "img"
|
|
crop_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
w, h = sat_img.size
|
|
cols = max(0, (w - crop_size) // stride + 1)
|
|
rows = max(0, (h - crop_size) // stride + 1)
|
|
|
|
crops_meta = []
|
|
|
|
for cx in range(cols):
|
|
for cy in range(rows):
|
|
px_x = cx * stride
|
|
px_y = cy * stride
|
|
|
|
crop_name = f"crop_{cx}_{cy}.png"
|
|
crop_path = crop_dir / crop_name
|
|
|
|
if not crop_path.exists():
|
|
box = (px_x, px_y, px_x + crop_size, px_y + crop_size)
|
|
patch = sat_img.crop(box)
|
|
patch_resized = patch.resize((target_size, target_size), Image.LANCZOS)
|
|
patch_resized.save(crop_path, "PNG")
|
|
patch.close()
|
|
patch_resized.close()
|
|
|
|
crops_meta.append({
|
|
"name": crop_name,
|
|
"x": cx,
|
|
"y": cy,
|
|
"px_x": px_x,
|
|
"px_y": px_y,
|
|
})
|
|
|
|
return crops_meta
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 5. Compute GPS for crops + match drone -> crops
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def compute_crop_gps(crops_meta: list[dict], bbox: dict, sat_w: int, sat_h: int, crop_size: int) -> list[dict]:
|
|
"""Add GPS center coordinates to each crop metadata entry."""
|
|
lt_lat = bbox["lt_lat"]
|
|
lt_lon = bbox["lt_lon"]
|
|
rb_lat = bbox["rb_lat"]
|
|
rb_lon = bbox["rb_lon"]
|
|
|
|
for crop in crops_meta:
|
|
center_px_x = crop["px_x"] + crop_size / 2
|
|
center_px_y = crop["px_y"] + crop_size / 2
|
|
|
|
crop["lon"] = lt_lon + (center_px_x / sat_w) * (rb_lon - lt_lon)
|
|
crop["lat"] = lt_lat + (center_px_y / sat_h) * (rb_lat - lt_lat)
|
|
|
|
return crops_meta
|
|
|
|
|
|
def _frame_id(filename: str) -> str:
|
|
"""Extract frame ID from drone filename: '01_0001.JPG' -> '0001'."""
|
|
stem = os.path.splitext(filename)[0] # '01_0001'
|
|
return stem.split("_", 1)[1] # '0001'
|
|
|
|
|
|
def match_drone_to_crops(
|
|
drone_entries: list[dict],
|
|
crops_meta: list[dict],
|
|
) -> tuple[dict, dict]:
|
|
"""Match each drone image to its positive and semi-positive crops.
|
|
|
|
positive: crop with minimum GPS distance to drone center.
|
|
semi-positive: all crops within ±1 grid step from the positive crop.
|
|
|
|
Returns:
|
|
positive_map: {frame_id: [crop_name]}
|
|
semi_positive_map: {frame_id: [crop_name, ...]}
|
|
"""
|
|
# Build arrays for vectorized distance computation
|
|
crop_lats = np.array([c["lat"] for c in crops_meta])
|
|
crop_lons = np.array([c["lon"] for c in crops_meta])
|
|
|
|
# Grid lookup: (x, y) -> crop_name for O(1) neighbor search
|
|
grid = {}
|
|
for crop in crops_meta:
|
|
grid[(crop["x"], crop["y"])] = crop["name"]
|
|
|
|
positive_map = {}
|
|
semi_positive_map = {}
|
|
|
|
for drone in drone_entries:
|
|
d_lat, d_lon = drone["lat"], drone["lon"]
|
|
fid = _frame_id(drone["filename"])
|
|
|
|
# Compute distances to all crops (vectorized haversine)
|
|
dlat = np.radians(crop_lats - d_lat)
|
|
dlon = np.radians(crop_lons - d_lon)
|
|
a = np.sin(dlat / 2) ** 2 + np.cos(math.radians(d_lat)) * np.cos(np.radians(crop_lats)) * np.sin(dlon / 2) ** 2
|
|
dists = 6_371_000 * 2 * np.arcsin(np.sqrt(a))
|
|
|
|
# Positive: closest crop
|
|
best_idx = int(np.argmin(dists))
|
|
best_crop = crops_meta[best_idx]
|
|
positive_map[fid] = [best_crop["name"]]
|
|
|
|
# Semi-positives: ±1 neighbors in grid via dict lookup
|
|
bx, by = best_crop["x"], best_crop["y"]
|
|
semi = []
|
|
for dx in range(-1, 2):
|
|
for dy in range(-1, 2):
|
|
if dx == 0 and dy == 0:
|
|
continue
|
|
neighbor = grid.get((bx + dx, by + dy))
|
|
if neighbor is not None:
|
|
semi.append(neighbor)
|
|
|
|
semi_positive_map[fid] = semi
|
|
|
|
return positive_map, semi_positive_map
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 6. Write metadata files
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def write_positive_json(dst: Path, route: str, positive_map: dict):
|
|
out_path = dst / route / "positive.json"
|
|
with open(out_path, "w") as f:
|
|
json.dump(positive_map, f, indent=2)
|
|
|
|
|
|
def write_semi_positive_json(dst: Path, route: str, semi_positive_map: dict):
|
|
out_path = dst / route / "semi_positive.json"
|
|
with open(out_path, "w") as f:
|
|
json.dump(semi_positive_map, f, indent=2)
|
|
|
|
|
|
def write_db_position(dst: Path, route: str, crops_meta: list[dict],
|
|
sat_w: int, sat_h: int, bbox: dict):
|
|
"""Write db_postion.txt with GPS coordinates and scale for each crop.
|
|
|
|
Format matches UAV-GeoLoc: name\tlon\tlat\tscale_lon\tscale_lat (tab-separated).
|
|
scale_lon/scale_lat = degrees per pixel in the satellite map.
|
|
Note: filename uses original UAV-GeoLoc spelling 'postion' (sic).
|
|
"""
|
|
scale_lon = (bbox["rb_lon"] - bbox["lt_lon"]) / sat_w
|
|
scale_lat = (bbox["rb_lat"] - bbox["lt_lat"]) / sat_h
|
|
|
|
out_path = dst / route / "DB" / "db_postion.txt"
|
|
with open(out_path, "w") as f:
|
|
for crop in crops_meta:
|
|
f.write(f"{crop['name']}\t{crop['lon']:.8f}\t{crop['lat']:.8f}"
|
|
f"\t{scale_lon:.2e}\t{scale_lat:.2e}\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 7. Generate Index files
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def generate_index_files(
|
|
dst: Path,
|
|
all_routes: list[str],
|
|
route_data: dict,
|
|
train_files: set[str],
|
|
test_files: set[str],
|
|
):
|
|
"""Generate train/test query and DB index files in UAV-GeoLoc format.
|
|
|
|
train_query.txt format:
|
|
route/drone/filename label positive_crop1 positive_crop2 ...
|
|
train_db.txt format:
|
|
route/DB/img/crop_name
|
|
"""
|
|
index_dir = dst / "Index"
|
|
index_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
train_query_lines = []
|
|
test_query_lines = []
|
|
all_db_set = set()
|
|
|
|
for route in all_routes:
|
|
data = route_data[route]
|
|
positive_map = data["positive_map"]
|
|
semi_positive_map = data["semi_positive_map"]
|
|
crops_meta = data["crops_meta"]
|
|
|
|
# All DB crops for this route go into the gallery
|
|
for crop in crops_meta:
|
|
db_path = f"{route}/DB/img/{crop['name']}"
|
|
all_db_set.add(db_path)
|
|
|
|
for drone_entry in data["drone_entries"]:
|
|
fname = drone_entry["filename"]
|
|
fid = _frame_id(fname)
|
|
drone_path = f"{route}/drone/{fname}"
|
|
|
|
positives = positive_map.get(fid, [])
|
|
semi_pos = semi_positive_map.get(fid, [])
|
|
all_pos = positives + semi_pos
|
|
|
|
pos_paths = [f"{route}/DB/img/{p}" for p in all_pos]
|
|
line = f"{drone_path} 0 {' '.join(pos_paths)}"
|
|
|
|
if fname in train_files:
|
|
train_query_lines.append(line)
|
|
elif fname in test_files:
|
|
test_query_lines.append(line)
|
|
|
|
# DB gallery is the same for train and test (split is by query, not by route)
|
|
sorted_db = sorted(all_db_set)
|
|
|
|
# Write files
|
|
_write_lines(index_dir / "train_query.txt", train_query_lines)
|
|
_write_lines(index_dir / "test_query.txt", test_query_lines)
|
|
_write_lines(index_dir / "train_db.txt", sorted_db)
|
|
_write_lines(index_dir / "test_db.txt", sorted_db)
|
|
_write_lines(index_dir / "all_db.txt", sorted_db)
|
|
|
|
# Route lists
|
|
_write_lines(index_dir / "train.txt", all_routes)
|
|
_write_lines(index_dir / "test.txt", all_routes)
|
|
|
|
return {
|
|
"train_queries": len(train_query_lines),
|
|
"test_queries": len(test_query_lines),
|
|
"total_db": len(all_db_set),
|
|
}
|
|
|
|
|
|
def _write_lines(path: Path, lines: list[str]):
|
|
with open(path, "w") as f:
|
|
for line in lines:
|
|
f.write(line + "\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def process_route(
|
|
src: Path,
|
|
dst: Path,
|
|
route: str,
|
|
bbox: dict,
|
|
crop_size: int,
|
|
stride: int,
|
|
target_size: int,
|
|
) -> dict:
|
|
"""Process a single route: resize drone, crop satellite, match pairs."""
|
|
print(f"\n{'='*60}")
|
|
print(f"Route {route}")
|
|
print(f"{'='*60}")
|
|
|
|
# 1. Read drone metadata
|
|
drone_entries = read_drone_metadata(src, route)
|
|
print(f" Drone images: {len(drone_entries)}")
|
|
|
|
# 2. Resize drone images
|
|
n_drone = resize_drone_images(src, dst, route, target_size)
|
|
print(f" Drone resized: {n_drone}")
|
|
|
|
# 3. Load and crop satellite
|
|
print(f" Loading satellite map...")
|
|
sat_img = load_satellite(src, route)
|
|
sat_w, sat_h = sat_img.size
|
|
print(f" Satellite size: {sat_w}x{sat_h}")
|
|
|
|
crops_meta = crop_satellite(sat_img, dst, route, crop_size, stride, target_size)
|
|
print(f" Satellite crops: {len(crops_meta)}")
|
|
|
|
# Free memory
|
|
del sat_img
|
|
|
|
# 4. Compute GPS for crops
|
|
route_bbox = bbox[route]
|
|
crops_meta = compute_crop_gps(crops_meta, route_bbox, sat_w, sat_h, crop_size)
|
|
|
|
# 5. Match drone -> crops
|
|
positive_map, semi_positive_map = match_drone_to_crops(drone_entries, crops_meta)
|
|
print(f" Positives matched: {len(positive_map)}")
|
|
|
|
# Stats
|
|
avg_semi = sum(len(v) for v in semi_positive_map.values()) / max(1, len(semi_positive_map))
|
|
print(f" Avg semi-positives per drone: {avg_semi:.1f}")
|
|
|
|
# 6. Write metadata
|
|
write_positive_json(dst, route, positive_map)
|
|
write_semi_positive_json(dst, route, semi_positive_map)
|
|
write_db_position(dst, route, crops_meta, sat_w, sat_h, route_bbox)
|
|
print(f" Metadata written.")
|
|
|
|
return {
|
|
"drone_entries": drone_entries,
|
|
"crops_meta": crops_meta,
|
|
"positive_map": positive_map,
|
|
"semi_positive_map": semi_positive_map,
|
|
}
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
src = Path(args.src)
|
|
dst = Path(args.dst)
|
|
|
|
print(f"Source: {src}")
|
|
print(f"Destination: {dst}")
|
|
print(f"Crop: {args.crop_size}x{args.crop_size}, stride: {args.stride}")
|
|
print(f"Target size: {args.target_size}x{args.target_size}")
|
|
print(f"Excluded routes: {EXCLUDED_ROUTES}")
|
|
|
|
dst.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Read satellite bounding boxes
|
|
bbox = read_satellite_bbox(src)
|
|
|
|
# Determine routes to process
|
|
all_routes = sorted([
|
|
d for d in os.listdir(src)
|
|
if os.path.isdir(src / d) and d.isdigit() and d not in EXCLUDED_ROUTES
|
|
])
|
|
if args.routes:
|
|
all_routes = [r for r in all_routes if r in args.routes]
|
|
|
|
print(f"Routes to process: {all_routes}")
|
|
|
|
# Read train/test splits
|
|
train_files = read_split(src, "train")
|
|
test_files = read_split(src, "test")
|
|
print(f"Train split: {len(train_files)} files")
|
|
print(f"Test split: {len(test_files)} files")
|
|
|
|
# Process each route
|
|
route_data = {}
|
|
total_drone = 0
|
|
total_crops = 0
|
|
|
|
for route in all_routes:
|
|
data = process_route(src, dst, route, bbox, args.crop_size, args.stride, args.target_size)
|
|
route_data[route] = data
|
|
total_drone += len(data["drone_entries"])
|
|
total_crops += len(data["crops_meta"])
|
|
|
|
# Generate Index files
|
|
print(f"\n{'='*60}")
|
|
print("Generating Index files...")
|
|
idx_stats = generate_index_files(dst, all_routes, route_data, train_files, test_files)
|
|
print(f" Train queries: {idx_stats['train_queries']}")
|
|
print(f" Test queries: {idx_stats['test_queries']}")
|
|
print(f" DB gallery (all crops): {idx_stats['total_db']}")
|
|
|
|
# Summary
|
|
print(f"\n{'='*60}")
|
|
print("SUMMARY")
|
|
print(f"{'='*60}")
|
|
print(f" Routes processed: {len(all_routes)}")
|
|
print(f" Total drone images (resized to {args.target_size}): {total_drone}")
|
|
print(f" Total satellite crops ({args.crop_size}→{args.target_size}): {total_crops}")
|
|
print(f" Output: {dst}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|