Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
223 lines
6.7 KiB
Markdown
223 lines
6.7 KiB
Markdown
# GeoLoc-UAV (GeoLoc-UAV-main)
|
|
|
|
Research code for UAV geo-localization / cross-view image retrieval on the **UAV-GeoLoc (World-UAV)** dataset.
|
|
|
|
- Dataset: available on Hugging Face: `https://huggingface.co/datasets/RingoWRW97/UAV-GeoLoc`
|
|
- This folder (`GeoLoc-UAV-main`) contains training/evaluation scripts and model definitions.
|
|
|
|
## What this code does
|
|
|
|
The task is formulated as **retrieval**: given a UAV query image, retrieve the matching geo-referenced database (DB) image(s).
|
|
|
|
- **Training**: contrastive classification via InfoNCE (implemented as `CrossEntropyLoss` on similarity matrix).
|
|
- **Evaluation**: extract global descriptors for queries and DB, then run FAISS `IndexFlatL2` search and report **Top-1 / Top-5 / Top-10** accuracy.
|
|
|
|
Two main modes exist in the code:
|
|
|
|
- **`vanilia`** (spelling in code): a standard CNN/ViT backbone (`resnet18`, `dinov2_*`) + an aggregation head (`multiconvap`, `convap`, optional `LPN`).
|
|
- **`group`**: a GroupNet-style encoder that uses a set of transformed views + point grids (scale/rotate sampling) before aggregation.
|
|
|
|
## Project structure
|
|
|
|
Key entry points:
|
|
|
|
- `train_vanilia.py`: train a vanilla backbone (ResNet) retrieval model.
|
|
- `train_vanilia_dino.py`: same, but with a DINOv2 backbone.
|
|
- `train_group.py`: train the group-based model.
|
|
- `train_group_dino.py`: group-based model with DINOv2 features.
|
|
- `preprocess_data.py`: helper to generate train index files from scene lists.
|
|
|
|
Evaluation:
|
|
|
|
- `eval_simidataset_parser.py`: evaluate on World-UAV-style splits (reads a list of scene folders).
|
|
- `eval_real_dataset.py`: evaluate on a “real” dataset layout (query_images/reference_images + gt CSV/NPY).
|
|
- `eval_denseuav.py`: evaluate on DenseUAV-style lists (query.txt/db.txt/gt.txt).
|
|
- `eval_real.sh`, `eval_rot.sh`: example command lines (paths are author-specific).
|
|
|
|
Core modules:
|
|
|
|
- `dataset/World.py`: dataset loaders for World-UAV, “real” and DenseUAV layouts.
|
|
- `models/`: backbones, aggregators, group networks.
|
|
- `eval/eval.py`: feature extraction + FAISS retrieval metrics.
|
|
|
|
## Installation
|
|
|
|
This repository does not ship a pinned `requirements.txt`. A typical working environment:
|
|
|
|
```bash
|
|
python -m venv .venv
|
|
source .venv/bin/activate
|
|
|
|
pip install -U pip
|
|
pip install torch torchvision transformers tensorboard tqdm pillow numpy pandas h5py matplotlib opencv-python
|
|
|
|
# FAISS (choose one):
|
|
pip install faiss-cpu
|
|
# or: pip install faiss-gpu (if your platform provides it)
|
|
```
|
|
|
|
Notes:
|
|
|
|
- Most scripts assume CUDA is available; CPU should work for evaluation but may be slow.
|
|
- Mixed precision is enabled by default in training scripts.
|
|
|
|
## Dataset format (World-UAV style)
|
|
|
|
For the World-UAV dataset loaders (`WorldDatasetEvalVanilia` / `WorldDatasetEvalGroup`), the code expects a **scene folder** with:
|
|
|
|
```text
|
|
<dataset_root_dir>/
|
|
<scene_name>/
|
|
positive.json
|
|
semi_positive.json
|
|
DB/
|
|
img/
|
|
*.png
|
|
query/
|
|
height100_rot0/footage/*.jpeg
|
|
height100_rot45/footage/*.jpeg
|
|
...
|
|
height150_rot315/footage/*.jpeg
|
|
```
|
|
|
|
Ground truth:
|
|
|
|
- `positive.json`: maps each query key to a list of positive DB image filenames (under `DB/img/`).
|
|
- `semi_positive.json`: optional extra positives (the code includes them in GT if present).
|
|
|
|
Important:
|
|
|
|
- Query images are collected across multiple height/rotation folders. The evaluator replicates GT to match the number of query images extracted (see `eval/eval.py`: `multi_num = ql.shape[0] / len(pos_gt)`).
|
|
|
|
## Training index files (World-UAV style)
|
|
|
|
Training loaders (`WorldDatasetTrainVanilia`, `WorldDatasetTrainGroup`) read a text file where each line contains:
|
|
|
|
```text
|
|
<relative_query_path> <label_int> <relative_db_path>
|
|
```
|
|
|
|
Example (paths are **relative** to `dataset_root_dir`):
|
|
|
|
```text
|
|
some_scene/query/height100_rot0/footage/height100_rot0_0001.jpeg 12 some_scene/DB/img/000123.png
|
|
```
|
|
|
|
The helper script `preprocess_data.py` can generate `*_query_all.txt` and `*_db_all.txt` from a file containing scene names (one per line). It is currently written with author-specific absolute paths and may require edits to:
|
|
|
|
- `root`
|
|
- `txt` (scene list file)
|
|
- `save_path`
|
|
|
|
## Quick start (training)
|
|
|
|
All training scripts contain a `Configuration` dataclass with hardcoded paths like:
|
|
|
|
- `dataset_root_dir`
|
|
- `train_query_txt`
|
|
- `val_index_txt`
|
|
- `test_index_txt`
|
|
|
|
Update them to your local paths before running.
|
|
|
|
### Vanilla (ResNet)
|
|
|
|
```bash
|
|
python train_vanilia.py
|
|
```
|
|
|
|
### Vanilla (DINOv2 backbone)
|
|
|
|
```bash
|
|
python train_vanilia_dino.py
|
|
```
|
|
|
|
### Group model
|
|
|
|
```bash
|
|
python train_group.py
|
|
```
|
|
|
|
### Group model (DINO variant)
|
|
|
|
```bash
|
|
python train_group_dino.py
|
|
```
|
|
|
|
Checkpoints:
|
|
|
|
- Training scripts create a timestamped directory under `config.model_path/config.model/<HHMMSS>/`
|
|
- The best checkpoint (by average Top-1 on validation scenes) is saved as `weights_e{epoch}_{score}.pth`
|
|
|
|
TensorBoard:
|
|
|
|
- Vanilla scripts write under `world_vanillia/...`
|
|
- Group script writes under `world/...`
|
|
|
|
## Quick start (evaluation)
|
|
|
|
### Evaluate on World-UAV scenes list
|
|
|
|
Use `eval_simidataset_parser.py`. It reads a text file where each line is a scene folder name (relative to `--dataset_root`):
|
|
|
|
```bash
|
|
python eval_simidataset_parser.py \
|
|
--mode vanilia \
|
|
--dataset_root "/path/to/WorldLoc" \
|
|
--test_txt "/path/to/WorldLoc/Index/test.txt" \
|
|
--save_txt "/tmp/results.txt" \
|
|
--checkpoint_path "/path/to/weights.pth" \
|
|
--backbone_arch resnet18 \
|
|
--pretrain_flag False \
|
|
--agg_in_channels 512 \
|
|
--agg_out_channels 512 \
|
|
--agg_LPN False
|
|
```
|
|
|
|
For the group model:
|
|
|
|
```bash
|
|
python eval_simidataset_parser.py \
|
|
--mode group \
|
|
--dataset_root "/path/to/WorldLoc" \
|
|
--test_txt "/path/to/WorldLoc/Index/test.txt" \
|
|
--save_txt "/tmp/results.txt" \
|
|
--checkpoint_path "/path/to/weights.pth" \
|
|
--agg_in_channels 256 \
|
|
--agg_out_channels 256
|
|
```
|
|
|
|
### Evaluate on DenseUAV lists
|
|
|
|
```bash
|
|
python eval_denseuav.py \
|
|
--mode vanilia \
|
|
--dataset_query "/path/to/query.txt" \
|
|
--dataset_db "/path/to/db.txt" \
|
|
--dataset_gt "/path/to/gt.txt" \
|
|
--checkpoint_path "/path/to/weights.pth"
|
|
```
|
|
|
|
### Evaluate on “real” query/reference folder layout
|
|
|
|
```bash
|
|
python eval_real_dataset.py \
|
|
--mode vanilia \
|
|
--dataset_root_dir "/path/to/test_set" \
|
|
--checkpoint_path "/path/to/weights.pth" \
|
|
--save_dir_path "/tmp/geoloc-uav-real-eval"
|
|
```
|
|
|
|
## Common pitfalls / required path fixes
|
|
|
|
- **Hardcoded absolute paths**: many scripts use `/media/...` paths. Replace them with your local paths.
|
|
- **Transform config path is hardcoded in code**:
|
|
- `dataset/World.py` and `models/group/groupnet_dino.py` load `transform_config.json` via an absolute `json_path`.
|
|
- For portability, point it to `configs/transform_config.json` in this repository.
|
|
- **Mode spelling**: the scripts use `vanilia` (not `vanilla`).
|
|
|
|
## Demo
|
|
|
|

|
|
|