From fd014a3155f12e40d8fd2c2b08711a4cf7c5dc59 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 17 Apr 2026 17:12:43 +0300 Subject: [PATCH] Dataset preparation script with documentation - 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) --- CLAUDE.md | 40 +++ README.md | 132 +++++++++ scripts/prepare_dataset.py | 559 +++++++++++++++++++++++++++++++++++++ 3 files changed, 731 insertions(+) create mode 100644 CLAUDE.md create mode 100644 README.md create mode 100644 scripts/prepare_dataset.py diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..f0207cc --- /dev/null +++ b/CLAUDE.md @@ -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) — косинусный эффект широты, не разная высота съёмки. diff --git a/README.md b/README.md new file mode 100644 index 0000000..1d3db57 --- /dev/null +++ b/README.md @@ -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. diff --git a/scripts/prepare_dataset.py b/scripts/prepare_dataset.py new file mode 100644 index 0000000..dfdf1e0 --- /dev/null +++ b/scripts/prepare_dataset.py @@ -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()