Pikaliov/UAV-VisLoc-prepare
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Dataset preparation script with documentation

Browse Source 
- 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>
This commit is contained in:
pikaliov
2026-04-17 17:12:43 +03:00
commit fd014a3155
3 changed files with 731 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
CLAUDE.md Normal file
View File

@@ -0,0 +1,40 @@
# UAV-VisLoc Dataset Preparation
## Пути
- **Исходный датасет:** `/home/servml/Документы/datasets/UAV_VisLoc_dataset/`
- **Обработанный датасет:** `/home/servml/Документы/datasets/UAV_VisLoc_processed/`
- **Скрипт:** `scripts/prepare_dataset.py`
- **Статус:** выполнен, данные готовы (2026-04-17)
## Результаты обработки
- Drone: 6,744 изображений resized 256x256
- Satellite кропов: 74,807 (512x512 -> 256x256)
- Train queries: 5,060 / Test queries: 1,684
- Gallery: 74,807 кропов (одинаковая для train и test)
- Median distance drone->crop: 25.9m, P99: 45.7m
## Формат данных
### Разделение positive / semi-positive / negative
В формате UAV-GeoLoc нет явного списка negative — negatives определяются неявно:
- **Positive** (1 на drone): ближайший crop по GPS. Хранится в `positive.json` и как ПЕРВЫЙ crop в строке `train_query.txt`
- **Semi-positive** (8 на drone): +-1 соседи positive crop в grid. Хранятся в `semi_positive.json` и как crops 2-9 в строке `train_query.txt`
- **Negative** (implicit): ВСЕ остальные кропы в gallery (~74K). Не хранятся отдельно — при contrastive learning in-batch negatives формируются из других пар в batch
### Как это используется в train_query.txt
```
01/drone/01_0001.JPG 0 01/DB/img/crop_5_18.png 01/DB/img/crop_4_17.png 01/DB/img/crop_4_18.png ...
^ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
label 1 positive + 8 semi-positives (все считаются positive matches)
```
Все 9 кропов в строке — positives для данного query. Всё что НЕ в этом списке — negatives.
## Известные проблемы
- Маршрут 07 исключён (спутник 3000x170, слишком узкий)
- 6 drone в маршруте 06 (06_0093-06_0098) за пределами спутниковой карты (distance >1000m)
- Нет val split (только train/test как в оригинальном UAV-VisLoc)
## GSD спутника
~0.30 м/px (единый zoom level). Вариации GSD по долготе (0.23-0.27 м/px) — косинусный эффект широты, не разная высота съёмки.

132
README.md Normal file
View File

@@ -0,0 +1,132 @@
# UAV-VisLoc Dataset Preparation
Prepare UAV-VisLoc dataset for cross-view geo-localization retrieval training.
Generates satellite crops, GPS-matched drone-crop pairs, and Index files
compatible with UAV-GeoLoc format.
## Pipeline
```
UAV_VisLoc_dataset/ UAV_VisLoc_processed/
├── 01/ ├── 01/
│ ├── drone/*.JPG (3976x2652) ---> │ ├── drone/*.JPG (256x256)
│ ├── satellite01.tif ---> │ ├── DB/img/crop_X_Y.png (256x256)
│ └── 01.csv │ ├── DB/db_postion.txt
│ │ ├── positive.json
├── ... │ └── semi_positive.json
├── 09/ ├── ...
│ ├── satellite09_01-01.tif --+ ├── 09/ (stitched from 4 tiles)
│ ├── satellite09_01-02.tif +--> │ ├── DB/img/ (~22K crops)
│ ├── satellite09_02-01.tif | │ └── ...
│ └── satellite09_02-02.tif --+ │
├── satellite_ coordinates_range.csv └── Index/
├── visloc_train.csv ├── train_query.txt
└── visloc_test.csv ├── test_query.txt
├── train_db.txt
├── test_db.txt
└── all_db.txt
```
## Quick Start
```bash
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
```
To process specific routes only:
```bash
python scripts/prepare_dataset.py --src ... --dst ... --routes 01 02 03
```
## Steps
1. **Resize drone images** -> 256x256 JPEG (quality=95)
2. **Stitch satellite tiles** for route 09 (4 tiles -> 44800x33280)
3. **Crop satellite maps** -> 512x512 patches, stride 256 (50% overlap), resize -> 256x256 PNG
4. **Compute GPS** for each crop center from satellite bbox + grid position
5. **Match drone -> crops** via vectorized haversine (positive = closest, semi-positive = +-1 in grid)
6. **Write metadata**: positive.json, semi_positive.json, db_postion.txt (per route)
7. **Generate Index files**: train_query.txt, test_query.txt, train_db.txt, test_db.txt, all_db.txt
## Output Format (UAV-GeoLoc compatible)
### Index files
**train_query.txt / test_query.txt:**
```
01/drone/01_0001.JPG 0 01/DB/img/crop_5_18.png 01/DB/img/crop_4_17.png ...
```
Format: `query_path label positive_crop semi_positive_crops...`
**train_db.txt / test_db.txt / all_db.txt:**
```
01/DB/img/crop_0_0.png
01/DB/img/crop_0_1.png
...
```
Full gallery (all 74,807 crops), identical for train and test (split is by query).
### Per-route metadata
**positive.json:**
```json
{"0001": ["crop_5_18.png"], "0002": ["crop_5_19.png"], ...}
```
Keys are frame IDs (without route prefix).
**semi_positive.json:**
```json
{"0001": ["crop_4_17.png", "crop_4_18.png", ...], ...}
```
8 neighbors (+-1 in grid) of the positive crop.
**db_postion.txt** (tab-separated, matching UAV-GeoLoc spelling):
```
crop_0_0.png 115.97197337 29.77349180 2.68e-06 -2.68e-06
```
Columns: name, longitude, latitude, scale_lon (deg/px), scale_lat (deg/px).
## Dataset Statistics
| Route | Drone | Crops | Region | Satellite (px) |
|-------|-------|-------|--------|----------------|
| 01 | 817 | 3,811 | Changjiang | 9774x26762 |
| 02 | 1,071 | 5,676 | Changjiang | 11482x34291 |
| 03 | 768 | 12,648 | Taizhou | 35092x24308 |
| 04 | 738 | 10,281 | Taizhou | 18093x38408 |
| 05 | 473 | 805 | Yunnan | 9394x6144 |
| 06 | 344 | 1,110 | Zhuxi | 8082x9780 |
| 07 | -- | -- | Excluded | 3000x170 (too narrow) |
| 08 | 1,033 | 10,416 | Huzhou | 43421x16294 |
| 09 | 766 | 22,446 | Huzhou | 44800x33280 (stitched) |
| 10 | 144 | 432 | Huailai | 6593x5077 |
| 11 | 590 | 7,182 | Shandan | 29592x16582 |
| **Total** | **6,744** | **74,807** | | |
Split: 5,060 train / 1,684 test queries. Gallery: 74,807 crops (shared).
## GPS Matching Quality
- Median distance drone -> positive crop: **25.9m**
- P99: **45.7m**
- Known issue: 6 drones in route 06 (06_0093-06_0098) are outside satellite coverage (~1,091m to nearest crop)
## Satellite Resolution
All maps: ~0.30 m/pixel (single Google Earth zoom level).
One crop 512x512 covers ~154x154m on the ground.
## Requirements
```
numpy
Pillow
```
## Memory
Peak RAM usage: ~8.7 GB (route 09 stitching: 4 tiles + merged image).
Other routes: 1-3 GB.

559
scripts/prepare_dataset.py Normal file
View File

@@ -0,0 +1,559 @@
#!/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()