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>

GeoLoc-UAV-main/.gitignore

@@ -0,0 +1 @@
+*.pyc

GeoLoc-UAV-main/README.md

@@ -0,0 +1,222 @@
+# 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
+
+![demo](demo_small.gif)
+

GeoLoc-UAV-main/configs/configs_dino.json

@@ -0,0 +1,26 @@
+{
+    "default_backbone_config": {
+        "backbone_arch": "dinov2_vitb14"
+    },
+    "default_agg_config":{
+        "agg_arch": "convap",
+        "agg_config":{
+            "in_channels": 768, 
+            "out_channels": 768, 
+            "s1": 1,
+            "s2": 1
+        }
+    },
+
+    "dataset_root_dir": "/media/guan/新加卷/EdgeBing/WorldLoc/",
+    "test_index_txt": "/media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt",
+    "test_rank_gt": "/media/guan/新加卷/EdgeBing/WorldLoc/test/test_rerank.txt",
+    "write_path": "/media/guan/新加卷/EdgeBing/WorldLoc/test",
+
+    "num_workers": 0,
+    "batch_size": 16,
+    "custom_sampling": "True",
+    "verbose":"True",
+
+    "device":"cuda"
+}

GeoLoc-UAV-main/configs/configs_group.json

@@ -0,0 +1,27 @@
+{
+    "default_group_config": {
+          "group_arch" : "groupnet",
+        "group_config": "none"
+    },
+    "default_agg_config":{
+        "agg_arch": "convap",
+        "agg_config":{
+            "in_channels": 256, 
+            "out_channels": 256, 
+            "s1": 1,
+            "s2": 1
+        }
+    },
+
+    "dataset_root_dir": "/media/guan/新加卷/EdgeBing/WorldLoc/",
+    "test_index_txt": "/media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt",
+    "test_rank_gt": "/media/guan/新加卷/EdgeBing/WorldLoc/test/test_rerank.txt",
+    "write_path": "/media/guan/新加卷/EdgeBing/WorldLoc/test",
+
+    "num_workers": 0,
+    "batch_size": 1,
+    "custom_sampling": "True",
+    "verbose":"True",
+
+    "device":"cuda"
+}

GeoLoc-UAV-main/configs/loftr.yml

@@ -0,0 +1,50 @@
+default: &default
+    class: 'LoFTR'
+    ckpt: '/home/guan/image-matching-toolbox/pretrained/loftr/outdoor_ds.ckpt'
+    match_threshold: 0.2
+    imsize: -1
+    no_match_upscale: False
+    eval_coarse: False
+example:
+    <<: *default
+    match_threshold: 0.5
+    imsize: -1
+hpatch:        
+    <<: *default
+    imsize: 480
+    no_match_upscale: True
+megadepth:        
+    <<: *default
+    imsize: 1024
+inloc:
+    <<: *default
+    match_threshold: 0.5
+    npts: 4096
+    imsize: 1024
+    pairs: 'pairs-query-netvlad40-temporal.txt'
+    rthres: 48
+    skip_matches: 20
+airloc:
+    <<: *default
+    match_threshold: 0.0 # Save all matches
+    pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt']
+    npts: 4096
+    imsize: 1024
+    qt_dthres: 4
+    qt_psize: 48
+    qt_unique: True
+    ransac_thres: [20]
+    sc_thres: 0.2 # Filtering during quantization
+    covis_cluster: True
+aachen:
+    <<: *default
+    match_threshold: 0.0 # Save all matches
+    pairs: ['pairs-db-covis20.txt', 'pairs-query-netvlad50.txt']
+    npts: 4096
+    imsize: 1024
+    qt_dthres: 4
+    qt_psize: 48
+    qt_unique: True
+    ransac_thres: [20]
+    sc_thres: 0.2 # Filtering during quantization
+    covis_cluster: True

GeoLoc-UAV-main/configs/transform_config copy.json

@@ -0,0 +1,11 @@
+{
+"transform_config" : {
+    "sample_scale_begin": 0,
+ "sample_scale_inter": 0.5, 
+ "sample_scale_num": 3, 
+ "sample_rotate_begin": 0,
+ "sample_rotate_inter": 90, 
+ "sample_rotate_num": 4
+}
+
+}

GeoLoc-UAV-main/configs/transform_config.json

@@ -0,0 +1,11 @@
+{
+"transform_config" : {
+    "sample_scale_begin": 0,
+ "sample_scale_inter": 0.5, 
+ "sample_scale_num": 3, 
+ "sample_rotate_begin": 0,
+ "sample_rotate_inter": 90, 
+ "sample_rotate_num": 4
+}
+
+}

GeoLoc-UAV-main/dataset/World.py

@@ -0,0 +1,964 @@
+import os
+import cv2
+import numpy as np
+from PIL import Image, UnidentifiedImageError
+from torch.utils.data import Dataset
+import copy
+from tqdm import tqdm
+import time
+import random
+import glob
+import json
+import pandas as pd
+
+from torch.utils.data import DataLoader
+import torchvision.transforms as T
+import json
+
+
+from utils.utils import TransformerCV
+
+# transform_config = {
+#      "sample_scale_begin": 0,
+#   "sample_scale_inter": 0.5, 
+#   "sample_scale_num": 3, 
+#   "sample_rotate_begin": 0,
+#   "sample_rotate_inter": 45, 
+#   "sample_rotate_num": 8,   
+# }
+json_path = "/media/guan/新加卷/Code/Code/configs/transform_config.json"
+with open(json_path, 'r', encoding='utf-8') as file:
+    data = json.load(file)
+transform_config = data["transform_config"]
+
+
+# transform_config = {
+#      "sample_scale_begin": 0,
+#   "sample_scale_inter": 0.5, 
+#   "sample_scale_num": 1, 
+#   "sample_rotate_begin": 0,
+#   "sample_rotate_inter": 0, 
+#   "sample_rotate_num": 1,   
+# }
+
+default_transform = T.Compose([
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+def get_data(txt):
+
+    data = {}
+    idx = 0
+    with open(txt, 'r') as f:
+        for line in f:
+            line_list = line.split(' ')[:-1]
+            data[idx] = line_list
+            idx += 1
+
+    return data
+
+class WorldDatasetTrainGroup(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][2])
+            self.pairs.append((label, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+    
+    def __getitem__(self, index):
+        
+        idx, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = self.image_loader(query_img_path)
+        # db
+        db_img = self.image_loader(db_img_path)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(query_img)
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(db_img)
+        
+        # return query_img, db_img, idx
+        # group
+        query_img *= 255
+        query_img, query_pt = self.transformImg(query_img)
+
+        db_img *= 255
+        db_img, db_pt = self.transformImg(db_img)
+        
+        return query_img, query_pt, db_img, db_pt, idx
+    
+    def transformImg(self, img):
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+        # print("First Element ID: {} - Last Element ID: {}".format(self.samples[0][0], self.samples[-1][0]))  
+
+class WorldDatasetTrainVanilia(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][2])
+            self.pairs.append((label, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+    
+    def __getitem__(self, index):
+        
+        idx, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = self.image_loader(query_img_path)
+        # db
+        db_img = self.image_loader(db_img_path)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(query_img)
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(db_img)
+        
+        return query_img, db_img, idx
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+
+class WorldDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 name,
+                 mode,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.transforms = transforms
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+
+        self.data_dir = data_dir
+        self.name = name 
+
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        positive = json.load(open(pos_json_path))
+
+        self.samples = []
+        if mode == 'query':
+            height_list = ["height100_rot0", "height100_rot45", "height100_rot90", "height100_rot135", "height100_rot180", "height100_rot225", "height100_rot270", "height100_rot315", 
+                           "height125_rot0", "height125_rot45", "height125_rot90", "height125_rot135", "height125_rot180", "height125_rot225", "height125_rot270", "height125_rot315",
+                            "height150_rot0", "height150_rot45", "height150_rot90", "height150_rot135", "height150_rot180", "height150_rot225", "height150_rot270", "height150_rot315"]
+            # height_list = ["height100_rot20", "height100_rot60", "height100_rot150", "height100_rot210"]
+            for i in height_list:
+                if os.path.exists(os.path.join(data_dir, name,'query',  i, 'footage')):
+                    temp_path = os.path.join(data_dir, name,'query', i, 'footage')
+                    temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
+                    if len(temp) != len(positive.keys()):
+                        filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
+                        self.samples.extend(filter_temp)
+                    else:
+                        self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir, name, 'DB', 'img')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+      
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        semi_pos_json_path = os.path.join(self.data_dir, self.name,  'semi_positive.json')
+        positive = json.load(open(pos_json_path))
+        semi_positive = json.load(open(semi_pos_json_path))
+
+        pos_gt = []
+        for key in positive.keys():
+            value = positive[key]
+            
+            temp_index = []
+            # pos
+            for one_value in value:
+                temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                temp_path = temp_path_dir + '/' + one_value
+                one_index = self.samples.index(temp_path)
+                temp_index.append(one_index)
+            # semi-pos
+            try:
+                semi_value = semi_positive[key]
+                for one_value in semi_value:
+                    temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                    temp_path = temp_path_dir + '/' + one_value
+                    one_index = self.samples.index(temp_path)
+                    temp_index.append(one_index)
+            except:
+                pos_gt.append([key, temp_index])
+                continue
+
+            pos_gt.append([key, temp_index])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class WorldDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 name,
+                 mode,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.transforms = transforms
+
+        self.data_dir = data_dir
+        self.name = name 
+
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        positive = json.load(open(pos_json_path))
+
+        self.samples = []
+        if mode == 'query':
+            height_list = ["height100_rot0", "height100_rot45", "height100_rot90", "height100_rot135", "height100_rot180", "height100_rot225", "height100_rot270", "height100_rot315", 
+                           "height125_rot0", "height125_rot45", "height125_rot90", "height125_rot135", "height125_rot180", "height125_rot225", "height125_rot270", "height125_rot315",
+                            "height150_rot0", "height150_rot45", "height150_rot90", "height150_rot135", "height150_rot180", "height150_rot225", "height150_rot270", "height150_rot315"]
+            # height_list = ["height100_rot20", "height100_rot60", "height100_rot150", "height100_rot210"]
+            for i in height_list:
+                if os.path.exists(os.path.join(data_dir, name,'query',  i, 'footage')): #query
+                    temp_path = os.path.join(data_dir, name,'query', i, 'footage')
+                    temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
+                    if len(temp) != len(positive.keys()):
+                        filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
+                        self.samples.extend(filter_temp)
+                    else:
+                        self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir, name, 'DB', 'img')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+      
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        semi_pos_json_path = os.path.join(self.data_dir, self.name,  'semi_positive.json')
+        positive = json.load(open(pos_json_path))
+        semi_positive = json.load(open(semi_pos_json_path))
+
+        pos_gt = []
+        for key in positive.keys():
+            value = positive[key]
+            
+            temp_index = []
+            # pos
+            for one_value in value:
+                temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                temp_path = temp_path_dir + '/' + one_value
+                one_index = self.samples.index(temp_path)
+                temp_index.append(one_index)
+            try:
+                semi_value = semi_positive[key]
+            # semi-pos
+                for one_value in semi_value:
+                    temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                    temp_path = temp_path_dir + '/' + one_value
+                    one_index = self.samples.index(temp_path)
+                    temp_index.append(one_index)
+            except:
+                pos_gt.append([key, temp_index])
+                continue
+
+            pos_gt.append([key, temp_index])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class AerialDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 mode,
+                 angle=0,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            temp_path = os.path.join(data_dir,  'query_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = angle
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir,  'reference_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = angle
+
+        self.transforms = transforms
+        self.data_dir = data_dir
+        self.mode = mode
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode, self.angle)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+        columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
+        pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
+        df  = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
+
+        pos_gt = []
+        for i in range(len(df)):
+            for j in range(df.shape[1]):
+                if j == 0:
+                    key = df.iloc[i, j]
+                    temp_index = []
+                else:
+                    value = df.iloc[i, j]
+                    temp_index.append(value)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+    
+    def get_gt_npy(self,):
+
+        data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
+        data = np.load(data_path, allow_pickle=True)
+        pos_gt = []
+        for i in range(data.shape[0]):
+            key = data[i, 0]
+            temp_index = []
+            temp_value = data[i, 1]
+            for j in temp_value:
+                temp_index.append(j)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode, angle):
+        try:
+            if mode == 'query':
+                img = Image.open(path)
+                if angle == 0:
+                    return img
+                rotated_image = img.rotate(angle,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        # Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class AerialDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 mode,
+                 angle=0,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            temp_path = os.path.join(data_dir,  'query_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = angle
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir,  'reference_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = 0
+
+        self.transforms = transforms
+        self.mode = mode
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+
+
+        self.data_dir = data_dir
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode, self.angle)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+        
+        # group
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+        columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
+        pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
+        df  = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
+
+        pos_gt = []
+        for i in range(len(df)):
+            for j in range(df.shape[1]):
+                if j == 0:
+                    key = df.iloc[i, j]
+                    temp_index = []
+                else:
+                    value = df.iloc[i, j]
+                    temp_index.append(value)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+    
+    def get_gt_npy(self,):
+
+        data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
+        data = np.load(data_path, allow_pickle=True)
+        pos_gt = []
+        for i in range(data.shape[0]):
+            key = data[i, 0]
+            temp_index = []
+            temp_value = data[i, 1]
+            for j in temp_value:
+                temp_index.append(j)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode, angle):
+        try:
+            if mode == 'query': 
+                img = Image.open(path)
+                rotated_image = img.rotate(angle,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+    
+
+class DenseUAVDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                txt, 
+                 mode,
+                 gt_txt,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+           
+        
+        if mode == 'DB':
+             with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+
+        self.transforms = transforms
+        self.mode = mode
+        self.gt_txt = gt_txt
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+        pos_gt = []
+        with open(self.gt_txt, 'r') as f_gt:
+            for info in f_gt:
+                key, values = eval(info.split(' ')[0]), info.strip().split(' ')[1:]
+                temp_value = []
+                for value in values:
+                    temp_value.append(eval(value))
+                pos_gt.append([key, temp_value])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode):
+        try:
+            if mode == 'query':
+                img = Image.open(path)
+                rotated_image = img.rotate(0,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        # Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class DenseUAVDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                txt, 
+                 mode,
+                 gt_txt,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+            
+        
+        if mode == 'DB':
+            with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+
+        self.transforms = transforms
+        self.mode = mode
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+        self.gt_txt = gt_txt
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+        
+        # group
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+        pos_gt = []
+        with open(self.gt_txt, 'r') as f_gt:
+            for info in f_gt:
+                key, values = eval(info.split(' ')[0]), info.strip().split(' ')[1:]
+                temp_value = []
+                for value in values:
+                    temp_value.append(eval(value))
+                pos_gt.append([key, temp_value])
+        
+        return pos_gt
+    
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode):
+        try:
+            if mode == 'query': 
+                img = Image.open(path)
+                rotated_image = img.rotate(270,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+# 测试代码
+
+# data_dir = "/media/guan/新加卷/EdgeBing/WorldLoc"
+# query_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/train_query.txt"
+
+# train_dataset = WorldDatasetTrain(data_dir, query_txt)
+
+# train_dataloader = DataLoader(train_dataset,
+#                                 batch_size=64,
+#                                 num_workers=0,
+#                                 shuffle=False,
+#                                 pin_memory=True)
+
+
+# train_dataloader.dataset.shuffle() 
+
+# for  query, query_pt, reference, reference_pt, idx in  tqdm(train_dataloader, total=len(train_dataloader)):
+
+#     print(1)

GeoLoc-UAV-main/dataset/World_ori.py

@@ -0,0 +1,174 @@
+import os
+import cv2
+import numpy as np
+from torch.utils.data import Dataset
+import copy
+from tqdm import tqdm
+import time
+import random
+
+from torch.utils.data import DataLoader
+import torchvision.transforms as T
+
+from utils.utils import TransformerCV
+
+transform_config = {
+     "sample_scale_begin": 0,
+  "sample_scale_inter": 0.5, 
+  "sample_scale_num": 5, 
+  "sample_rotate_begin": -45,
+  "sample_rotate_inter": 45, 
+  "sample_rotate_num": 8,   
+}
+
+default_transform = T.Compose([
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+def get_data(txt):
+
+    data = {}
+    idx = 0
+    with open(txt, 'r') as f:
+        for line in f:
+            line_list = line.split(' ')[:-1]
+            data[idx] = line_list
+            idx += 1
+
+    return data
+
+class WorldDatasetTrain(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][2])
+            self.pairs.append((label, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+    
+    def __getitem__(self, index):
+        
+        idx, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = cv2.imread(query_img_path)
+        query_img = cv2.cvtColor(query_img, cv2.COLOR_BGR2RGB)
+        # db
+        db_img = cv2.imread(db_img_path)
+        db_img = cv2.cvtColor(db_img, cv2.COLOR_BGR2RGB)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(image=query_img)['image']
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(image=db_img)['image']
+        
+        return query_img, db_img, idx
+    
+    
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+        # print("First Element ID: {} - Last Element ID: {}".format(self.samples[0][0], self.samples[-1][0]))  
+
+
+# 测试代码
+
+# data_dir = "/media/guan/新加卷/EdgeBing/WorldLoc"
+# query_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/train_query.txt"
+
+# train_dataset = WorldDatasetTrain(data_dir, query_txt)
+
+# train_dataloader = DataLoader(train_dataset,
+#                                 batch_size=64,
+#                                 num_workers=0,
+#                                 shuffle=False,
+#                                 pin_memory=True)
+
+
+# train_dataloader.dataset.shuffle() 
+
+# for  query, reference, idx in  tqdm(train_dataloader, total=len(train_dataloader)):
+
+#     print(1)

GeoLoc-UAV-main/dataset/World_rot.py

@@ -0,0 +1,957 @@
+import os
+import cv2
+import numpy as np
+from PIL import Image, UnidentifiedImageError
+from torch.utils.data import Dataset
+import copy
+from tqdm import tqdm
+import time
+import random
+import glob
+import json
+import pandas as pd
+
+from torch.utils.data import DataLoader
+import torchvision.transforms as T
+import json
+
+
+from utils.utils import TransformerCV
+
+# transform_config = {
+#      "sample_scale_begin": 0,
+#   "sample_scale_inter": 0.5, 
+#   "sample_scale_num": 3, 
+#   "sample_rotate_begin": 0,
+#   "sample_rotate_inter": 45, 
+#   "sample_rotate_num": 8,   
+# }
+json_path = "/media/guan/新加卷/Code/Code/configs/transform_config.json"
+with open(json_path, 'r', encoding='utf-8') as file:
+    data = json.load(file)
+transform_config = data["transform_config"]
+
+
+# transform_config = {
+#      "sample_scale_begin": 0,
+#   "sample_scale_inter": 0.5, 
+#   "sample_scale_num": 1, 
+#   "sample_rotate_begin": 0,
+#   "sample_rotate_inter": 0, 
+#   "sample_rotate_num": 1,   
+# }
+
+default_transform = T.Compose([
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+def get_data(txt):
+
+    data = {}
+    idx = 0
+    with open(txt, 'r') as f:
+        for line in f:
+            line_list = line.split(' ')[:-1]
+            data[idx] = line_list
+            idx += 1
+
+    return data
+
+class WorldDatasetTrainGroup(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][2])
+            self.pairs.append((label, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+    
+    def __getitem__(self, index):
+        
+        idx, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = self.image_loader(query_img_path)
+        # db
+        db_img = self.image_loader(db_img_path)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(query_img)
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(db_img)
+        
+        # return query_img, db_img, idx
+        # group
+        query_img *= 255
+        query_img, query_pt = self.transformImg(query_img)
+
+        db_img *= 255
+        db_img, db_pt = self.transformImg(db_img)
+        
+        return query_img, query_pt, db_img, db_pt, idx
+    
+    def transformImg(self, img):
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+        # print("First Element ID: {} - Last Element ID: {}".format(self.samples[0][0], self.samples[-1][0]))  
+
+class WorldDatasetTrainVanilia(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][2])
+            self.pairs.append((label, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+    
+    def __getitem__(self, index):
+        
+        idx, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = self.image_loader(query_img_path)
+        # db
+        db_img = self.image_loader(db_img_path)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(query_img)
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(db_img)
+        
+        return query_img, db_img, idx
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+
+class WorldDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 name,
+                 mode,
+                 height_mode=None,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.transforms = transforms
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+
+        self.data_dir = data_dir
+        self.name = name 
+
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        positive = json.load(open(pos_json_path))
+
+        self.samples = []
+        if mode == 'query':
+            if os.path.exists(os.path.join(data_dir, name,'query',  height_mode, 'footage')):
+                temp_path = os.path.join(data_dir, name,'query', height_mode, 'footage')
+                temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
+                if len(temp) != len(positive.keys()):
+                    filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
+                    self.samples.extend(filter_temp)
+                else:
+                    self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir, name, 'DB', 'img')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+      
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        semi_pos_json_path = os.path.join(self.data_dir, self.name,  'semi_positive.json')
+        positive = json.load(open(pos_json_path))
+        semi_positive = json.load(open(semi_pos_json_path))
+
+        pos_gt = []
+        for key in positive.keys():
+            value = positive[key]
+            
+            temp_index = []
+            # pos
+            for one_value in value:
+                temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                temp_path = temp_path_dir + '/' + one_value
+                one_index = self.samples.index(temp_path)
+                temp_index.append(one_index)
+            # semi-pos
+            try:
+                semi_value = semi_positive[key]
+                for one_value in semi_value:
+                    temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                    temp_path = temp_path_dir + '/' + one_value
+                    one_index = self.samples.index(temp_path)
+                    temp_index.append(one_index)
+            except:
+                pos_gt.append([key, temp_index])
+                continue
+
+            pos_gt.append([key, temp_index])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class WorldDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 name,
+                 mode,
+                 height_mode=None,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.transforms = transforms
+
+        self.data_dir = data_dir
+        self.name = name 
+
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        positive = json.load(open(pos_json_path))
+
+        self.samples = []
+        if mode == 'query':
+           
+            if os.path.exists(os.path.join(data_dir, name,'query',  height_mode, 'footage')): #query
+                temp_path = os.path.join(data_dir, name,'query', height_mode, 'footage')
+                temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
+                if len(temp) != len(positive.keys()):
+                    filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
+                    self.samples.extend(filter_temp)
+                else:
+                    self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir, name, 'DB', 'img')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+      
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        semi_pos_json_path = os.path.join(self.data_dir, self.name,  'semi_positive.json')
+        positive = json.load(open(pos_json_path))
+        semi_positive = json.load(open(semi_pos_json_path))
+
+        pos_gt = []
+        for key in positive.keys():
+            value = positive[key]
+            
+            temp_index = []
+            # pos
+            for one_value in value:
+                temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                temp_path = temp_path_dir + '/' + one_value
+                one_index = self.samples.index(temp_path)
+                temp_index.append(one_index)
+            try:
+                semi_value = semi_positive[key]
+            # semi-pos
+                for one_value in semi_value:
+                    temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                    temp_path = temp_path_dir + '/' + one_value
+                    one_index = self.samples.index(temp_path)
+                    temp_index.append(one_index)
+            except:
+                pos_gt.append([key, temp_index])
+                continue
+
+            pos_gt.append([key, temp_index])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class AerialDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 mode,
+                 angle=0,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            temp_path = os.path.join(data_dir,  'query_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = angle
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir,  'reference_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = angle
+
+        self.transforms = transforms
+        self.data_dir = data_dir
+        self.mode = mode
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode, self.angle)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+        columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
+        pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
+        df  = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
+
+        pos_gt = []
+        for i in range(len(df)):
+            for j in range(df.shape[1]):
+                if j == 0:
+                    key = df.iloc[i, j]
+                    temp_index = []
+                else:
+                    value = df.iloc[i, j]
+                    temp_index.append(value)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+    
+    def get_gt_npy(self,):
+
+        data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
+        data = np.load(data_path, allow_pickle=True)
+        pos_gt = []
+        for i in range(data.shape[0]):
+            key = data[i, 0]
+            temp_index = []
+            temp_value = data[i, 1]
+            for j in temp_value:
+                temp_index.append(j)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode, angle):
+        try:
+            if mode == 'query':
+                img = Image.open(path)
+                if angle == 0:
+                    return img
+                rotated_image = img.rotate(angle,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        # Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class AerialDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 mode,
+                 angle=0,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            temp_path = os.path.join(data_dir,  'query_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = angle
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir,  'reference_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+            self.angle = 0
+
+        self.transforms = transforms
+        self.mode = mode
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+
+
+        self.data_dir = data_dir
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode, self.angle)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+        
+        # group
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+        columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
+        pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
+        df  = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
+
+        pos_gt = []
+        for i in range(len(df)):
+            for j in range(df.shape[1]):
+                if j == 0:
+                    key = df.iloc[i, j]
+                    temp_index = []
+                else:
+                    value = df.iloc[i, j]
+                    temp_index.append(value)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+    
+    def get_gt_npy(self,):
+
+        data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
+        data = np.load(data_path, allow_pickle=True)
+        pos_gt = []
+        for i in range(data.shape[0]):
+            key = data[i, 0]
+            temp_index = []
+            temp_value = data[i, 1]
+            for j in temp_value:
+                temp_index.append(j)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode, angle):
+        try:
+            if mode == 'query': 
+                img = Image.open(path)
+                rotated_image = img.rotate(angle,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+    
+
+class DenseUAVDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                txt, 
+                 mode,
+                 gt_txt,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+           
+        
+        if mode == 'DB':
+             with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+
+        self.transforms = transforms
+        self.mode = mode
+        self.gt_txt = gt_txt
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+        pos_gt = []
+        with open(self.gt_txt, 'r') as f_gt:
+            for info in f_gt:
+                key, values = eval(info.split(' ')[0]), info.strip().split(' ')[1:]
+                temp_value = []
+                for value in values:
+                    temp_value.append(eval(value))
+                pos_gt.append([key, temp_value])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode):
+        try:
+            if mode == 'query':
+                img = Image.open(path)
+                rotated_image = img.rotate(0,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        # Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class DenseUAVDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                txt, 
+                 mode,
+                 gt_txt,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+            
+        
+        if mode == 'DB':
+            with open(txt, 'r') as f:
+                for i in f:
+                    self.samples.append(i.strip())
+
+        self.transforms = transforms
+        self.mode = mode
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+        self.gt_txt = gt_txt
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path, self.mode)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+        
+        # group
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+        pos_gt = []
+        with open(self.gt_txt, 'r') as f_gt:
+            for info in f_gt:
+                key, values = eval(info.split(' ')[0]), info.strip().split(' ')[1:]
+                temp_value = []
+                for value in values:
+                    temp_value.append(eval(value))
+                pos_gt.append([key, temp_value])
+        
+        return pos_gt
+    
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path, mode):
+        try:
+            if mode == 'query': 
+                img = Image.open(path)
+                rotated_image = img.rotate(270,expand=True)
+                return rotated_image
+            else:
+                return Image.open(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+# 测试代码
+
+# data_dir = "/media/guan/新加卷/EdgeBing/WorldLoc"
+# query_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/train_query.txt"
+
+# train_dataset = WorldDatasetTrain(data_dir, query_txt)
+
+# train_dataloader = DataLoader(train_dataset,
+#                                 batch_size=64,
+#                                 num_workers=0,
+#                                 shuffle=False,
+#                                 pin_memory=True)
+
+
+# train_dataloader.dataset.shuffle() 
+
+# for  query, query_pt, reference, reference_pt, idx in  tqdm(train_dataloader, total=len(train_dataloader)):
+
+#     print(1)

GeoLoc-UAV-main/dataset/World_weight.py

@@ -0,0 +1,748 @@
+import os
+import cv2
+import numpy as np
+from PIL import Image, UnidentifiedImageError
+from torch.utils.data import Dataset
+import copy
+from tqdm import tqdm
+import time
+import random
+import glob
+import json
+import pandas as pd
+
+from torch.utils.data import DataLoader
+import torchvision.transforms as T
+
+
+from utils.utils import TransformerCV
+
+transform_config = {
+     "sample_scale_begin": 0,
+  "sample_scale_inter": 0.5, 
+  "sample_scale_num": 3, 
+  "sample_rotate_begin": 0,
+  "sample_rotate_inter": 45, 
+  "sample_rotate_num": 8,   
+}
+
+default_transform = T.Compose([
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+def get_data(txt):
+
+    data = {}
+    idx = 0
+    with open(txt, 'r') as f:
+        for line in f:
+            # line_list = line.split(' ')[:-1]
+            line_list = line.split(' ')
+            data[idx] = line_list
+            idx += 1
+
+    return data
+
+class WorldDatasetTrainGroup(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][3])
+            weight = eval(idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][4][:-1])
+            self.pairs.append((label, weight, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+    
+    def __getitem__(self, index):
+        
+        idx, weight, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = self.image_loader(query_img_path)
+        # db
+        db_img = self.image_loader(db_img_path)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(query_img)
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(db_img)
+        
+        # return query_img, db_img, idx
+        # group
+        query_img *= 255
+        query_img, query_pt = self.transformImg(query_img)
+
+        db_img *= 255
+        db_img, db_pt = self.transformImg(db_img)
+        
+        return query_img, query_pt, db_img, db_pt, idx, weight
+    
+    def transformImg(self, img):
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+        # print("First Element ID: {} - Last Element ID: {}".format(self.samples[0][0], self.samples[-1][0]))  
+
+class WorldDatasetTrainVanilia(Dataset):
+    def __init__(self, 
+                 data_dir, 
+                 query_txt,
+                 transforms_query=default_transform,
+                 transforms_db=default_transform,
+                 shuffle_batch_size=64):
+        super().__init__()
+
+        self.pairs = []
+        self.data = get_data(query_txt)
+
+        for idx in self.data.items():
+            query_img_path = os.path.join(data_dir, idx[1][0])
+            label = eval(idx[1][1])
+            db_image_path = os.path.join(data_dir, idx[1][2])
+            self.pairs.append((label, query_img_path, db_image_path))
+        
+        self.transforms_query = transforms_query
+        self.transforms_db = transforms_db
+        self.shuffle_batch_size = shuffle_batch_size
+        
+        self.samples = copy.deepcopy(self.pairs)
+
+    
+    def __getitem__(self, index):
+        
+        idx, query_img_path, db_img_path = self.samples[index]
+        # query
+        query_img = self.image_loader(query_img_path)
+        # db
+        db_img = self.image_loader(db_img_path)
+         # image transforms
+        if self.transforms_query is not None:
+            query_img = self.transforms_query(query_img)
+            
+        if self.transforms_db is not None:
+            db_img = self.transforms_db(db_img)
+        
+        return query_img, db_img, idx
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+    def shuffle(self, ):
+
+        """
+        generate unique class_id
+        """
+        print("\n Shuffle Dataset")
+
+        pair_pool = copy.deepcopy(self.pairs)
+        #shuffle
+        random.shuffle(pair_pool)
+
+        pairs_epoch = set()   
+        label_batch = set()
+
+        current_batch = []
+        batches = []
+
+         # progressbar
+        pbar = tqdm()
+
+        while True:
+            pbar.update()
+            if len(pair_pool) > 0:
+                pair = pair_pool.pop(0)
+
+                label, _, _ = pair
+
+                if label not in label_batch and pair not in pairs_epoch:
+
+                    label_batch.add(label)
+                    current_batch.append(pair)
+                    pairs_epoch.add(pair)
+
+                    break_counter = 0
+                
+                else:
+                    if pair not in pairs_epoch:
+                        pair_pool.append(pair)
+                    
+                    break_counter += 1
+                
+                if break_counter >= 5000:
+                        break
+            
+            else:
+                break
+
+            if len(current_batch) >= self.shuffle_batch_size:
+                batches.extend(current_batch)
+                label_batch = set()
+                current_batch = []
+        
+        pbar.close()
+
+        time.sleep(0.3)
+
+        self.samples = batches
+
+        print("Original Length: {} - Length after Shuffle: {}".format(len(self.pairs), len(self.samples))) 
+        print("Break Counter:", break_counter)
+        print("Pairs left out of last batch to avoid creating noise:", len(self.pairs) - len(self.samples))
+
+class WorldDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 name,
+                 mode,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.transforms = transforms
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+
+        self.data_dir = data_dir
+        self.name = name 
+
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        positive = json.load(open(pos_json_path))
+
+        self.samples = []
+        if mode == 'query':
+            height_list = ["height100_rot0", "height100_rot45", "height100_rot90", "height100_rot135", "height100_rot180", "height100_rot225", "height100_rot270", "height100_rot315", 
+                           "height125_rot0", "height125_rot45", "height125_rot90", "height125_rot135", "height125_rot180", "height125_rot225", "height125_rot270", "height125_rot315",
+                            "height150_rot0", "height150_rot45", "height150_rot90", "height150_rot135", "height150_rot180", "height150_rot225", "height150_rot270", "height150_rot315"]
+            for i in height_list:
+                if os.path.exists(os.path.join(data_dir, name,'query',  i, 'footage')):
+                    temp_path = os.path.join(data_dir, name,'query', i, 'footage')
+                    temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
+                    if len(temp) != len(positive.keys()):
+                        filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
+                        self.samples.extend(filter_temp)
+                    else:
+                        self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir, name, 'DB', 'img')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+      
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        semi_pos_json_path = os.path.join(self.data_dir, self.name,  'semi_positive.json')
+        positive = json.load(open(pos_json_path))
+        semi_positive = json.load(open(semi_pos_json_path))
+
+        pos_gt = []
+        for key in positive.keys():
+            value = positive[key]
+            
+            temp_index = []
+            # pos
+            for one_value in value:
+                temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                temp_path = temp_path_dir + '/' + one_value
+                one_index = self.samples.index(temp_path)
+                temp_index.append(one_index)
+            # semi-pos
+            try:
+                semi_value = semi_positive[key]
+                for one_value in semi_value:
+                    temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                    temp_path = temp_path_dir + '/' + one_value
+                    one_index = self.samples.index(temp_path)
+                    temp_index.append(one_index)
+            except:
+                pos_gt.append([key, temp_index])
+                continue
+
+            pos_gt.append([key, temp_index])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class WorldDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 name,
+                 mode,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.transforms = transforms
+
+        self.data_dir = data_dir
+        self.name = name 
+
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        positive = json.load(open(pos_json_path))
+
+        self.samples = []
+        if mode == 'query':
+            height_list = ["height100_rot0", "height100_rot45", "height100_rot90", "height100_rot135", "height100_rot180", "height100_rot225", "height100_rot270", "height100_rot315", 
+                           "height125_rot0", "height125_rot45", "height125_rot90", "height125_rot135", "height125_rot180", "height125_rot225", "height125_rot270", "height125_rot315",
+                            "height150_rot0", "height150_rot45", "height150_rot90", "height150_rot135", "height150_rot180", "height150_rot225", "height150_rot270", "height150_rot315"]
+            for i in height_list:
+                if os.path.exists(os.path.join(data_dir, name,'query',  i, 'footage')):
+                    temp_path = os.path.join(data_dir, name,'query', i, 'footage')
+                    temp = sorted(glob.glob(f'{temp_path}/{"*.jpeg"}'))
+                    if len(temp) != len(positive.keys()):
+                        filter_temp = [image for image in temp if image.split('/')[-1].split('.')[0].split('_')[-1] in positive.keys()]
+                        self.samples.extend(filter_temp)
+                    else:
+                        self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir, name, 'DB', 'img')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+      
+        pos_json_path = os.path.join(self.data_dir, self.name,  'positive.json')
+        semi_pos_json_path = os.path.join(self.data_dir, self.name,  'semi_positive.json')
+        positive = json.load(open(pos_json_path))
+        semi_positive = json.load(open(semi_pos_json_path))
+
+        pos_gt = []
+        for key in positive.keys():
+            value = positive[key]
+            
+            temp_index = []
+            # pos
+            for one_value in value:
+                temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                temp_path = temp_path_dir + '/' + one_value
+                one_index = self.samples.index(temp_path)
+                temp_index.append(one_index)
+            try:
+                semi_value = semi_positive[key]
+            # semi-pos
+                for one_value in semi_value:
+                    temp_path_dir = os.path.join(self.data_dir, self.name, 'DB', 'img')
+                    temp_path = temp_path_dir + '/' + one_value
+                    one_index = self.samples.index(temp_path)
+                    temp_index.append(one_index)
+            except:
+                pos_gt.append([key, temp_index])
+                continue
+
+            pos_gt.append([key, temp_index])
+        
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class AerialDatasetEvalVanilia(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 mode,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            temp_path = os.path.join(data_dir,  'query_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir,  'reference_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        self.transforms = transforms
+        self.data_dir = data_dir
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img
+        
+    
+    def get_gt(self,):
+
+        columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
+        pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
+        df  = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
+
+        pos_gt = []
+        for i in range(len(df)):
+            for j in range(df.shape[1]):
+                if j == 0:
+                    key = df.iloc[i, j]
+                    temp_index = []
+                else:
+                    value = df.iloc[i, j]
+                    temp_index.append(value)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+    
+    def get_gt_npy(self,):
+
+        data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
+        data = np.load(data_path, allow_pickle=True)
+        pos_gt = []
+        for i in range(data.shape[0]):
+            key = data[i, 0]
+            temp_index = []
+            temp_value = data[i, 1]
+            for j in temp_value:
+                temp_index.append(j)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            img = Image.open(path)
+            rotated_image = img.rotate(270)
+            return rotated_image
+        # Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+
+class AerialDatasetEvalGroup(Dataset):
+    def __init__(self, 
+                data_dir, 
+                 mode,
+                 transforms=default_transform
+                 ):
+        super().__init__()
+
+        self.samples = []
+        if mode == 'query':
+            temp_path = os.path.join(data_dir,  'query_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+        
+        if mode == 'DB':
+            temp_path = os.path.join(data_dir,  'reference_images')
+            temp = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            self.samples.extend(temp)
+
+        self.transforms = transforms
+      
+        self.group_transformer = TransformerCV(transform_config)
+        self.pts_step = 5
+
+        self.data_dir = data_dir
+
+
+    def __getitem__(self, index):
+        
+        img_path = self.samples[index]
+        # query
+        img = self.image_loader(img_path)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+        
+        # group
+        img *= 255
+        img, pt = self.transformImg(img)
+        
+        return img, pt
+    
+    def transformImg(self, img):
+
+        xs, ys = np.meshgrid(np.arange(self.pts_step,img.size()[1]-self.pts_step,self.pts_step), np.arange(self.pts_step,img.size()[2]-self.pts_step,self.pts_step))
+        xs=xs.reshape(-1,1)
+        ys = ys.reshape(-1,1)
+        pts = np.hstack((xs,ys))
+        img = img.permute(1,2,0).detach().numpy()
+        transformed_imgs=self.group_transformer.transform(img,pts)
+        data_img, data_pt = self.group_transformer.postprocess_transformed_imgs(transformed_imgs)
+        return data_img, data_pt
+    
+    def get_gt(self,):
+
+        columns_to_use_by_index = [1, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,]
+        pos_cvs_path = os.path.join(self.data_dir, 'gt_matches.csv')
+        df  = pd.read_csv(pos_cvs_path, usecols=columns_to_use_by_index)
+
+        pos_gt = []
+        for i in range(len(df)):
+            for j in range(df.shape[1]):
+                if j == 0:
+                    key = df.iloc[i, j]
+                    temp_index = []
+                else:
+                    value = df.iloc[i, j]
+                    temp_index.append(value)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+    
+    def get_gt_npy(self,):
+
+        data_path = os.path.join(self.data_dir, 'vpair_gt.npy')
+        data = np.load(data_path, allow_pickle=True)
+        pos_gt = []
+        for i in range(data.shape[0]):
+            key = data[i, 0]
+            temp_index = []
+            temp_value = data[i, 1]
+            for j in temp_value:
+                temp_index.append(j)
+        
+            pos_gt.append([key, temp_index])
+        return pos_gt
+
+
+    def getitem(self, index):
+
+        return self.samples[index]
+
+
+    
+    @staticmethod
+    def image_loader(path):
+        try:
+            return Image.open(path)
+            # return imread(path)
+        except UnidentifiedImageError:
+            print(f'Image {path} could not be loaded')
+            return Image.new('RGB', (224, 224))
+
+    def __len__(self):
+
+        return len(self.samples)
+    
+# 测试代码
+
+# data_dir = "/media/guan/新加卷/EdgeBing/WorldLoc"
+# query_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/train_query.txt"
+
+# train_dataset = WorldDatasetTrain(data_dir, query_txt)
+
+# train_dataloader = DataLoader(train_dataset,
+#                                 batch_size=64,
+#                                 num_workers=0,
+#                                 shuffle=False,
+#                                 pin_memory=True)
+
+
+# train_dataloader.dataset.shuffle() 
+
+# for  query, query_pt, reference, reference_pt, idx in  tqdm(train_dataloader, total=len(train_dataloader)):
+
+#     print(1)

GeoLoc-UAV-main/eval/eval.py

@@ -0,0 +1,379 @@
+import time
+import torch
+import numpy as np
+from tqdm import tqdm
+from torch.cuda.amp import autocast
+import torch.nn.functional as F
+import faiss
+import faiss.contrib.torch_utils
+import h5py
+import os
+
+def predict(train_config, model, dataloader):
+    
+    model.eval()
+    
+    # wait before starting progress bar
+    # time.sleep(0.1)
+    bar = tqdm(dataloader, total=len(dataloader))
+    # if train_config.verbose:
+    #     bar = tqdm(dataloader, total=len(dataloader))
+    # else:
+    #     bar = dataloader
+        
+    img_features_list = []
+    # import time
+    # torch.cuda.synchronize()
+    # st = time.time()
+    with torch.no_grad():
+        
+        for img, pt in bar:
+        
+            with autocast():
+         
+                img_feature, _ = model(img, pt)
+                # print(f"Initial memory allocated: {torch.cuda.memory_allocated()} bytes")
+            
+            # save features in fp32 for sim calculation
+            img_features_list.append(img_feature.to(torch.float32))
+        # torch.cuda.synchronize()
+        # et = time.time()
+        print('---------------------------------time---------------------------------')
+        # print('time cost: ', (et - st)/len(dataloader))
+        # keep Features on GPU
+        img_features = torch.cat(img_features_list, dim=0) 
+        
+    # if train_config.verbose:
+    bar.close()
+        
+    return img_features
+
+def predict_rerank(train_config, model, dataloader, name, mode):
+    
+    model.eval()
+    
+    # wait before starting progress bar
+    time.sleep(0.1)
+    
+    if train_config.verbose:
+        bar = tqdm(dataloader, total=len(dataloader))
+    else:
+        bar = dataloader
+        
+    img_features_list = []
+    h5_name = str(name)+'_'+mode + '.h5'
+    
+    
+    with torch.no_grad():
+        
+        for img, pt, img_path in bar:
+        
+            with autocast():
+         
+                img_feature, geat_list = model(img, pt)
+                 # save features in fp32 for sim calculation
+                img_features_list.append(img_feature.to(torch.float32))
+                # average_geats = torch.mean(geat_list, dim=2)
+                # average_geats = average_geats.reshape(geat_list.shape[1], geat_list.shape[3], geat_list.shape[4]).cpu()
+                feature_geats = geat_list.squeeze(0).cpu()
+                feature_geats = feature_geats[::60, :, :, :].reshape(-1, 24)
+                
+                # if os.path.exists(h5_name):
+                #     pass
+                with h5py.File(h5_name, 'a', libver='latest') as fd:
+                    if img_path[0] in fd:
+                        continue
+                    grp = fd.create_group(img_path[0])
+                    grp.create_dataset('global_feature', data=feature_geats.cpu())
+        
+        # keep Features on GPU
+        img_features = torch.cat(img_features_list, dim=0) 
+        print('---------------------------------save h5 file---------------------------------')
+        
+    if train_config.verbose:
+        bar.close()
+        
+    return img_features
+
+def predict_backbone(train_config, model, dataloader, LPN):
+    
+    model.eval()
+    
+    # wait before starting progress bar
+    # time.sleep(0.1)
+    bar = tqdm(dataloader, total=len(dataloader))
+    
+    # if train_config.verbose:
+    #     bar = tqdm(dataloader, total=len(dataloader))
+    # else:
+    #     bar = dataloader
+        
+    img_features_list = []
+    # import time
+    # torch.cuda.synchronize()
+    # st = time.time()
+    with torch.no_grad():
+        
+        for img in bar:
+        
+            with autocast():
+         
+                # img_feature = model(img)
+                # img_feature = model(img.to(train_config.device).half())
+                img_feature = model(img.to(train_config["device"]).half())
+                # img_feature = model(img.to(train_config["device"]))
+            
+            # save features in fp32 for sim calculation
+            if LPN:
+                img_feature_tensor = torch.stack(img_feature, dim=2).reshape(img_feature[0].shape[0], -1)
+                img_features_list.append(img_feature_tensor.to(torch.float32))
+            else:
+                
+                img_features_list.append(img_feature.to(torch.float32))
+            
+            
+        # torch.cuda.synchronize()
+        # et = time.time()
+        # print('---------------------------------time---------------------------------')
+        # print('time cost: ', (et - st)/len(dataloader))
+        # keep Features on GPU
+        img_features = torch.cat(img_features_list, dim=0) 
+        
+    # if train_config.verbose:
+    bar.close()
+        
+    return img_features
+
+def evaluate_reank(config,
+                  model,
+                  query_loader,
+                  gallery_loader,
+                  pos_gt,
+                  ranks=[1, 5, 10],
+                  name = None,
+                  cleanup=True):
+    # 需要保存下来group中的特征，故重新书写此代码
+    
+    
+    print("Extract Features:")
+    img_features_query  = predict_rerank(config, model, query_loader, name, 'query')
+    img_features_gallery  = predict_rerank(config, model, gallery_loader, name, 'gallery')
+
+    
+    gl = img_features_gallery.cpu()
+    ql = img_features_query.cpu()
+
+    # -------------------------init------------------------------------------
+    faiss_index = faiss.IndexFlatL2(gl.shape[1])
+    # add references
+    faiss_index.add(gl)
+
+        # search for queries in the index
+    _, predictions = faiss_index.search(ql, max(ranks))
+
+    
+
+    correct_at_rank = np.zeros(len(ranks))
+
+    multi_num = ql.shape[0] / len(pos_gt)
+    really_pos_gt = pos_gt * int(multi_num)
+
+    for q_idx, pred in enumerate(predictions):
+        for i, n in enumerate(ranks):
+            if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1])):
+                correct_at_rank[i] += 1
+
+    correct_at_rank = correct_at_rank / len(predictions)
+   
+    
+    return correct_at_rank, predictions,really_pos_gt
+
+
+def evaluate(config,
+                  model,
+                  query_loader,
+                  gallery_loader,
+                  pos_gt,
+                  mode,
+                  LPN,
+                  ranks=[1, 5, 10],
+                  name = None,
+                  cleanup=True):
+    
+    
+    print("Extract Features:")
+    
+    if mode == 'group':
+        
+        img_features_query = predict(config, model, query_loader)
+        img_features_gallery = predict(config, model, gallery_loader)
+    elif mode == 'vanilia':
+        
+        img_features_query = predict_backbone(config, model, query_loader, LPN)
+        img_features_gallery = predict_backbone(config, model, gallery_loader, LPN)
+    
+    gl = img_features_gallery.cpu()
+    ql = img_features_query.cpu()
+    # t-sne
+    # import numpy as np
+    # from sklearn.manifold import TSNE
+    # from sklearn.preprocessing import StandardScaler
+    # import matplotlib.pyplot as plt
+    
+    # ql_stand = StandardScaler().fit_transform(ql)
+    # num = int(ql_stand.shape[0] / 76)
+    # t_sne_save = config.dataset_root_dir + '/' + name + '/'
+    # y =  list(range(0,10))
+    # reap_y = np.array([item for item in y for _ in range(num)])
+
+    # f_1 = ql_stand[::76, :]
+    # f_2 = ql_stand[5::76, :]
+    # f_3 = ql_stand[10::76, :]
+    # f_4 = ql_stand[15::76, :]
+    # f_5 = ql_stand[20::76, :]
+    # f_6 = ql_stand[25::76, :]
+    # f_7 = ql_stand[30::76, :]
+    # f_8 = ql_stand[35::76, :]
+    # f_9 = ql_stand[40::76, :]
+    # f_10 = ql_stand[45::76, :]
+
+    # x_stand = np.concatenate((f_1, f_2, f_3, f_4, f_5, f_6, f_7, f_8, f_9,f_10), axis=0)
+    
+    # tsne = TSNE(n_components=2, perplexity=num-1, n_iter=5000, n_jobs=-1)
+    # X_tsne = tsne.fit_transform(x_stand)
+    # plt.figure(figsize=(8, 8))
+    # # 归一化颜色值
+    # norm = plt.Normalize(reap_y.min(), reap_y.max())
+    # # 选择不同的颜色映射
+    # cmap = plt.get_cmap('plasma')
+
+    # # 转换颜色值到[0, 1]区间内
+    # colors = cmap(norm(reap_y))
+    # scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1],  c=colors, alpha=0.7)
+    # plt.colorbar(scatter)
+
+    # plt.savefig(t_sne_save + 't_sne_' + 'dinov2'+ '.png')
+
+
+    # -------------------------init------------------------------------------
+    faiss_index = faiss.IndexFlatL2(gl.shape[1])
+    # add references
+    faiss_index.add(gl)
+
+        # search for queries in the index
+    _, predictions = faiss_index.search(ql, max(ranks))
+
+    
+
+    correct_at_rank = np.zeros(len(ranks))
+
+    multi_num = ql.shape[0] / len(pos_gt)
+    really_pos_gt = pos_gt * int(multi_num)
+
+    for q_idx, pred in enumerate(predictions):
+        for i, n in enumerate(ranks):
+            # if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1][:ranks[i]])): # test_40
+            if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1])):
+                correct_at_rank[i] += 1
+          
+    # 测试是问题，设置一个train小样本，快速迭代
+              
+
+    correct_at_rank = correct_at_rank / len(predictions)
+   
+    
+    return correct_at_rank, predictions,really_pos_gt
+
+def evaluate_other(config,
+                  model,
+                  query_loader,
+                  gallery_loader,
+                  pos_gt,
+                  ranks=[1, 5, 10],
+                  name = None,
+                  cleanup=True,
+                  LPN=False):
+    
+    
+    print("Extract Features:")
+    # img_features_query = predict(config, model, query_loader)
+    # img_features_gallery = predict(config, model, gallery_loader)
+    img_features_query = predict_backbone(config, model, query_loader)
+    img_features_gallery = predict_backbone(config, model, gallery_loader)
+    
+    gl = img_features_gallery.cpu()
+    ql = img_features_query.cpu()
+
+    # t-sne
+    # import numpy as np
+    # from sklearn.manifold import TSNE
+    # from sklearn.preprocessing import StandardScaler
+    # import matplotlib.pyplot as plt
+    
+    # ql_stand = StandardScaler().fit_transform(ql)
+    # num = int(ql_stand.shape[0] / 76)
+    # t_sne_save = config.dataset_root_dir + '/' + name + '/'
+    # y =  list(range(0,10))
+    # reap_y = np.array([item for item in y for _ in range(num)])
+
+    # f_1 = ql_stand[::76, :]
+    # f_2 = ql_stand[5::76, :]
+    # f_3 = ql_stand[10::76, :]
+    # f_4 = ql_stand[15::76, :]
+    # f_5 = ql_stand[20::76, :]
+    # f_6 = ql_stand[25::76, :]
+    # f_7 = ql_stand[30::76, :]
+    # f_8 = ql_stand[35::76, :]
+    # f_9 = ql_stand[40::76, :]
+    # f_10 = ql_stand[45::76, :]
+
+    # x_stand = np.concatenate((f_1, f_2, f_3, f_4, f_5, f_6, f_7, f_8, f_9,f_10), axis=0)
+    
+    # tsne = TSNE(n_components=2, perplexity=num-1, n_iter=5000, n_jobs=-1)
+    # X_tsne = tsne.fit_transform(x_stand)
+    # plt.figure(figsize=(8, 8))
+    # # 归一化颜色值
+    # norm = plt.Normalize(reap_y.min(), reap_y.max())
+    # # 选择不同的颜色映射
+    # cmap = plt.get_cmap('plasma')
+
+    # # 转换颜色值到[0, 1]区间内
+    # colors = cmap(norm(reap_y))
+    # scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1],  c=colors, alpha=0.7)
+    # plt.colorbar(scatter)
+
+    # plt.savefig(t_sne_save + 't_sne_' + 'dinov2'+ '.png')
+
+
+    # -------------------------init------------------------------------------
+    faiss_index = faiss.IndexFlatL2(gl.shape[1])
+    # add references
+    faiss_index.add(gl)
+
+        # search for queries in the index
+    _, predictions = faiss_index.search(ql, max(ranks))
+
+    
+
+    correct_at_rank = np.zeros(len(ranks))
+
+  
+    really_pos_gt = pos_gt 
+
+    for q_idx, pred in enumerate(predictions):
+        for i, n in enumerate(ranks):
+            if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1])):
+                correct_at_rank[i] += 1
+
+              
+
+    correct_at_rank = correct_at_rank / len(predictions)
+   
+    
+    return correct_at_rank, predictions,really_pos_gt
+
+   
+
+    
+
+    

GeoLoc-UAV-main/eval_anyloc.py

@@ -0,0 +1,191 @@
+import os
+import sys
+import torch
+import argparse
+import torch
+from eval import eval
+from torchvision import transforms as T
+import numpy as np
+import glob
+from torch.utils.data import DataLoader
+from dataset.World import  DenseUAVDatasetEvalVanilia
+from dataset.World import AerialDatasetEvalVanilia
+from models.anyloc import AnyModel
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='dinov2_vitg14', help='Model architecture')
+    parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
+   
+    # Dataset Paths
+    parser.add_argument('--dataset_query', type=str, default='/media/guan/新加卷/DenseUAV/DenseUAV/test/query.txt', help='Root directory of the dataset')
+    parser.add_argument('--dataset_db', type=str, default='/media/guan/新加卷/DenseUAV/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--dataset_gt', type=str, default='/media/guan/新加卷/DenseUAV/DenseUAV/test/gt.txt', help='Root directory of the dataset')
+    parser.add_argument('--dataset_root_dir', type=str, default='/media/guan/新加卷/EdgeBing/TestData/test_40_midref_rot90/', help='Root directory of the dataset')
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default="/media/guan/新加卷/Code(1)/Code/vit_base_eva_gta_same_area.pth", help='Path to start from a checkpoint')
+    parser.add_argument('--save_dir_path', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(0,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+    config = {
+        "mode": args.mode,
+        "model": args.model,
+        # "dataset_query": args.dataset_query,
+        # "dataset_db": args.dataset_db,
+        # "dataset_gt": args.dataset_gt,
+        "dataset_root_dir":args.dataset_root_dir,
+        "checkpoint_path": args.checkpoint_path,
+        "save_dir_path":args.save_dir_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end,
+        "LPN":False
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = AnyModel(model_name=config['mode'],
+                    pretrained=True)
+ 
+model = model.to(config["device"])
+ 
+# eva_dataset_query = DenseUAVDatasetEvalVanilia(txt=config['dataset_query'],
+#                     mode='query',
+#                     gt_txt=config["dataset_gt"],
+#                     transforms=eval_transform)
+
+# eval_dataloader_query = DataLoader(eva_dataset_query,
+#                     batch_size=config['batch_size'],
+#                     num_workers=config['num_workers'],
+#                     shuffle=not config['custom_sampling'],
+#                     pin_memory=True)
+
+# eva_dataset_db = DenseUAVDatasetEvalVanilia(txt=config['dataset_db'],
+#                         mode='DB',
+#                         gt_txt=config["dataset_gt"],
+#                         transforms=eval_transform)
+
+# eval_dataloader_db = DataLoader(eva_dataset_db,
+#                     batch_size=config['batch_size'],
+#                     num_workers=config['num_workers'],
+#                     shuffle=not config['custom_sampling'],
+#                     pin_memory=True)
+
+
+
+# pos_gt = eval_dataloader_db.dataset.get_gt()
+
+if not os.path.exists(config['save_dir_path']):
+    os.mkdir(config['save_dir_path'])
+
+# test angle
+angle_list = list(range(0, 1))
+for angle in angle_list:
+    eva_dataset_query = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'],
+                            mode='query',
+                            angle=angle,
+                            transforms=eval_transform)
+
+    eval_dataloader_query = DataLoader(eva_dataset_query,
+                            batch_size=config['batch_size'],
+                            num_workers=config['num_workers'],
+                            shuffle=not config['custom_sampling'],
+                            pin_memory=True)
+
+    eva_dataset_db = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'],
+                                mode='DB',
+                                transforms=eval_transform)
+
+    eval_dataloader_db = DataLoader(eva_dataset_db,
+                            batch_size=config['batch_size'],
+                            num_workers=config['num_workers'],
+                            shuffle=not config['custom_sampling'],
+                            pin_memory=True)
+    pos_gt = eval_dataloader_db.dataset.get_gt()    
+    result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["model"], LPN=config['LPN'])
+    print(config['checkpoint_path'])
+    print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+    save_result_txt = config['save_dir_path'] + '/' + str(angle) + '.txt'
+    with open(save_result_txt, 'w') as f_w:
+        info = 'top 1: '+ str(round(result[0]*100,2)) + ' top 5: ' +str(round(result[1]*100,2)) + ' top 10: ' + str(round(result[2]*100,2))
+        f_w.write(info + '\n')
+    f_w.close()
+
+
+# with open("/media/guan/新加卷/Code/result/anyloc/denseuav_g.txt", "w") as f_w:
+#     for i in range(predictions.shape[0]):
+#         query_path = eval_dataloader_query.dataset.getitem(i)
+#         if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+#             num = 1
+#         else:
+#             num = 0
+#         pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+#         info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+#         f_w.write(info)
+

GeoLoc-UAV-main/eval_denseuav.py

@@ -0,0 +1,286 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  DenseUAVDatasetEvalVanilia,DenseUAVDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='group', help='Model architecture')
+    parser.add_argument('--model_path', type=str, default='./world', help='Path to save model checkpoints')
+    
+    # Group Config
+    parser.add_argument('--group_arch', type=str, default='groupdinonet', help='Group architecture')
+    parser.add_argument('--group_config', type=str, default='none', help='Group configuration')
+
+    # Backbone Config
+    parser.add_argument('--backbone_arch', type=str, default='dinov2_vits14', help='Backbone architecture')
+    parser.add_argument('--pretrain_flag', type=bool, default=True, help='Flag to use pre-trained weights')
+
+    # Agg Config
+    parser.add_argument('--agg_arch', type=str, default='multiconvap', help='Aggregation architecture')
+    parser.add_argument('--agg_in_channels', type=int, default=384, help='Input channels for aggregation')
+    parser.add_argument('--agg_out_channels', type=int, default=384, help='Output channels for aggregation')
+    parser.add_argument('--agg_s1', type=int, default=1, help='Aggregation s1 parameter')
+    parser.add_argument('--agg_s2', type=int, default=1, help='Aggregation s2 parameter')
+    parser.add_argument('--agg_LPN', type=bool, default=False, help='Use LPN for aggregation')
+
+    # Dataset Paths
+    parser.add_argument('--dataset_query', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/query.txt', help='Root directory of the dataset')
+    parser.add_argument('--dataset_db', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--dataset_gt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/gt.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=5, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(0,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    # Build the config dictionaries dynamically based on parsed args
+    group_config = {
+        "group_arch": args.group_arch,
+        "group_config": {args.group_config}
+    }
+
+    backbone_config = {
+        "backbone_arch": args.backbone_arch,
+        "pretrain_flag": args.pretrain_flag
+    }
+
+    agg_config = {
+        "agg_arch": args.agg_arch,
+        "agg_config": {
+            "in_channels": args.agg_in_channels,
+            "out_channels": args.agg_out_channels,
+            "s1": args.agg_s1,
+            "s2": args.agg_s2,
+            "LPN": args.agg_LPN
+        }
+    }
+
+    config = {
+        "mode": args.mode,
+        "model_path": args.model_path,
+        "group": group_config,
+        "backbone": backbone_config,
+        "agg": agg_config,
+        "dataset_query": args.dataset_query,
+        "dataset_db": args.dataset_db,
+        "dataset_gt": args.dataset_gt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+if config["mode"] == "vanilia":
+    model = model.BackboneGlobal(config['backbone']['backbone_arch'],
+                             config['backbone']['pretrain_flag'],
+                               config['agg']['agg_arch'],
+                               config['agg']['agg_config'])
+    
+    eva_dataset_query = DenseUAVDatasetEvalVanilia(txt=config['dataset_query'],
+                        mode='query',
+                        gt_txt=config["dataset_gt"],
+                        transforms=eval_transform)
+
+    eval_dataloader_query = DataLoader(eva_dataset_query,
+                        batch_size=config['batch_size'],
+                        num_workers=config['num_workers'],
+                        shuffle=not config['custom_sampling'],
+                        pin_memory=True)
+
+    eva_dataset_db = DenseUAVDatasetEvalVanilia(txt=config['dataset_db'],
+                            mode='DB',
+                            gt_txt=config["dataset_gt"],
+                            transforms=eval_transform)
+
+    eval_dataloader_db = DataLoader(eva_dataset_db,
+                        batch_size=config['batch_size'],
+                        num_workers=config['num_workers'],
+                        shuffle=not config['custom_sampling'],
+                        pin_memory=True)
+else:
+    model = model.GrounpDinoGlobal(config['group']['group_arch'],
+                               config['agg']['agg_arch'],
+                               config['agg']['agg_config'])
+    eva_dataset_query = DenseUAVDatasetEvalGroup(txt=config["dataset_query"],
+                        mode='query',
+                        gt_txt=config['dataset_gt'],
+                        transforms=eval_transform)
+
+    eval_dataloader_query = DataLoader(eva_dataset_query,
+                        batch_size=config['batch_size'],
+                        num_workers=config['num_workers'],
+                        shuffle=not config['custom_sampling'],
+                        pin_memory=True)
+
+    eva_dataset_db = DenseUAVDatasetEvalGroup(txt=config["dataset_db"],
+                            mode='DB',
+                            gt_txt=config['dataset_gt'],
+                            transforms=eval_transform)
+
+    eval_dataloader_db = DataLoader(eva_dataset_db,
+                        batch_size=config['batch_size'],
+                        num_workers=config['num_workers'],
+                        shuffle=not config['custom_sampling'],
+                        pin_memory=True)
+
+
+
+model_state_dict = torch.load(config['checkpoint_path'], map_location=config['device'])
+model.load_state_dict(model_state_dict, strict=False)
+
+model = model.to(config['device'])
+
+
+pos_gt = eval_dataloader_db.dataset.get_gt()
+result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["mode"],LPN=config['agg']['agg_config']['LPN'])
+print(config['checkpoint_path'])
+print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+
+
+
+            # vis and save retrieval results
+# save_vis_dir = config.dataset_root_dir + '/'  + 'vis' + '/'
+# if not os.path.exists(save_vis_dir):
+#     os.makedirs(save_vis_dir)
+
+# temp_path = os.path.join(config.dataset_root_dir,  'reference_images')
+# DB_path = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+
+# # save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+#     for i in range(predictions.shape[0]):
+#         query_path = eval_dataloader_query.dataset.getitem(i)
+#         if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+#             num = 1
+#         else:
+#             num = 0
+#         pred_path = DB_path[predictions[i,0]]
+#         info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+#         f.write(info)
+
+
+# for i in range(predictions.shape[0]):
+#     query_path = eval_dataloader_query.dataset.getitem(i)
+#     fig, axs = plt.subplots(2, 6, figsize=(15, 5)) 
+#     query_img = plt.imread(query_path)
+#     for j in range(2):
+#         for k in range(6):
+#             if j == 0 and k == 0:
+#                 axs[j, k].imshow(query_img)
+#                 axs[j, k].axis('off')  # 不显示坐标轴
+#             elif j==0 and k != 0:
+#                 if np.any(np.in1d(predictions[i,k], really_pos_gt[i][1] )):
+
+#                     db_img_path = DB_path[predictions[i,k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                         # 创建一个矩形框
+#                     rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='blue', facecolor='none')
+                    
+#                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+#                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+#                     axs[j, k].add_patch(rect)
+#                     axs[j,k].axis('off')  # 不显示坐标轴
+#                 else:
+#                     db_img_path = DB_path[predictions[i,k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                         # 创建一个矩形框
+#                     rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='red', facecolor='none')
+                    
+#                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+#                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+#                     axs[j, k].add_patch(rect)
+#                     axs[j, k].axis('off')  # 不显示坐标轴
+#             if j ==1:
+#                 try:
+#                     db_img_path = DB_path[really_pos_gt[i][1][k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                     axs[j, k].axis('off')  # 不显示坐标轴
+#                 except:
+#                     break
+            
+    # save_one_path = save_vis_dir + str(i) + '.png'
+    # plt.savefig(save_one_path, dpi=300)
+                       
+                    

GeoLoc-UAV-main/eval_game4loc.py

@@ -0,0 +1,190 @@
+import os
+import sys
+import torch
+import argparse
+import torch
+from eval import eval
+from torchvision import transforms as T
+import numpy as np
+import glob
+from torch.utils.data import DataLoader
+from dataset.World import  DenseUAVDatasetEvalVanilia
+from dataset.World import AerialDatasetEvalVanilia
+from models.game4loc import DesModel
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='vit_base_patch16_rope_reg1_gap_256.sbb_in1k', help='Model architecture')
+    parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
+   
+    # Dataset Paths
+    # parser.add_argument('--dataset_query', type=str, default='/media/guan/新加卷/DenseUAV/DenseUAV/test/query.txt', help='Root directory of the dataset')
+    # parser.add_argument('--dataset_db', type=str, default='/media/guan/新加卷/DenseUAV/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    # parser.add_argument('--dataset_gt', type=str, default='/media/guan/新加卷/DenseUAV/DenseUAV/test/gt.txt', help='Root directory of the dataset')
+    parser.add_argument('--dataset_root_dir', type=str, default='/media/guan/新加卷/EdgeBing/TestData/test_40_midref_rot0', help='Root directory of the dataset')
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default="/media/guan/新加卷/Code/Code/vit_base_eva_gta_same_area.pth", help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=5, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(0,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+    config = {
+        "mode": args.mode,
+        "model": args.model,
+        # "dataset_query": args.dataset_query,
+        # "dataset_db": args.dataset_db,
+        # "dataset_gt": args.dataset_gt,
+        "dataset_root_dir":args.dataset_root_dir,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end,
+        "LPN":False
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((384, 384), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = DesModel(model_name='vit_base_patch16_rope_reg1_gap_256.sbb_in1k',
+                    pretrained=True,
+                    img_size=384,
+                    share_weights=True)
+
+if config["checkpoint_path"] is not None:  
+    print("Start from:", config["checkpoint_path"])
+    model_state_dict = torch.load(config["checkpoint_path"])  
+    model.load_state_dict(model_state_dict, strict=True)     
+
+    # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    # if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+    #     model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+            
+    # Model to device   
+    model = model.to(config["device"])
+
+eva_dataset_query = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'],
+                        mode='query',
+                        transforms=eval_transform)
+
+eval_dataloader_query = DataLoader(eva_dataset_query,
+                        batch_size=config['batch_size'],
+                        num_workers=config['num_workers'],
+                        shuffle=not config['custom_sampling'],
+                        pin_memory=True)
+
+eva_dataset_db = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'],
+                            mode='DB',
+                            transforms=eval_transform)
+
+eval_dataloader_db = DataLoader(eva_dataset_db,
+                        batch_size=config['batch_size'],
+                        num_workers=config['num_workers'],
+                        shuffle=not config['custom_sampling'],
+                        pin_memory=True)
+pos_gt = eval_dataloader_db.dataset.get_gt()    
+# eva_dataset_query = DenseUAVDatasetEvalVanilia(txt=config['dataset_query'],
+#                     mode='query',
+#                     gt_txt=config["dataset_gt"],
+#                     transforms=eval_transform)
+
+# eval_dataloader_query = DataLoader(eva_dataset_query,
+#                     batch_size=config['batch_size'],
+#                     num_workers=config['num_workers'],
+#                     shuffle=not config['custom_sampling'],
+#                     pin_memory=True)
+
+# eva_dataset_db = DenseUAVDatasetEvalVanilia(txt=config['dataset_db'],
+#                         mode='DB',
+#                         gt_txt=config["dataset_gt"],
+#                         transforms=eval_transform)
+
+# eval_dataloader_db = DataLoader(eva_dataset_db,
+#                     batch_size=config['batch_size'],
+#                     num_workers=config['num_workers'],
+#                     shuffle=not config['custom_sampling'],
+#                     pin_memory=True)
+
+
+
+# pos_gt = eval_dataloader_db.dataset.get_gt()
+
+
+result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["model"], LPN=config['LPN'])
+print(config['checkpoint_path'])
+print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+
+with open("/media/guan/新加卷/Code/result/Game4loc/nadir_g.txt", "w") as f_w:
+    for i in range(predictions.shape[0]):
+        query_path = eval_dataloader_query.dataset.getitem(i)
+        if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+            num = 1
+        else:
+            num = 0
+        pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+        info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+        f_w.write(info)

GeoLoc-UAV-main/eval_other_data.py

@@ -0,0 +1,233 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  AerialDatasetEvalGroup, AerialDatasetEvalVanilia
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+
+def default_group_config():
+
+    return {
+          "group_arch" : "groupdinonet", #group
+        "group_config": {
+            "none"
+        }
+    }
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "dinov2_vits14",  #dinov2_vitb14,resnet18
+          "pretrain_flag":True
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "multiconvap", #convap
+    "agg_config": {
+      "in_channels": 384, #256 #512，768
+      "out_channels": 384, #256
+      "s1": 1,
+      "s2": 1,
+      'LPN':False
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "resnet18"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world"
+
+    # model config
+    group:dict = field(default_factory=default_group_config)
+    backbone:dict = field(default_factory=default_backbone_config)
+
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc/TestData/vpair"
+    train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/WorldLoc/Index/train_query.txt"
+
+    # val_index
+    val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/WorldLoc/Index/val.txt"
+
+    # test_index
+    test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/WorldLoc/Index/test.txt"
+    save_pred_txt = "/media/Shen/Data/RingoData/WorldLoc/txt/rot270/divo-s-frozen.txt"
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 if os.name == 'nt' else 4 
+    
+    # train on GPU if available
+    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 30
+    batch_size: int = 10                # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (1,)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = model.BackboneGlobal(config.backbone['backbone_arch'],
+                             config.backbone['pretrain_flag'],
+                               config.agg['agg_arch'],
+                               config.agg['agg_config'])
+
+# model = model.GrounpGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+
+# model = model.GrounpDinoGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+
+model_state_dict = torch.load("/media/Shen/Data/RingoData/WorldLoc/Code/world_vanilia/dinos-info-data-aug-multi-frozen-/122040/weights_e1_0.4058.pth")
+model.load_state_dict(model_state_dict, strict=False)
+
+model = model.to(config.device)
+
+
+eva_dataset_query = AerialDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                        mode='query',
+                        transforms=eval_transform)
+
+eval_dataloader_query = DataLoader(eva_dataset_query,
+                    batch_size=config.batch_size,
+                    num_workers=config.num_workers,
+                    shuffle=not config.custom_sampling,
+                    pin_memory=True)
+
+eva_dataset_db = AerialDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                        mode='DB',
+                        transforms=eval_transform)
+
+eval_dataloader_db = DataLoader(eva_dataset_db,
+                    batch_size=config.batch_size,
+                    num_workers=config.num_workers,
+                    shuffle=not config.custom_sampling,
+                    pin_memory=True)
+
+pos_gt = eval_dataloader_db.dataset.get_gt_npy() #get_gt()#
+result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='vanilia',LPN=config.agg['agg_config']['LPN'])
+print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+
+
+            # vis and save retrieval results
+# save_vis_dir = config.dataset_root_dir + '/'  + 'vis' + '/'
+# if not os.path.exists(save_vis_dir):
+#     os.makedirs(save_vis_dir)
+
+temp_path = os.path.join(config.dataset_root_dir,  'reference_images')
+DB_path = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+
+# save top 1 flase or wrong
+with open(config.save_pred_txt, 'w') as f:
+    for i in range(predictions.shape[0]):
+        query_path = eval_dataloader_query.dataset.getitem(i)
+        if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+            num = 1
+        else:
+            num = 0
+        pred_path = DB_path[predictions[i,0]]
+        info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+        f.write(info)
+
+
+# for i in range(predictions.shape[0]):
+#     query_path = eval_dataloader_query.dataset.getitem(i)
+#     fig, axs = plt.subplots(2, 6, figsize=(15, 5)) 
+#     query_img = plt.imread(query_path)
+#     for j in range(2):
+#         for k in range(6):
+#             if j == 0 and k == 0:
+#                 axs[j, k].imshow(query_img)
+#                 axs[j, k].axis('off')  # 不显示坐标轴
+#             elif j==0 and k != 0:
+#                 if np.any(np.in1d(predictions[i,k], really_pos_gt[i][1] )):
+
+#                     db_img_path = DB_path[predictions[i,k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                         # 创建一个矩形框
+#                     rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='blue', facecolor='none')
+                    
+#                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+#                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+#                     axs[j, k].add_patch(rect)
+#                     axs[j,k].axis('off')  # 不显示坐标轴
+#                 else:
+#                     db_img_path = DB_path[predictions[i,k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                         # 创建一个矩形框
+#                     rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='red', facecolor='none')
+                    
+#                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+#                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+#                     axs[j, k].add_patch(rect)
+#                     axs[j, k].axis('off')  # 不显示坐标轴
+#             if j ==1:
+#                 try:
+#                     db_img_path = DB_path[really_pos_gt[i][1][k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                     axs[j, k].axis('off')  # 不显示坐标轴
+#                 except:
+#                     break
+            
+    # save_one_path = save_vis_dir + str(i) + '.png'
+    # plt.savefig(save_one_path, dpi=300)
+                       
+                    

GeoLoc-UAV-main/eval_real.sh

@@ -0,0 +1,86 @@
+# test-40 use terrain .pth
+# python eval_real_dataset.py --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/resnet18_fine/weights_e10_0.2697.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/res18_finetune
+# python eval_real_dataset.py --mode vanilia --backbone_arch resnet18 --pretrain_flag True --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/resnet18_frozen/weights_e10_0.2752.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/res18_f
+# python eval_real_dataset.py --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/resnet18_lpn/weights_e13_0.2686.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/res18_lpn
+# python eval_real_dataset.py --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/dinos-finetune/weights_e3_0.4102.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/dinos_finetune
+# python eval_real_dataset.py --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag True --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/dinos_frozen/weights_e3_0.3964.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/dinos_f
+# python eval_real_dataset.py --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/dinos_lpn/weights_e1_0.3376.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/dinos_lpn
+# python eval_real_dataset.py --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-terrain-s3r4/132119/weights_e1_0.5073.pth --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/our
+
+
+# DenseUAV 
+# python eval_denseuav.py --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/resnet_fine/weights_e5_0.4895.pth
+# python eval_denseuav.py --mode vanilia --backbone_arch resnet18 --pretrain_flag True --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/resnet_frozen/weights_e5_0.4828.pth
+# python eval_denseuav.py --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/resnet_lpn/weights_e5_0.3936.pth
+# python eval_denseuav.py --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/dinos_fine/weights_e1_0.5286.pth
+# python eval_denseuav.py --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag True --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/dinos_frozen/weights_e4_0.5813.pth
+# python eval_denseuav.py --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/dinos_lpn/weights_e5_0.5437.pth
+# python eval_denseuav.py --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-city-s3r4/120253/weights_e10_0.7412.pth
+
+# our dataset terrain
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/terrain/dino_frozen.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag True --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/dinos_frozen/weights_e3_0.3964.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/terrain/dino_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/dinos-finetune/weights_e3_0.4102.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/terrain/dino_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/dinos_lpn/weights_e1_0.3376.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/terrain/resnet_frozen.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag True --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/resnet18_frozen/weights_e10_0.2752.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/terrain/resnet_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/resnet18_fine/weights_e10_0.2697.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/terrain/resnet_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/terrain/resnet18_lpn/weights_e13_0.2686.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/result/terrain/our.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-terrain-s3r4/132119/weights_e4_0.5587.pth
+
+
+
+# our dataset city
+# python eval_simidataset_parser.py --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag True --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_frozen/weights_e6_0.7023.pth
+# python eval_simidataset_parser.py --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_finetune/weights_e5_0.6370.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/city/dino_frozen.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag True --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_frozen/weights_e6_0.7023.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/city/dino_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_finetune/weights_e5_0.6370.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/city/dino_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_LPN/weights_e5_0.5323.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/city/resnet_frozen.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_country.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag True --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/resnet_frozen/weights_e5_0.6413.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/city/resnet_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/resnet_finetune/weights_e5_0.5685.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/city/resnet_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/resnet_LPN/weights_e5_0.4468.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/result/city/our.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-city-s3r4/120253/weights_e10_0.7412.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/city/group_our_epoch1.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_country.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-country/210021/weights_e1_0.6177.pth
+
+
+# our dataset mix
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/Mix/dino_frozen.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag True --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/dinos_frozen/weights_e4_0.5813.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/Mix/dino_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/dinos_fine/weights_e1_0.5286.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/Mix/dino_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/dinos_lpn/weights_e5_0.5437.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/Mix/resnet_frozen.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag True --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/resnet_frozen/weights_e5_0.4828.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/Mix/resnet_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/resnet_fine/weights_e5_0.4895.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/worldloc_result/Mix/resnet_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/resnet_lpn/weights_e5_0.3936.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/Code/result/mix/our.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/PTH/all/our/weights_e9_0.6299.pth
+
+# ablation
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s2r2.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-all-s2r2/174124/weights_e9_0.6315.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s3r6.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-all-s3r6/201207/weights_e4_0.5876.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s3r4.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-all-s3r4/100424/weights_e6_0.6471.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s3r3.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-all-s3r3/101328/weights_e9_0.5981.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s2r2.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-all-s2r2/174124/weights_e9_0.6315.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s1r1.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-all-s1r1/231526/weights_e3_0.5810.pth
+# python eval_simidataset_parser.py --save_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/txt/worldloc_result/ablation/s3r8.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/Index/test_all.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path world/groupdino-new-all/190104/weights_e9_0.6299.pth
+
+# anyloc small - base - large -giant
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vits14 --save_txt /media/guan/新加卷/Code/result/terrain/dinos_vis.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitb14 --save_txt /media/guan/新加卷/Code/result/terrain/dinob.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitl14 --save_txt /media/guan/新加卷/Code/result/terrain/dinol_vis.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitg14 --save_txt /media/guan/新加卷/Code/result/terrain/dinog_vis.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vits14 --save_txt /media/guan/新加卷/Code/result/city/dinos.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitb14 --save_txt /media/guan/新加卷/Code/result/city/dinob.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitl14 --save_txt /media/guan/新加卷/Code/result/city/dinol.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitg14 --save_txt /media/guan/新加卷/Code/result/city/dinog.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_country.txt
+
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vits14 --save_txt /media/guan/新加卷/Code/result/mix/dinos.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitb14 --save_txt /media/guan/新加卷/Code/result/mix/dinob.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitl14 --save_txt /media/guan/新加卷/Code/result/mix/dinol.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt
+# python eval_simidataset_parser_anyloc.py --mode dinov2_vitg14 --save_txt /media/guan/新加卷/Code/result/mix/dinog.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_all.txt
+
+# python eval_simidataset_parser_game4loc.py --save_txt /media/guan/新加卷/Code/result/Game4loc/Game4loc_vis0.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+# python eval_simidataset_parser_game4loc.py --save_txt /media/guan/新加卷/Code/result/Game4loc/Game4loc_vis.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+# python eval_simidataset_parser_game4loc.py --save_txt /media/guan/新加卷/Code/result/Game4loc/Game4loc_vis.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/Index/test_vis.txt
+
+
+# anyloc
+# python eval_anyloc.py --mode dinov2_vits14 --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/anyloc_small
+# python eval_anyloc.py --mode dinov2_vitl14 --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/anyloc_base
+# python eval_anyloc.py --mode dinov2_vitg14 --save_dir_path /media/guan/新加卷/Code/worldloc_result/Rot/anyloc_large

GeoLoc-UAV-main/eval_real_dataset.py

@@ -0,0 +1,308 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  AerialDatasetEvalGroup, AerialDatasetEvalVanilia
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+from tqdm import tqdm
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='group', help='Model architecture')
+    parser.add_argument('--model_path', type=str, default='./world', help='Path to save model checkpoints')
+    
+    # Group Config
+    parser.add_argument('--group_arch', type=str, default='groupdinonet', help='Group architecture')
+    parser.add_argument('--group_config', type=str, default='none', help='Group configuration')
+
+    # Backbone Config
+    parser.add_argument('--backbone_arch', type=str, default='dinov2_vits14', help='Backbone architecture')
+    parser.add_argument('--pretrain_flag', type=bool, default=True, help='Flag to use pre-trained weights')
+
+    # Agg Config
+    parser.add_argument('--agg_arch', type=str, default='multiconvap', help='Aggregation architecture')
+    parser.add_argument('--agg_in_channels', type=int, default=384, help='Input channels for aggregation')
+    parser.add_argument('--agg_out_channels', type=int, default=384, help='Output channels for aggregation')
+    parser.add_argument('--agg_s1', type=int, default=1, help='Aggregation s1 parameter')
+    parser.add_argument('--agg_s2', type=int, default=1, help='Aggregation s2 parameter')
+    parser.add_argument('--agg_LPN', type=bool, default=False, help='Use LPN for aggregation')
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root_dir', type=str, default='/media/guan/新加卷/EdgeBing/TestData/test_40_midref_rot0/', help='Root directory of the dataset')
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
+    parser.add_argument('--save_dir_path', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    # Build the config dictionaries dynamically based on parsed args
+    group_config = {
+        "group_arch": args.group_arch,
+        "group_config": {args.group_config}
+    }
+
+    backbone_config = {
+        "backbone_arch": args.backbone_arch,
+        "pretrain_flag": args.pretrain_flag
+    }
+
+    agg_config = {
+        "agg_arch": args.agg_arch,
+        "agg_config": {
+            "in_channels": args.agg_in_channels,
+            "out_channels": args.agg_out_channels,
+            "s1": args.agg_s1,
+            "s2": args.agg_s2,
+            "LPN": args.agg_LPN
+        }
+    }
+
+    config = {
+        "mode": args.mode,
+        "model_path": args.model_path,
+        "group": group_config,
+        "backbone": backbone_config,
+        "agg": agg_config,
+        "dataset_root_dir": args.dataset_root_dir,
+        "checkpoint_path": args.checkpoint_path,
+        "save_dir_path":args.save_dir_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Test Config
+#-------------------------------------------------------------------------------------------#
+config = parse_config()
+
+if not os.path.exists(config['save_dir_path']):
+    os.mkdir(config['save_dir_path'])
+
+# test angle
+angle_list = list(range(0, 1))
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+
+if config["mode"] == "vanilia":
+    model = model.BackboneGlobal(config['backbone']['backbone_arch'],
+                            config['backbone']['pretrain_flag'],
+                            config['agg']['agg_arch'],
+                            config['agg']['agg_config'])
+else:
+    model = model.GrounpDinoGlobal(config['group']['group_arch'],
+                            config['agg']['agg_arch'],
+                            config['agg']['agg_config'])
+
+ 
+for angle in tqdm(angle_list):
+    if config["mode"] == "vanilia":
+        eva_dataset_query = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'],
+                            mode='query',
+                            angle=angle,
+                            transforms=eval_transform)
+
+        eval_dataloader_query = DataLoader(eva_dataset_query,
+                            batch_size=config['batch_size'],
+                            num_workers=config['num_workers'],
+                            shuffle=not config['custom_sampling'],
+                            pin_memory=True)
+
+        eva_dataset_db = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'],
+                                mode='DB',
+                                transforms=eval_transform)
+
+        eval_dataloader_db = DataLoader(eva_dataset_db,
+                            batch_size=config['batch_size'],
+                            num_workers=config['num_workers'],
+                            shuffle=not config['custom_sampling'],
+                            pin_memory=True)
+    else:
+        
+        eva_dataset_query = AerialDatasetEvalGroup(data_dir=config["dataset_root_dir"],
+                            mode='query',
+                            angle=angle,
+                            transforms=eval_transform)
+
+        eval_dataloader_query = DataLoader(eva_dataset_query,
+                            batch_size=config['batch_size'],
+                            num_workers=config['num_workers'],
+                            shuffle=not config['custom_sampling'],
+                            pin_memory=True)
+
+        eva_dataset_db = AerialDatasetEvalGroup(data_dir=config["dataset_root_dir"],
+                                mode='DB',
+                                transforms=eval_transform)
+
+        eval_dataloader_db = DataLoader(eva_dataset_db,
+                            batch_size=config['batch_size'],
+                            num_workers=config['num_workers'],
+                            shuffle=not config['custom_sampling'],
+                            pin_memory=True)
+
+
+
+
+    # model = model.GrounpGlobal(config.group['group_arch'],
+    #                                config.agg['agg_arch'],
+    #                                config.agg['agg_config'])
+
+
+
+    model_state_dict = torch.load(config['checkpoint_path'], map_location='cuda:0')
+    model.load_state_dict(model_state_dict, strict=False)
+
+    model = model.to(config['device'])
+
+
+
+    # pos_gt = eval_dataloader_db.dataset.get_gt_npy()
+    pos_gt = eval_dataloader_db.dataset.get_gt()
+    
+    result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["mode"],LPN=config['agg']['agg_config']['LPN'])
+    print(config['checkpoint_path'])
+    print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+    save_result_txt = config['save_dir_path'] + '/' + str(angle) + '.txt'
+    with open(save_result_txt, 'w') as f_w:
+        info = 'top 1: '+ str(round(result[0]*100,2)) + ' top 5: ' +str(round(result[1]*100,2)) + ' top 10: ' + str(round(result[2]*100,2))
+        f_w.write(info + '\n')
+    f_w.close()
+
+
+            # vis and save retrieval results
+# save_vis_dir = config.dataset_root_dir + '/'  + 'vis' + '/'
+# if not os.path.exists(save_vis_dir):
+#     os.makedirs(save_vis_dir)
+
+# temp_path = os.path.join(config.dataset_root_dir,  'reference_images')
+# DB_path = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+
+# # save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+#     for i in range(predictions.shape[0]):
+#         query_path = eval_dataloader_query.dataset.getitem(i)
+#         if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+#             num = 1
+#         else:
+#             num = 0
+#         pred_path = DB_path[predictions[i,0]]
+#         info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+#         f.write(info)
+
+
+# for i in range(predictions.shape[0]):
+#     query_path = eval_dataloader_query.dataset.getitem(i)
+#     fig, axs = plt.subplots(2, 6, figsize=(15, 5)) 
+#     query_img = plt.imread(query_path)
+#     for j in range(2):
+#         for k in range(6):
+#             if j == 0 and k == 0:
+#                 axs[j, k].imshow(query_img)
+#                 axs[j, k].axis('off')  # 不显示坐标轴
+#             elif j==0 and k != 0:
+#                 if np.any(np.in1d(predictions[i,k], really_pos_gt[i][1] )):
+
+#                     db_img_path = DB_path[predictions[i,k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                         # 创建一个矩形框
+#                     rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='blue', facecolor='none')
+                    
+#                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+#                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+#                     axs[j, k].add_patch(rect)
+#                     axs[j,k].axis('off')  # 不显示坐标轴
+#                 else:
+#                     db_img_path = DB_path[predictions[i,k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                         # 创建一个矩形框
+#                     rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='red', facecolor='none')
+                    
+#                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+#                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+#                     axs[j, k].add_patch(rect)
+#                     axs[j, k].axis('off')  # 不显示坐标轴
+#             if j ==1:
+#                 try:
+#                     db_img_path = DB_path[really_pos_gt[i][1][k]]
+#                     db_img = plt.imread(db_img_path)
+#                     axs[j, k].imshow(db_img)
+#                     axs[j, k].axis('off')  # 不显示坐标轴
+#                 except:
+#                     break
+            
+    # save_one_path = save_vis_dir + str(i) + '.png'
+    # plt.savefig(save_one_path, dpi=300)
+                       
+                    

GeoLoc-UAV-main/eval_rot.sh

@@ -0,0 +1,14 @@
+
+# python eval_simidataset_parser_rot.py --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/dino_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_finetune/weights_e5_0.6370.pth
+# python eval_simidataset_parser_rot.py --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/dino_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt --mode vanilia --backbone_arch dinov2_vits14 --pretrain_flag False --agg_in_channels 384 --agg_out_channels 384 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/dinos_LPN/weights_e5_0.5323.pth
+# python eval_simidataset_parser_rot.py --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/resnet_finetune.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN False --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/resnet_finetune/weights_e5_0.5685.pth
+# python eval_simidataset_parser_rot.py --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/resnet_lpn.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt --mode vanilia --backbone_arch resnet18 --pretrain_flag False --agg_in_channels 512 --agg_out_channels 512 --agg_LPN True --checkpoint_path /media/guan/新加卷/Code/Code/PTH/city/resnet_LPN/weights_e5_0.4468.pth
+# python eval_simidataset_parser_rot.py --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/our.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt --mode group --agg_in_channels 256 --agg_out_channels 256 --checkpoint_path /media/guan/新加卷/Code/Code/world/groupdino-new-city-s3r4/120253/weights_e10_0.7412.pth
+
+
+python eval_simidataset_parser_anyloc_rot.py --mode dinov2_vits14 --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/dinos.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt
+python eval_simidataset_parser_anyloc_rot.py --mode dinov2_vitb14 --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/dinob.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt
+python eval_simidataset_parser_anyloc_rot.py --mode dinov2_vitl14 --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/dinol.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt
+python eval_simidataset_parser_anyloc_rot.py --mode dinov2_vitg14 --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/dinog.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt
+
+python eval_simidataset_parser_game4loc_rot.py --save_txt /media/guan/新加卷/Code/worldloc_result/Rot/game4loc.txt --test_txt /media/guan/新加卷/EdgeBing/WorldLoc/ya1/test_country.txt

GeoLoc-UAV-main/eval_simidataset.py

@@ -0,0 +1,230 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+
+def default_group_config():
+
+    return {
+          "group_arch" : "groupdinonet", #group
+        "group_config": {
+            "none"
+        }
+    }
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "resnet18",  #dinov2_vitb14,resnet18
+          "pretrain_flag":True
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "multiconvap", #convap
+    "agg_config": {
+      "in_channels": 512, #256 #512，768
+      "out_channels": 512, #256
+      "s1": 1,
+      "s2": 1,
+      'LPN':False
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "resnet18"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world"
+
+    # model config
+    group:dict = field(default_factory=default_group_config)
+    backbone:dict = field(default_factory=default_backbone_config)
+
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc"
+    train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/Index/train_query.txt"
+
+    # val_index
+    val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/val.txt"
+
+    # test_index
+    test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/test_country.txt"
+    save_pred_txt = "/media/Shen/Data/RingoData/WorldLoc/txt/new_rot/dinos-finetune.txt"
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 if os.name == 'nt' else 4 
+    
+    # train on GPU if available
+    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 30
+    batch_size: int = 10                # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (1,)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = model.BackboneGlobal(config.backbone['backbone_arch'],
+                             config.backbone['pretrain_flag'],
+                               config.agg['agg_arch'],
+                               config.agg['agg_config'])
+
+# model = model.GrounpGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+
+# model = model.GrounpDinoGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+
+model_state_dict = torch.load("PTH/city/resnet_frozen/weights_e5_0.6413.pth", map_location='cuda:1')
+model.load_state_dict(model_state_dict, strict=False)
+
+model = model.to(config.device)
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+result_list_recall = []
+result_list_precision = []
+with open(config.test_index_txt,"r") as val_test:
+    for line in val_test:
+        eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                                name=line.strip('\n'),
+                                mode='query',
+                                transforms=eval_transform)
+
+        eval_dataloader_query = DataLoader(eva_dataset_query,
+                            batch_size=config.batch_size,
+                            num_workers=config.num_workers,
+                            shuffle=not config.custom_sampling,
+                            pin_memory=True)
+        
+        eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                                name=line.strip('\n'),
+                                mode='DB',
+                                transforms=eval_transform)
+
+        eval_dataloader_db = DataLoader(eva_dataset_db,
+                            batch_size=config.batch_size,
+                            num_workers=config.num_workers,
+                            shuffle=not config.custom_sampling,
+                            pin_memory=True)
+        
+        # eva_dataset_query = WorldDatasetEvalGroup(data_dir=config.dataset_root_dir,
+        #                         name=line.strip('\n'),
+        #                         mode='query',
+        #                         transforms=eval_transform)
+
+        # eval_dataloader_query = DataLoader(eva_dataset_query,
+        #                     batch_size=config.batch_size,
+        #                     num_workers=config.num_workers,
+        #                     shuffle=not config.custom_sampling,
+        #                     pin_memory=True)
+        
+        # eva_dataset_db = WorldDatasetEvalGroup(data_dir=config.dataset_root_dir,
+        #                         name=line.strip('\n'),
+        #                         mode='DB',
+        #                         transforms=eval_transform)
+
+        # eval_dataloader_db = DataLoader(eva_dataset_db,
+        #                     batch_size=config.batch_size,
+        #                     num_workers=config.num_workers,
+        #                     shuffle=not config.custom_sampling,
+        #                     pin_memory=True)
+        
+        pos_gt = eval_dataloader_db.dataset.get_gt()
+        result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='vanilia',LPN=config.agg['agg_config']['LPN'])
+        print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+      
+
+        # ap@5
+        ap_list = []
+        for i in range(predictions.shape[0]):
+            ex = np.isin(predictions[i, 5:], really_pos_gt[i][1])
+            num_all = np.sum(ex) / 5 * 100
+            ap_list.append(num_all)
+        average_ap = np.mean(np.array(ap_list))
+        
+        result_list_recall.append(result)
+        result_list_precision.append(average_ap)
+
+    result_array = np.array(result_list_recall)
+    average_result = np.mean(result_array, axis=0)
+    print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+    
+   
+    result_precision = np.array(result_list_precision)
+    av_p = np.mean(result_precision)
+    print('AP@5 is', round(av_p,2))
+
+# save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+#     for i in range(predictions.shape[0]):
+#         query_path = eval_dataloader_query.dataset.getitem(i)
+#         if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+#             num = 1
+#         else:
+#             num = 0
+#         pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+#         info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+#         f.write(info)
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_simidataset_parser.py

@@ -0,0 +1,273 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='group', help='Model architecture')
+    parser.add_argument('--model_path', type=str, default='./world', help='Path to save model checkpoints')
+    
+    # Group Config
+    parser.add_argument('--group_arch', type=str, default='groupdinonet', help='Group architecture')
+    parser.add_argument('--group_config', type=str, default='none', help='Group configuration')
+
+    # Backbone Config
+    parser.add_argument('--backbone_arch', type=str, default='dinov2_vits14', help='Backbone architecture')
+    parser.add_argument('--pretrain_flag', type=bool, default=True, help='Flag to use pre-trained weights')
+
+    # Agg Config
+    parser.add_argument('--agg_arch', type=str, default='multiconvap', help='Aggregation architecture')
+    parser.add_argument('--agg_in_channels', type=int, default=384, help='Input channels for aggregation')
+    parser.add_argument('--agg_out_channels', type=int, default=384, help='Output channels for aggregation')
+    parser.add_argument('--agg_s1', type=int, default=1, help='Aggregation s1 parameter')
+    parser.add_argument('--agg_s2', type=int, default=1, help='Aggregation s2 parameter')
+    parser.add_argument('--agg_LPN', type=bool, default=False, help='Use LPN for aggregation')
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc/ya1/', help='Root directory of the dataset')
+    parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    # Build the config dictionaries dynamically based on parsed args
+    group_config = {
+        "group_arch": args.group_arch,
+        "group_config": {args.group_config}
+    }
+
+    backbone_config = {
+        "backbone_arch": args.backbone_arch,
+        "pretrain_flag": args.pretrain_flag
+    }
+
+    agg_config = {
+        "agg_arch": args.agg_arch,
+        "agg_config": {
+            "in_channels": args.agg_in_channels,
+            "out_channels": args.agg_out_channels,
+            "s1": args.agg_s1,
+            "s2": args.agg_s2,
+            "LPN": args.agg_LPN
+        }
+    }
+
+    config = {
+        "mode": args.mode,
+        "model_path": args.model_path,
+        "group": group_config,
+        "backbone": backbone_config,
+        "agg": agg_config,
+        "dataset_root_dir": args.dataset_root,
+        "test_index_txt": args.test_txt,
+        "save_txt":args.save_txt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+if config["mode"] == 'vanilia':
+
+    model = model.BackboneGlobal(config["backbone"]['backbone_arch'],
+                             config["backbone"]['pretrain_flag'],
+                               config["agg"]['agg_arch'],
+                               config["agg"]['agg_config'])
+    model_state_dict = torch.load(config['checkpoint_path'], map_location=config['device'])
+    model.load_state_dict(model_state_dict, strict=False)
+
+    model = model.to(config['device'])
+
+# model = model.GrounpGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+else:
+    model = model.GrounpDinoGlobal(config["group"]['group_arch'],
+                                config["agg"]['agg_arch'],
+                                config["agg"]['agg_config'])
+
+    model_state_dict = torch.load(config['checkpoint_path'], map_location=config['device'])
+    model.load_state_dict(model_state_dict, strict=False)
+
+    model = model.to(config['device'])
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+result_list_recall = []
+result_list_precision = []
+with open(config['save_txt'], 'w') as f_w:
+    with open(config["test_index_txt"],"r") as val_test:
+        for line in val_test:
+            if config["mode"] == 'vanilia':
+                eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='query',
+                                        transforms=eval_transform)
+
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='DB',
+                                        transforms=eval_transform)
+
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+            else:
+                eva_dataset_query = WorldDatasetEvalGroup(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='query',
+                                        transforms=eval_transform)
+
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalGroup(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='DB',
+                                        transforms=eval_transform)
+
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+            
+            pos_gt = eval_dataloader_db.dataset.get_gt()
+            result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["mode"],LPN=config["agg"]['agg_config']['LPN'])
+            print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+            f_w.write(line + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
+
+
+            # ap@5
+            ap_list = []
+            for i in range(predictions.shape[0]):
+                ex = np.isin(predictions[i, 5:], really_pos_gt[i][1])
+                num_all = np.sum(ex) / 5 * 100
+                ap_list.append(num_all)
+            average_ap = np.mean(np.array(ap_list))
+            
+            result_list_recall.append(result)
+            result_list_precision.append(average_ap)
+            
+
+        result_array = np.array(result_list_recall)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        
+    
+        result_precision = np.array(result_list_precision)
+        av_p = np.mean(result_precision)
+        print('AP@5 is', round(av_p,2))
+
+        info1 = 'Average' + 'top 1: ' + str(round(average_result[0]*100,2)) + 'top 5: ' + str(round(average_result[1]*100,2)) + 'top 10:' + str(round(average_result[2]*100,2)) + '\n'
+        f_w.write(info1)
+        f_w.write('AP@5 is'+str(round(av_p,2)))
+
+        
+
+# save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+#     for i in range(predictions.shape[0]):
+#         query_path = eval_dataloader_query.dataset.getitem(i)
+#         if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+#             num = 1
+#         else:
+#             num = 0
+#         pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+#         info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+#         f.write(info)
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_simidataset_parser_anyloc.py

@@ -0,0 +1,195 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+from models.anyloc import AnyModel
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='dinov2_vitl14', help='Model architecture')
+    parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
+    
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc', help='Root directory of the dataset')
+    parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+
+    config = {
+        "mode": args.mode,
+        "model": args.model,
+        "dataset_root_dir": args.dataset_root,
+        "test_index_txt": args.test_txt,
+        "save_txt":args.save_txt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = AnyModel(model_name=config['mode'],
+                    pretrained=True)
+ 
+model = model.to(config["device"])
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+result_list_recall = []
+result_list_precision = []
+with open(config['save_txt'], 'w') as f_w:
+    with open(config["test_index_txt"],"r") as val_test:
+        for line in val_test:
+            if config["model"] == 'vanilia':
+                eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='query',
+                                        transforms=eval_transform)
+
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='DB',
+                                        transforms=eval_transform)
+
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+            
+            pos_gt = eval_dataloader_db.dataset.get_gt()
+            result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["model"],LPN=False)
+            print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+            f_w.write(line + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
+
+
+            # ap@5
+            ap_list = []
+            for i in range(predictions.shape[0]):
+                ex = np.isin(predictions[i, 5:], really_pos_gt[i][1])
+                num_all = np.sum(ex) / 5 * 100
+                ap_list.append(num_all)
+            average_ap = np.mean(np.array(ap_list))
+            
+            result_list_recall.append(result)
+            result_list_precision.append(average_ap)
+            
+
+        result_array = np.array(result_list_recall)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        
+    
+        result_precision = np.array(result_list_precision)
+        av_p = np.mean(result_precision)
+        print('AP@5 is', round(av_p,2))
+
+        info1 = 'Average' + 'top 1: ' + str(round(average_result[0]*100,2)) + 'top 5: ' + str(round(average_result[1]*100,2)) + 'top 10: ' + str(round(average_result[2]*100,2)) + '\n'
+        f_w.write(info1)
+        f_w.write('AP@5 is'+str(round(av_p,2)))
+
+        
+
+# save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+    for i in range(predictions.shape[0]):
+        query_path = eval_dataloader_query.dataset.getitem(i)
+        if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+            num = 1
+        else:
+            num = 0
+        pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+        info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+        f_w.write(info)
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_simidataset_parser_anyloc_rot.py

@@ -0,0 +1,161 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World_rot import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+from models.anyloc import AnyModel
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='dinov2_vitl14', help='Model architecture')
+    parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
+    
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc/ya1/', help='Root directory of the dataset')
+    parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+
+    config = {
+        "mode": args.mode,
+        "model": args.model,
+        "dataset_root_dir": args.dataset_root,
+        "test_index_txt": args.test_txt,
+        "save_txt":args.save_txt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = AnyModel(model_name=config['mode'],
+                    pretrained=True)
+ 
+model = model.to(config["device"])
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+rotation_angles = list(range(0, 360, 5))
+height_rot_list = [f"height100_rot{angle}" for angle in rotation_angles]
+with open(config['save_txt'], 'w') as f_w:
+    with open(config["test_index_txt"],"r") as val_test:
+        for line in val_test:
+            for height_mode in height_rot_list:
+                if config["model"] == 'vanilia':
+                    eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='query',
+                                            height_mode=height_mode,
+                                            transforms=eval_transform)
+
+                    eval_dataloader_query = DataLoader(eva_dataset_query,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                    
+                    eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='DB',
+                                            transforms=eval_transform)
+
+                    eval_dataloader_db = DataLoader(eva_dataset_db,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["model"],LPN=False)
+                print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+                f_w.write(str(height_mode) + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
+
+
+           
+
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_simidataset_parser_game4loc.py

@@ -0,0 +1,209 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+from models.game4loc import DesModel
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='vit_base_patch16_rope_reg1_gap_256.sbb_in1k', help='Model architecture')
+    parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
+    
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc/', help='Root directory of the dataset')
+    parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default="/media/guan/新加卷/Code/Code/vit_base_eva_gta_same_area.pth", help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+
+    config = {
+        "mode": args.mode,
+        "model": args.model,
+        "dataset_root_dir": args.dataset_root,
+        "test_index_txt": args.test_txt,
+        "save_txt":args.save_txt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((384, 384), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = DesModel(model_name='vit_base_patch16_rope_reg1_gap_256.sbb_in1k',
+                    pretrained=True,
+                    img_size=384,
+                    share_weights=True)
+
+if config["checkpoint_path"] is not None:  
+    print("Start from:", config["checkpoint_path"])
+    model_state_dict = torch.load(config["checkpoint_path"])  
+    model.load_state_dict(model_state_dict, strict=True)     
+
+    # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    # if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+    #     model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+            
+    # Model to device   
+    model = model.to(config["device"])
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+result_list_recall = []
+result_list_precision = []
+with open(config['save_txt'], 'w') as f_w:
+    with open(config["test_index_txt"],"r") as val_test:
+        for line in val_test:
+            if config["model"] == 'vanilia':
+                eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='query',
+                                        transforms=eval_transform)
+
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                        name=line.strip('\n'),
+                                        mode='DB',
+                                        transforms=eval_transform)
+
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                    batch_size=config["batch_size"],
+                                    num_workers=config["num_workers"],
+                                    shuffle=not config["custom_sampling"],
+                                    pin_memory=True)
+            
+            pos_gt = eval_dataloader_db.dataset.get_gt()
+            result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["model"],LPN=False)
+            
+            print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+            f_w.write(line + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
+
+
+            # ap@5
+            ap_list = []
+            for i in range(predictions.shape[0]):
+                ex = np.isin(predictions[i, 5:], really_pos_gt[i][1])
+                num_all = np.sum(ex) / 5 * 100
+                ap_list.append(num_all)
+            average_ap = np.mean(np.array(ap_list))
+            
+            result_list_recall.append(result)
+            result_list_precision.append(average_ap)
+            
+
+        result_array = np.array(result_list_recall)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        
+    
+        result_precision = np.array(result_list_precision)
+        av_p = np.mean(result_precision)
+        print('AP@5 is', round(av_p,2))
+
+        info1 = 'Average' + 'top 1: ' + str(round(average_result[0]*100,2)) + 'top 5: ' + str(round(average_result[1]*100,2)) +  'top 10' + str(round(average_result[2]*100,2)) + '\n'
+        f_w.write(info1)
+        f_w.write('AP@5 is'+str(round(av_p,2)))
+
+        
+
+# save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+    for i in range(predictions.shape[0]):
+        query_path = eval_dataloader_query.dataset.getitem(i)
+        if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+            num = 1
+        else:
+            num = 0
+        pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+        info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+        f_w.write(info)
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_simidataset_parser_game4loc_rot.py

@@ -0,0 +1,172 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World_rot import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+from models.game4loc import DesModel
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='vit_base_patch16_rope_reg1_gap_256.sbb_in1k', help='Model architecture')
+    parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
+    
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc/ya1/', help='Root directory of the dataset')
+    parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default="/media/guan/新加卷/Code/Code/vit_base_eva_gta_same_area.pth", help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+
+    config = {
+        "mode": args.mode,
+        "model": args.model,
+        "dataset_root_dir": args.dataset_root,
+        "test_index_txt": args.test_txt,
+        "save_txt":args.save_txt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((384, 384), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = DesModel(model_name='vit_base_patch16_rope_reg1_gap_256.sbb_in1k',
+                    pretrained=True,
+                    img_size=384,
+                    share_weights=True)
+
+if config["checkpoint_path"] is not None:  
+    print("Start from:", config["checkpoint_path"])
+    model_state_dict = torch.load(config["checkpoint_path"])  
+    model.load_state_dict(model_state_dict, strict=True)     
+
+    # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    # if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+    #     model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+            
+    # Model to device   
+    model = model.to(config["device"])
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+rotation_angles = list(range(0, 360, 5))
+height_rot_list = [f"height100_rot{angle}" for angle in rotation_angles]
+with open(config['save_txt'], 'w') as f_w:
+    with open(config["test_index_txt"],"r") as val_test:
+        for line in val_test:
+            for height_mode in height_rot_list:
+                if config["model"] == 'vanilia':
+                    eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='query',
+                                            height_mode=height_mode,
+                                            transforms=eval_transform)
+
+                    eval_dataloader_query = DataLoader(eva_dataset_query,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                    
+                    eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='DB',
+                                            transforms=eval_transform)
+
+                    eval_dataloader_db = DataLoader(eva_dataset_db,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["model"],LPN=False)
+                
+                print('top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+                f_w.write(height_mode + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
+
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_simidataset_parser_rot.py

@@ -0,0 +1,253 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World_rot import  WorldDatasetEvalVanilia, WorldDatasetEvalGroup
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import argparse
+
+def get_parser():
+    parser = argparse.ArgumentParser(description="Configuration for training the model")
+
+    # Model Configurations
+    parser.add_argument('--mode', type=str, default='group', help='Model architecture')
+    parser.add_argument('--model_path', type=str, default='./world', help='Path to save model checkpoints')
+    
+    # Group Config
+    parser.add_argument('--group_arch', type=str, default='groupdinonet', help='Group architecture')
+    parser.add_argument('--group_config', type=str, default='none', help='Group configuration')
+
+    # Backbone Config
+    parser.add_argument('--backbone_arch', type=str, default='dinov2_vits14', help='Backbone architecture')
+    parser.add_argument('--pretrain_flag', type=bool, default=True, help='Flag to use pre-trained weights')
+
+    # Agg Config
+    parser.add_argument('--agg_arch', type=str, default='multiconvap', help='Aggregation architecture')
+    parser.add_argument('--agg_in_channels', type=int, default=384, help='Input channels for aggregation')
+    parser.add_argument('--agg_out_channels', type=int, default=384, help='Output channels for aggregation')
+    parser.add_argument('--agg_s1', type=int, default=1, help='Aggregation s1 parameter')
+    parser.add_argument('--agg_s2', type=int, default=1, help='Aggregation s2 parameter')
+    parser.add_argument('--agg_LPN', type=bool, default=False, help='Use LPN for aggregation')
+
+    # Dataset Paths
+    parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc/ya1/', help='Root directory of the dataset')
+    parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+    parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
+   
+   #'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
+    # Checkpoint Config
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
+
+    # Training Parameters
+    parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
+    parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
+    parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
+    parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
+
+    # Training Settings
+    parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
+    parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
+    parser.add_argument('--seed', type=int, default=1, help='Random seed')
+    parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
+    parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
+    parser.add_argument('--gpu_ids', type=tuple, default=(1,), help='GPU IDs for training')
+
+    # Optimizer Config
+    parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
+    parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
+    parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
+
+    # Loss Config
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
+
+    # Learning Rate
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
+    parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
+    parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
+    parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
+
+    return parser
+
+def parse_config():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    # Build the config dictionaries dynamically based on parsed args
+    group_config = {
+        "group_arch": args.group_arch,
+        "group_config": {args.group_config}
+    }
+
+    backbone_config = {
+        "backbone_arch": args.backbone_arch,
+        "pretrain_flag": args.pretrain_flag
+    }
+
+    agg_config = {
+        "agg_arch": args.agg_arch,
+        "agg_config": {
+            "in_channels": args.agg_in_channels,
+            "out_channels": args.agg_out_channels,
+            "s1": args.agg_s1,
+            "s2": args.agg_s2,
+            "LPN": args.agg_LPN
+        }
+    }
+
+    config = {
+        "mode": args.mode,
+        "model_path": args.model_path,
+        "group": group_config,
+        "backbone": backbone_config,
+        "agg": agg_config,
+        "dataset_root_dir": args.dataset_root,
+        "test_index_txt": args.test_txt,
+        "save_txt":args.save_txt,
+        "checkpoint_path": args.checkpoint_path,
+        "num_workers": args.num_workers,
+        "device": args.device,
+        "cudnn_benchmark": args.cudnn_benchmark,
+        "cudnn_deterministic": args.cudnn_deterministic,
+        "mixed_precision": args.mixed_precision,
+        "custom_sampling": args.custom_sampling,
+        "seed": args.seed,
+        "epochs": args.epochs,
+        "batch_size": args.batch_size,
+        "verbose": args.verbose,
+        "gpu_ids": args.gpu_ids,
+        "clip_grad": args.clip_grad,
+        "decay_exclude_bias": args.decay_exclude_bias,
+        "grad_checkpointing": args.grad_checkpointing,
+        "label_smoothing": args.label_smoothing,
+        "lr": args.lr,
+        "scheduler": args.scheduler,
+        "warmup_epochs": args.warmup_epochs,
+        "lr_end": args.lr_end
+    }
+
+    return args, config
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+args, config = parse_config()
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+if config["mode"] == 'vanilia':
+
+    model = model.BackboneGlobal(config["backbone"]['backbone_arch'],
+                             config["backbone"]['pretrain_flag'],
+                               config["agg"]['agg_arch'],
+                               config["agg"]['agg_config'])
+    model_state_dict = torch.load(config['checkpoint_path'], map_location=config['device'])
+    model.load_state_dict(model_state_dict, strict=False)
+
+    model = model.to(config['device'])
+
+# model = model.GrounpGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+else:
+    model = model.GrounpDinoGlobal(config["group"]['group_arch'],
+                                config["agg"]['agg_arch'],
+                                config["agg"]['agg_config'])
+
+    model_state_dict = torch.load(config['checkpoint_path'], map_location=config['device'])
+    model.load_state_dict(model_state_dict, strict=False)
+
+    model = model.to(config['device'])
+
+ #------------------------------------------------------------Eval---------------------------------------------------------------------#
+rotation_angles = list(range(0, 360, 5))
+height_rot_list = [f"height100_rot{angle}" for angle in rotation_angles]
+with open(config['save_txt'], 'w') as f_w:
+    with open(config["test_index_txt"],"r") as val_test:
+        for line in val_test:
+            for height_mode in height_rot_list:
+                if config["mode"] == 'vanilia':
+                    eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='query',
+                                            height_mode=height_mode,
+                                            transforms=eval_transform)
+
+                    eval_dataloader_query = DataLoader(eva_dataset_query,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                    
+                    eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='DB',
+                                            transforms=eval_transform)
+
+                    eval_dataloader_db = DataLoader(eva_dataset_db,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                else:
+                    eva_dataset_query = WorldDatasetEvalGroup(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='query',
+                                            height_mode=height_mode,
+                                            transforms=eval_transform)
+
+                    eval_dataloader_query = DataLoader(eva_dataset_query,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                    
+                    eva_dataset_db = WorldDatasetEvalGroup(data_dir=config["dataset_root_dir"],
+                                            name=line.strip('\n'),
+                                            mode='DB',
+                                            transforms=eval_transform)
+
+                    eval_dataloader_db = DataLoader(eva_dataset_db,
+                                        batch_size=config["batch_size"],
+                                        num_workers=config["num_workers"],
+                                        shuffle=not config["custom_sampling"],
+                                        pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["mode"],LPN=config["agg"]['agg_config']['LPN'])
+                print(height_mode,' ','top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
+                f_w.write(height_mode + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
+
+
+           
+
+        
+
+# save top 1 flase or wrong
+# with open(config.save_pred_txt, 'w') as f:
+#     for i in range(predictions.shape[0]):
+#         query_path = eval_dataloader_query.dataset.getitem(i)
+#         if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
+#             num = 1
+#         else:
+#             num = 0
+#         pred_path = eval_dataloader_db.dataset.samples[predictions[i,0]]
+#         info = query_path +  ' ' + pred_path + ' ' + str(num) + '\n'
+#         f.write(info)
+
+
+
+                       
+                    

GeoLoc-UAV-main/eval_vis.py

@@ -0,0 +1,214 @@
+from torch.utils.data import DataLoader
+from dataclasses import dataclass,field
+from eval import eval
+import os
+import torch
+from torchvision import transforms as T
+from dataset.World import  WorldDatasetEval
+from models import model
+import glob
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+
+def default_group_config():
+
+    return {
+          "group_arch" : "groupnet", #group
+        "group_config": {
+            "none"
+        }
+    }
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "dinov2_vitb14", 
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "convap", #convap
+    "agg_config": {
+      "in_channels": 768, #256 #512
+      "out_channels": 768, #256
+      "s1": 1,
+      "s2": 1
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "resnet18"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world"
+
+    # model config
+    group:dict = field(default_factory=default_group_config)
+    backbone:dict = field(default_factory=default_backbone_config)
+
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/guan/新加卷/EdgeBing/WorldLoc"
+    train_query_txt: str = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/train_query.txt"
+
+    # val_index
+    val_index_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/val.txt"
+
+    # test_index
+    test_index_txt = "/media/guan/新加卷/EdgeBing/WorldLoc/Index/test.txt"
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 if os.name == 'nt' else 4 
+    
+    # train on GPU if available
+    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 30
+    batch_size: int = 28                # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (0,1,2,3)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+    
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                    'std': [0.229, 0.224, 0.225]}
+eval_transform = T.Compose([
+        T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+        T.ToTensor(),
+        T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+    ])
+
+model = model.BackboneGlobal(config.backbone['backbone_arch'],
+                               config.agg['agg_arch'],
+                               config.agg['agg_config'])
+
+# model = model.GrounpGlobal(config.group['group_arch'],
+#                                config.agg['agg_arch'],
+#                                config.agg['agg_config'])
+
+model_state_dict = torch.load("world/dinov2-base/094102/weights_e18_0.1987.pth")
+model.load_state_dict(model_state_dict, strict=False)
+
+model = model.to(config.device)
+
+with open(config.val_index_txt,"r") as val_test:
+        for line in val_test:
+            eva_dataset_query = WorldDatasetEval(data_dir=config.dataset_root_dir,
+                                    name=line.strip('\n'),
+                                    mode='query',
+                                    transforms=eval_transform)
+
+            eval_dataloader_query = DataLoader(eva_dataset_query,
+                                batch_size=config.batch_size,
+                                num_workers=config.num_workers,
+                                shuffle=not config.custom_sampling,
+                                pin_memory=True)
+            
+            eva_dataset_db = WorldDatasetEval(data_dir=config.dataset_root_dir,
+                                    name=line.strip('\n'),
+                                    mode='DB',
+                                    transforms=eval_transform)
+
+            eval_dataloader_db = DataLoader(eva_dataset_db,
+                                batch_size=config.batch_size,
+                                num_workers=config.num_workers,
+                                shuffle=not config.custom_sampling,
+                                pin_memory=True)
+            
+            pos_gt = eval_dataloader_db.dataset.get_gt()
+            result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, ranks=[1, 5, 10], name=line.strip('\n'))
+            print(line.strip('\n'), 'top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2))
+
+
+            # vis and save retrieval results
+            # save_vis_dir = config.dataset_root_dir + '/' + line.strip('\n') + '/' + 'resnet18' + '/'
+            # if not os.path.exists(save_vis_dir):
+            #     os.makedirs(save_vis_dir)
+            
+            # temp_path = os.path.join(config.dataset_root_dir,  line.strip('\n'), 'DB', 'img')
+            # DB_path = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
+            # for i in range(predictions.shape[0]):
+            #     query_path = eval_dataloader_query.dataset.getitem(i)
+            #     fig, axs = plt.subplots(2, 6, figsize=(15, 5)) 
+            #     query_img = plt.imread(query_path)
+            #     for j in range(2):
+            #         for k in range(6):
+            #             if j == 0 and k == 0:
+            #                 axs[j, k].imshow(query_img)
+            #                 axs[j, k].axis('off')  # 不显示坐标轴
+            #             elif j==0 and k != 0:
+            #                 if np.any(np.in1d(predictions[i,k], really_pos_gt[i][1] )):
+
+            #                     db_img_path = DB_path[predictions[i,k]]
+            #                     db_img = plt.imread(db_img_path)
+            #                     axs[j, k].imshow(db_img)
+            #                       # 创建一个矩形框
+            #                     rect = patches.Rectangle((0, 0), 1, 1, linewidth=10, edgecolor='green', facecolor='none')
+                                
+            #                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+            #                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+            #                     axs[j, k].add_patch(rect)
+            #                     axs[j,k].axis('off')  # 不显示坐标轴
+            #                 else:
+            #                     db_img_path = DB_path[predictions[i,k]]
+            #                     db_img = plt.imread(db_img_path)
+            #                     axs[j, k].imshow(db_img)
+            #                       # 创建一个矩形框
+            #                     rect = patches.Rectangle((0, 0), 1, 1, linewidth=10, edgecolor='red', facecolor='none')
+                                
+            #                     # 将矩形框添加到图像上，根据图像尺寸调整框的大小
+            #                     rect.set_transform(axs[j, k].transData)  # 将框的坐标系设置为数据坐标系
+            #                     axs[j, k].add_patch(rect)
+            #                     axs[j, k].axis('off')  # 不显示坐标轴
+            #             if j ==1:
+            #                 try:
+            #                     db_img_path = DB_path[really_pos_gt[i][1][k]]
+            #                     db_img = plt.imread(db_img_path)
+            #                     axs[j, k].imshow(db_img)
+            #                     axs[j, k].axis('off')  # 不显示坐标轴
+            #                 except:
+            #                     break
+                       
+            #     save_one_path = save_vis_dir + str(i) + '.png'
+            #     plt.savefig(save_one_path, dpi=300)
+                       
+                    

GeoLoc-UAV-main/models/aggregators/LPN.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import math
+from torch.nn import functional as F
+
+def get_part_pool(x, block=4, no_overlap=True):
+        result = []
+        H, W = x.size(2), x.size(3)
+        c_h, c_w = int(H/2), int(W/2)
+        per_h, per_w = H/(2*block),W/(2*block)
+        if per_h < 1 and per_w < 1:
+            new_H, new_W = H+(block-c_h)*2, W+(block-c_w)*2
+            x = nn.functional.interpolate(x, size=[new_H,new_W], mode='bilinear', align_corners=True)
+            H, W = x.size(2), x.size(3)
+            c_h, c_w = int(H/2), int(W/2)
+            per_h, per_w = H/(2*block),W/(2*block)
+        per_h, per_w = math.floor(per_h), math.floor(per_w)
+        for i in range(block):
+            i = i + 1
+            if i < block:
+                x_curr = x[:,:,(c_h-i*per_h):(c_h+i*per_h),(c_w-i*per_w):(c_w+i*per_w)]
+                if no_overlap and i > 1:
+                    x_pre = x[:,:,(c_h-(i-1)*per_h):(c_h+(i-1)*per_h),(c_w-(i-1)*per_w):(c_w+(i-1)*per_w)] 
+                    x_pad = F.pad(x_pre,(per_h,per_h,per_w,per_w),"constant",0)
+                    x_curr = x_curr - x_pad
+                result.append(x_curr)
+            else:
+                if no_overlap and i > 1:
+                    x_pre = x[:,:,(c_h-(i-1)*per_h):(c_h+(i-1)*per_h),(c_w-(i-1)*per_w):(c_w+(i-1)*per_w)]
+                    pad_h = c_h-(i-1)*per_h
+                    pad_w = c_w-(i-1)*per_w
+                    # x_pad = F.pad(x_pre,(pad_h,pad_h,pad_w,pad_w),"constant",0)
+                    if x_pre.size(2)+2*pad_h == H:
+                        x_pad = F.pad(x_pre,(pad_h,pad_h,pad_w,pad_w),"constant",0)
+                    else:
+                        ep = H - (x_pre.size(2)+2*pad_h)
+                        x_pad = F.pad(x_pre,(pad_h+ep,pad_h,pad_w+ep,pad_w),"constant",0)
+                    x = x - x_pad
+                result.append(x_curr)
+        return result

GeoLoc-UAV-main/models/aggregators/__init__.py

@@ -0,0 +1,3 @@
+from .gem import GeMPool
+from .convap import ConvAP
+from .multiconvap import MulConvAP

GeoLoc-UAV-main/models/aggregators/convap.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+
+class ConvAP(nn.Module):
+    """Implementation of ConvAP as of https://arxiv.org/pdf/2210.10239.pdf
+
+    Args:
+        in_channels (int): number of channels in the input of ConvAP
+        out_channels (int, optional): number of channels that ConvAP outputs. Defaults to 512.
+        s1 (int, optional): spatial height of the adaptive average pooling. Defaults to 2.
+        s2 (int, optional): spatial width of the adaptive average pooling. Defaults to 2.
+    """
+    def __init__(self, in_channels, out_channels=512, s1=2, s2=2):
+        super(ConvAP, self).__init__()
+        self.channel_pool = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=True)
+        self.AAP = nn.AdaptiveAvgPool2d((s1, s2))
+
+    def forward(self, x):
+        #
+        x, t = x #dinov2专属
+        # x = self.channel_pool(x)
+        x = self.AAP(x)
+        x = F.normalize(x.flatten(1), p=2, dim=1)
+        return x
+    
+
+if __name__ == '__main__':
+    x = torch.randn(4, 2048, 10, 10)
+    m = ConvAP(2048, 512)
+    r = m(x)
+    print(r.shape)

GeoLoc-UAV-main/models/aggregators/gem.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+class GeMPool(nn.Module):
+    """Implementation of GeM as in https://github.com/filipradenovic/cnnimageretrieval-pytorch
+    we add flatten and norm so that we can use it as one aggregation layer.
+    """
+    def __init__(self, p=3, eps=1e-6):
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1)*p)
+        self.eps = eps
+
+    def forward(self, x):
+        x = F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1./self.p)
+        x = x.flatten(1)
+        return F.normalize(x, p=2, dim=1)

GeoLoc-UAV-main/models/aggregators/multiconvap.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+from models.aggregators.LPN import get_part_pool
+
+class L2Norm(nn.Module):
+    def __init__(self, dim=1):
+        super().__init__()
+        self.dim = dim
+    def forward(self, x):
+        return F.normalize(x, p=2, dim=self.dim)
+
+class GeMPool(nn.Module):
+    """Implementation of GeM as in https://github.com/filipradenovic/cnnimageretrieval-pytorch
+    we add flatten and norm so that we can use it as one aggregation layer.
+    """
+    def __init__(self, p=3, eps=1e-6):
+        super().__init__()
+        self.p = nn.Parameter(torch.ones(1)*p)
+        self.eps = eps
+
+    def forward(self, x):
+        x = F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1./self.p)
+        x = x.flatten(1)
+        return F.normalize(x, p=2, dim=1)
+
+class MulConvAP(nn.Module):
+    """Implementation of ConvAP as of https://arxiv.org/pdf/2210.10239.pdf
+
+    Args:
+        in_channels (int): number of channels in the input of ConvAP
+        out_channels (int, optional): number of channels that ConvAP outputs. Defaults to 512.
+        s1 (int, optional): spatial height of the adaptive average pooling. Defaults to 2.
+        s2 (int, optional): spatial width of the adaptive average pooling. Defaults to 2.
+    """
+    def __init__(self, in_channels, out_channels=512, s1=2, s2=2, LPN=False):
+        super(MulConvAP, self).__init__()
+        self.out_channels = out_channels
+        self.channel_pool_1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1, bias=True)
+        self.channel_pool_3 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=3, padding=1,bias=True)
+        self.channel_pool_5 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=5, padding=2,bias=True)
+      
+        # self.AAP = nn.AdaptiveAvgPool2d((s1, s2))
+        self.AAP = nn.Sequential(L2Norm(), GeMPool())
+
+        # using LPN
+        if LPN == True:
+            self.LPN = True
+        else:
+            self.LPN = False
+    def forward(self, x):
+
+        if self.LPN == False:
+            # x, t = x #dinov2专属
+            x1 = self.channel_pool_1(x)
+            x3 = self.channel_pool_3(x)
+            x5 = self.channel_pool_5(x)
+
+            x1 = self.AAP(x1)
+            x3 = self.AAP(x3)
+            x5 = self.AAP(x5)
+
+            x = [i for i in [x1, x3, x5]]
+            x = torch.cat(x,dim=1)
+
+            # x = self.AAP(x)
+            x = F.normalize(x.flatten(1), p=2, dim=1)
+            return x
+        else:
+            partition_feature = get_part_pool(x)
+            partition_feature_list = []
+            for one_feature in partition_feature:
+                x1 = self.channel_pool_1(one_feature)
+                x3 = self.channel_pool_3(one_feature)
+                x5 = self.channel_pool_5(one_feature)
+
+                x1 = self.AAP(x1)
+                x3 = self.AAP(x3)
+                x5 = self.AAP(x5)
+
+                x = [i for i in [x1, x3, x5]]
+                x = torch.cat(x,dim=1)
+
+                x = F.normalize(x.flatten(1), p=2, dim=1)
+                partition_feature_list.append(x)
+            # partition_feature_tensor = torch.stack(partition_feature_list, dim=2).reshape(x.shape[0], -1)
+            
+            return partition_feature_list
+
+
+if __name__ == '__main__':
+    x = torch.randn(4, 2048, 10, 10)
+    # m = ConvAP(2048, 512)
+    # r = m(x)
+    # print(r.shape)

GeoLoc-UAV-main/models/anyloc.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from models import aggregators
+
+DINOV2_ARCHS = {
+    'dinov2_vits14': 384,
+    'dinov2_vitb14': 768,
+    'dinov2_vitl14': 1024,
+    'dinov2_vitg14': 1536,
+}
+
+class AnyModel(nn.Module):
+
+    def __init__(self, 
+                 model_name='dinov2_vitb14',
+                 pretrained=True,
+                 ):
+                 
+        super(AnyModel, self).__init__()
+
+        assert model_name in DINOV2_ARCHS.keys(), f'Unknown model name {model_name}'
+        self.model = torch.hub.load('facebookresearch/dinov2', model_name)
+        self.num_channels = DINOV2_ARCHS[model_name]
+        self.gem = aggregators.GeMPool()
+
+    def forward(self, x):
+
+        B, C, H, W = x.shape
+
+        x = self.model.prepare_tokens_with_masks(x)
+        
+        # First blocks are frozen
+        with torch.no_grad():
+            for blk in self.model.blocks:
+                x = blk(x)
+        x = x.detach()
+
+        t = x[:, 0]
+        f = x[:, 1:]
+
+        # Reshape to (B, C, H, W)
+        f = f.reshape((B, H // 14, W // 14, self.num_channels)).permute(0, 3, 1, 2)
+
+        g = self.gem(f)
+
+        return g

GeoLoc-UAV-main/models/backbone/__init__.py

@@ -0,0 +1,2 @@
+from .resnet import ResNet
+from .dinov2 import DINOv2

GeoLoc-UAV-main/models/backbone/dinov2.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+
+DINOV2_ARCHS = {
+    'dinov2_vits14': 384,
+    'dinov2_vitb14': 768,
+    'dinov2_vitl14': 1024,
+    'dinov2_vitg14': 1536,
+}
+
+class DINOv2(nn.Module):
+    """
+    DINOv2 model
+
+    Args:
+        model_name (str): The name of the model architecture 
+            should be one of ('dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14')
+        num_trainable_blocks (int): The number of last blocks in the model that are trainable.
+        norm_layer (bool): If True, a normalization layer is applied in the forward pass.
+        return_token (bool): If True, the forward pass returns both the feature map and the token.
+    """
+    def __init__(
+            self,
+            model_name='dinov2_vitb14',
+            num_trainable_blocks=2,
+            norm_layer=False,
+            return_token=False,
+            pretrain_flag=False
+        ):
+        super().__init__()
+
+        assert model_name in DINOV2_ARCHS.keys(), f'Unknown model name {model_name}'
+        
+        self.model = torch.hub.load('facebookresearch/dinov2', "dinov2_vits14")
+        # torch.hub.load('/home/Shen/.cache/torch/hub/facebookresearch_dinov2_main/',
+                                    # model_name,
+            # source='local')
+
+        # self.model = torch.hub.load('facebookresearch/dinov2', model_name)
+        self.num_channels = DINOV2_ARCHS[model_name]
+        self.num_trainable_blocks = num_trainable_blocks
+        self.norm_layer = norm_layer
+        self.return_token = return_token
+        self.flag = pretrain_flag
+
+
+    def forward(self, x):
+        """
+        The forward method for the DINOv2 class
+
+        Parameters:
+            x (torch.Tensor): The input tensor [B, 3, H, W]. H and W should be divisible by 14.
+
+        Returns:
+            f (torch.Tensor): The feature map [B, C, H // 14, W // 14].
+            t (torch.Tensor): The token [B, C]. This is only returned if return_token is True.
+        """
+
+        B, C, H, W = x.shape
+
+        x = self.model.prepare_tokens_with_masks(x)
+        if self.flag:
+        # When flag is True, freeze all parameters
+            for param in self.model.parameters():
+                param.requires_grad = False
+            with torch.no_grad():
+                for blk in self.model.blocks:
+                    x = blk(x)
+        else:
+            # When flag is False, freeze part of the parameters (e.g., first blocks)
+            for param in self.model.parameters():
+                param.requires_grad = False  # Freeze all layers initially
+            # Unfreeze the last few blocks (trainable)
+            for param in self.model.blocks[-self.num_trainable_blocks:].parameters():
+                param.requires_grad = True
+
+            with torch.no_grad():
+                for blk in self.model.blocks[:-self.num_trainable_blocks]:  # Freeze these blocks
+                    x = blk(x)
+            # Last blocks are trained
+            for blk in self.model.blocks[-self.num_trainable_blocks:]:  # Train these blocks
+                x = blk(x)
+            if self.norm_layer:
+                x = self.model.norm(x)
+        
+        t = x[:, 0]
+        f = x[:, 1:]
+
+        # Reshape to (B, C, H, W)
+        f = f.reshape((B, H // 14, W // 14, self.num_channels)).permute(0, 3, 1, 2)
+
+        if self.return_token:
+            return f, t
+        return f

GeoLoc-UAV-main/models/backbone/resnet.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torchvision
+import numpy as np
+
+class ResNet(nn.Module):
+    def __init__(self,
+                 model_name='resnet50',
+                 pretrained=True,
+                 layers_to_freeze=2,
+                 layers_to_crop=[],
+                 pretrain_flag = False
+                 ):
+        """Class representing the resnet backbone used in the pipeline
+        we consider resnet network as a list of 5 blocks (from 0 to 4),
+        layer 0 is the first conv+bn and the other layers (1 to 4) are the rest of the residual blocks
+        we don't take into account the global pooling and the last fc
+
+        Args:
+            model_name (str, optional): The architecture of the resnet backbone to instanciate. Defaults to 'resnet50'.
+            pretrained (bool, optional): Whether pretrained or not. Defaults to True.
+            layers_to_freeze (int, optional): The number of residual blocks to freeze (starting from 0) . Defaults to 2.
+            layers_to_crop (list, optional): Which residual layers to crop, for example [3,4] will crop the third and fourth res blocks. Defaults to [].
+
+        Raises:
+            NotImplementedError: if the model_name corresponds to an unknown architecture. 
+        """
+        super().__init__()
+        self.model_name = model_name.lower()
+        self.layers_to_freeze = layers_to_freeze
+        self.flag = pretrain_flag
+
+        if pretrained:
+            # the new naming of pretrained weights, you can change to V2 if desired.
+            weights = 'IMAGENET1K_V1'
+        else:
+            weights = None
+
+        if 'swsl' in model_name or 'ssl' in model_name:
+            # These are the semi supervised and weakly semi supervised weights from Facebook
+            self.model = torch.hub.load(
+                'facebookresearch/semi-supervised-ImageNet1K-models', model_name)
+        else:
+            if 'resnext50' in model_name:
+                self.model = torchvision.models.resnext50_32x4d(
+                    weights=weights)
+            elif 'resnet50' in model_name:
+                self.model = torchvision.models.resnet50(weights=weights)
+            elif '101' in model_name:
+                self.model = torchvision.models.resnet101(weights=weights)
+            elif '152' in model_name:
+                self.model = torchvision.models.resnet152(weights=weights)
+            elif '34' in model_name:
+                self.model = torchvision.models.resnet34(weights=weights)
+            elif '18' in model_name:
+                # self.model = torchvision.models.resnet18(pretrained=False)
+                self.model = torchvision.models.resnet18(weights=weights)
+            elif 'wide_resnet50_2' in model_name:
+                self.model = torchvision.models.wide_resnet50_2(
+                    weights=weights)
+            else:
+                raise NotImplementedError(
+                    'Backbone architecture not recognized!')
+
+        # freeze only if the model is pretrained
+        if pretrained and self.flag:
+            if layers_to_freeze >= 0:
+                self.model.conv1.requires_grad_(False)
+                self.model.bn1.requires_grad_(False)
+            if layers_to_freeze >= 1:
+                self.model.layer1.requires_grad_(False)
+            if layers_to_freeze >= 2:
+                self.model.layer2.requires_grad_(False)
+            if layers_to_freeze >= 3:
+                self.model.layer3.requires_grad_(False)
+
+        # remove the avgpool and most importantly the fc layer
+        self.model.avgpool = None
+        self.model.fc = None
+
+        if 4 in layers_to_crop:
+            self.model.layer4 = None
+        if 3 in layers_to_crop:
+            self.model.layer3 = None
+
+        out_channels = 2048
+        if '34' in model_name or '18' in model_name:
+            out_channels = 256
+            
+        self.out_channels = out_channels // 2 if self.model.layer4 is None else out_channels
+        self.out_channels = self.out_channels // 2 if self.model.layer3 is None else self.out_channels
+
+    def forward(self, x):
+       
+        x = self.model.conv1(x)
+        x = self.model.bn1(x)
+        x = self.model.relu(x)
+        x = self.model.maxpool(x)
+        x = self.model.layer1(x)
+        x = self.model.layer2(x)
+        if self.model.layer3 is not None:
+            x = self.model.layer3(x)
+        if self.model.layer4 is not None:
+            x = self.model.layer4(x)
+
+        return x
+      

GeoLoc-UAV-main/models/game4loc.py

@@ -0,0 +1,167 @@
+import torch
+import timm
+import numpy as np
+import torch.nn as nn
+from PIL import Image
+from urllib.request import urlopen
+from thop import profile
+
+
+class MLP(nn.Module):
+    def __init__(self, input_size=2048, hidden_size=512, output_size=2):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_size, hidden_size // 2)
+        self.fc3 = nn.Linear(hidden_size // 2, output_size)
+    
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.relu(x)
+        x = self.fc3(x)
+        return x
+
+
+class DesModel(nn.Module):
+
+    def __init__(self, 
+                 model_name='vit',
+                 pretrained=True,
+                 img_size=384,
+                 share_weights=True,
+                 train_with_recon=False,
+                 train_with_offset=False,
+                 model_hub='timm'):
+                 
+        super(DesModel, self).__init__()
+        self.share_weights = share_weights
+        self.model_name = model_name
+        self.img_size = img_size
+        if share_weights:
+            if "vit" in model_name or "swin" in model_name:
+                # automatically change interpolate pos-encoding to img_size
+                self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0, img_size=img_size) 
+            else:
+                self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
+        else:
+            if "vit" in model_name or "swin" in model_name:
+                self.model1 = timm.create_model(model_name, pretrained=pretrained, num_classes=0, img_size=img_size)
+                self.model2 = timm.create_model(model_name, pretrained=pretrained, num_classes=0, img_size=img_size) 
+            else:
+                self.model1 = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
+                self.model2 = timm.create_model(model_name, pretrained=pretrained, num_classes=0)
+
+        if train_with_offset:
+            self.MLP = MLP()
+        
+        self.logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        
+
+    def get_config(self,):
+        if self.share_weights:
+            data_config = timm.data.resolve_model_data_config(self.model)
+        else:
+            data_config = timm.data.resolve_model_data_config(self.model1)
+        return data_config
+    
+    
+    def set_grad_checkpointing(self, enable=True):
+        if self.share_weights:
+            self.model.set_grad_checkpointing(enable)
+        else:
+            self.model1.set_grad_checkpointing(enable)
+            self.model2.set_grad_checkpointing(enable)
+
+    def freeze_layers(self, frozen_blocks=10, frozen_stages=[0,0,0,0]):
+        pass
+
+    def forward(self, img1=None, img2=None):
+
+        if self.share_weights:
+            if img1 is not None and img2 is not None:
+                image_features1 = self.model(img1)     
+                image_features2 = self.model(img2)
+                return image_features1, image_features2            
+            elif img1 is not None:
+                image_features = self.model(img1)
+                return image_features
+            else:
+                image_features = self.model(img2)
+                return image_features
+        else:
+            if img1 is not None and img2 is not None:
+                image_features1 = self.model1(img1)     
+                image_features2 = self.model2(img2)
+                return image_features1, image_features2            
+            elif img1 is not None:
+                image_features = self.model1(img1)
+                return image_features
+            else:
+                image_features = self.model2(img2)
+                return image_features
+
+    def offset_pred(self, img_feature1, img_feature2):
+        offset = self.MLP(torch.cat((img_feature1, img_feature2), dim=1))
+        return offset
+
+
+if __name__ == '__main__':
+    # model = TimmModel(model_name='timm/vit_large_patch16_384.augreg_in21k_ft_in1k')
+    # # model = TimmModel(model_name='timm/vit_base_patch16_224.augreg_in1k')
+    # # from timm.models.vision_transformer import vit_base_patch16_224
+    # # model = vit_base_patch16_224(img_size=384, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, num_classes=0)
+
+
+    # model = DesModel(model_name='timm/resnet101.tv_in1k', img_size=384)
+    # model = DesModel(model_name='convnext_base.fb_in22k_ft_in1k_384', img_size=384)
+    model = DesModel(model_name='timm/swin_base_patch4_window7_224.ms_in22k_ft_in1k', img_size=384)
+    # # model = TimmModel(model_name='vit_base_patch16_rope_reg1_gap_256.sbb_in1k')
+    # # model = TimmModel(model_name='timm/vit_medium_patch16_rope_reg1_gap_256.sbb_in1k')
+    # # model = TimmModel(model_name='timm/vit_medium_patch16_gap_256.sw_in12k_ft_in1k')
+    # # model = TimmModel(model_name='timm/resnet101.tv_in1k') 
+    # # img = Image.open(urlopen(
+    # # 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
+    # # ))
+    x = torch.rand((1, 3, 384, 384))
+    x = x.cuda()
+    model.cuda()
+    x = model(x)
+    print(x.shape)
+
+    # flops, params = profile(model, inputs=(x,))
+    # # print(img.size)
+    # # img = transform(img)
+    # # print(img.size)
+
+    # # print(model1)
+    # print('flops(G)', flops/1e9, 'params(M)', params/1e6)
+
+    # from transformers import CLIPProcessor, CLIPModel
+    # model = CLIPModel.from_pretrained("/home/xmuairmud/jyx/clip-vit-base-patch16")
+    # vision_model = model.vision_model
+    # print(vision_model)
+
+    # dinov2_vitb14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg')
+    # print(dinov2_vitb14_reg.set_grad_checkpointing(True))
+
+    # from transformers import ViTModel, ViTImageProcessor, AutoModelForImageClassification, AutoConfig
+    # config = AutoConfig.from_pretrained('facebook/dino-vitb16')
+    # config.image_size = 384
+    # model = ViTModel.from_pretrained('facebook/dino-vitb16', config=config, ignore_mismatched_sizes=True)
+    # model = timm.create_model('vit_base_patch14_reg4_dinov2.lvd142m', pretrained=True, img_size=(384, 384))
+    # data_config = timm.data.resolve_model_data_config(model)
+    # print(data_config)
+    # processor = ViTImageProcessor.from_pretrained('facebook/dino-vitb16')
+
+
+    # x = torch.rand((1, 3, 384, 384))
+    # inputs = processor(images=x, return_tensors="pt")
+    # print(inputs['pixel_values'].shape)
+    # outputs = model(**inputs)
+    # print(outputs.pooler_output.shape)
+    # print(model(x).shape)
+    # flops, params = profile(dinov2_vitb14_reg, inputs=(x,))
+    # print('flops(G)', flops/1e9, 'params(M)', params/1e6)
+

GeoLoc-UAV-main/models/group/__init__.py

@@ -0,0 +1,2 @@
+from .groupnet import GroupNet
+from .groupnet_dino import GroupDinoNet

GeoLoc-UAV-main/models/group/groupnet.py

@@ -0,0 +1,182 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from utils.utils import dim_extend,interpolate_feats,l2_normalize
+
+class GroupNetConfig:
+    def __init__(self):
+        self.sample_scale_begin = 0
+        self.sample_scale_inter = 0.5 
+        self.sample_scale_num = 1
+
+        self.sample_rotate_begin = 0
+        self.sample_rotate_inter = 45
+        self.sample_rotate_num = 8
+
+group_config = GroupNetConfig()
+
+class VanillaLightCNN(nn.Module):
+    def __init__(self):
+        super(VanillaLightCNN, self).__init__()
+        self.conv0=nn.Sequential(
+            nn.Conv2d(3,16,5,1,2,bias=False),
+            nn.InstanceNorm2d(16),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(16,32,5,1,2,bias=False),
+            nn.InstanceNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.AvgPool2d(2, 2),
+
+            # 修改
+            nn.Conv2d(32,64,5,1,2,bias=False),
+            nn.InstanceNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.AvgPool2d(2, 2),
+
+        )
+        # 原来 32
+        self.conv1=nn.Sequential(
+            nn.Conv2d(64,64,5,1,2,bias=False),
+            nn.InstanceNorm2d(64),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(64,64,5,1,2,bias=False),
+            nn.InstanceNorm2d(64),
+        ) 
+
+    def forward(self, x):
+        x=self.conv1(self.conv0(x))
+        x=l2_normalize(x,axis=1)  # [1,c,w//2, h//2]
+        return x
+
+class ExtractorWrapper(nn.Module):
+    def __init__(self,scale_num, rotation_num):
+        super(ExtractorWrapper, self).__init__()
+        self.extractor=VanillaLightCNN()
+        self.sn, self.rn = scale_num, rotation_num 
+
+    def forward(self,img_list,pts_list):
+        '''
+
+        :param img_list:  list of [b,3,h,w]
+        :param pts_list:  list of [b,n,2]
+        :return:gefeats [b,n,f,sn,rn]
+        '''
+        assert(len(img_list)==self.rn*self.sn)
+        gfeats_list,neg_gfeats_list=[],[]
+        # feature extraction
+        for img_index,img in enumerate(img_list):
+            # extract feature
+            feats=self.extractor(img)
+            gfeats_list.append(interpolate_feats(img,pts_list[img_index],feats)[:,:,:,None])
+            
+
+        gfeats_list=torch.cat(gfeats_list,3)  # b,n,f,sn*rn
+        b,n,f,_=gfeats_list.shape
+        gfeats_list=gfeats_list.reshape(b,n,f,self.sn,self.rn)
+        
+        return gfeats_list
+
+class BilinearGCNN(nn.Module):
+    def __init__(self, scale_num, rotation_num):
+        super(BilinearGCNN, self).__init__()
+
+        self.r, self.s = rotation_num, scale_num
+
+        self.network1_embed1 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network1_embed1_short = nn.Conv2d(64, 64, 1, 1)
+        self.network1_embed1_relu = nn.ReLU(True)
+
+        self.network1_embed2 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network1_embed2_short = nn.Conv2d(64, 64, 1, 1)
+        self.network1_embed2_relu = nn.ReLU(True)
+
+        self.network1_embed3 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 16, 3, 1, 1),
+        )
+
+        ###########################
+        self.network2_embed1 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network2_embed1_short = nn.Conv2d(64, 64, 1, 1)
+        self.network2_embed1_relu = nn.ReLU(True)
+
+        self.network2_embed2 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network2_embed2_short = nn.Conv2d(64, 64, 1, 1)
+        self.network2_embed2_relu = nn.ReLU(True)
+
+        self.network2_embed3 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 16, 3, 1, 1),
+        )
+
+    def forward(self, x):
+        '''
+
+        :param x:  b,n,f,ssn,srn
+        :return:
+        '''
+        b, n, f, ssn, srn = x.shape
+        assert (ssn == self.s and srn == self.r)
+        x = x.reshape(b * n, f, ssn, srn)
+
+        x1 = self.network1_embed1_relu(self.network1_embed1(x) + self.network1_embed1_short(x))
+        x1 = self.network1_embed2_relu(self.network1_embed2(x1) + self.network1_embed2_short(x1))
+        x1 = self.network1_embed3(x1)
+
+        x2 = self.network2_embed1_relu(self.network2_embed1(x) + self.network2_embed1_short(x))
+        x2 = self.network2_embed2_relu(self.network2_embed2(x2) + self.network2_embed2_short(x2))
+        x2 = self.network2_embed3(x2)
+
+        x1 = x1.reshape(b * n, 16, self.s * self.r)
+        x2 = x2.reshape(b * n, 16, self.s * self.r).permute(0, 2, 1)  # b*n,25,16
+        x = torch.bmm(x1, x2).reshape(b * n, 256)  # b*n,8,25
+        assert (x.shape[1] == 256)
+        x=x.reshape(b,n,256)
+        x=l2_normalize(x,axis=2)
+        return x
+
+class EmbedderWrapper(nn.Module):
+    def __init__(self, scale_num, rotation_num):
+        super(EmbedderWrapper, self).__init__()
+        self.embedder=BilinearGCNN(scale_num, rotation_num)
+
+    def forward(self, gfeats):
+        # group cnns
+        gefeats=self.embedder(gfeats) # b,n,f
+        return gefeats
+
+class GroupNet(nn.Module):
+    def __init__(self, config=group_config):
+        super(GroupNet, self).__init__()
+        self.scale_num = config.sample_scale_num
+        self.rotation_num = config.sample_rotate_num
+       
+
+        self.extractor=ExtractorWrapper(self.scale_num, self.rotation_num).cuda()
+        self.embedder=EmbedderWrapper(self.scale_num, self.rotation_num).cuda()
+
+    def forward(self, img_list, pts_list):
+        gfeats=self.extractor(dim_extend(img_list),dim_extend(pts_list))
+        efeats=self.embedder(gfeats)
+        return efeats, gfeats

GeoLoc-UAV-main/models/group/groupnet_dino.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from utils.utils import dim_extend,interpolate_feats,l2_normalize
+import json
+
+json_path = "/media/guan/新加卷/Code/Code/configs/transform_config.json"
+with open(json_path, 'r', encoding='utf-8') as file:
+    data = json.load(file)
+group_config = data["transform_config"]
+
+# class GroupNetConfig:
+#     def __init__(self):
+#         self.sample_scale_begin = 0
+#         self.sample_scale_inter = 0.5 
+#         self.sample_scale_num = 3
+
+#         self.sample_rotate_begin = -45
+#         self.sample_rotate_inter = 45
+#         self.sample_rotate_num = 8
+
+# class GroupNetConfig:
+#     def __init__(self):
+#         self.sample_scale_begin = 0
+#         self.sample_scale_inter = 1 
+#         self.sample_scale_num = 1
+
+#         self.sample_rotate_begin = 0
+#         self.sample_rotate_inter = 0
+#         self.sample_rotate_num = 1
+# group_config = GroupNetConfig()
+
+class VanillaLightCNN(nn.Module):
+    def __init__(self):
+        super(VanillaLightCNN, self).__init__()
+        self.conv0 = nn.Sequential(
+            nn.Conv2d(384,384//2,1,1,bias=False),
+            nn.InstanceNorm2d(384//2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384//2,384//4,1,1,bias=False),
+            nn.InstanceNorm2d(384//4),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384//4,64,1,1,bias=False),
+            nn.InstanceNorm2d(64),
+        )
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(3,16,5,1,2,bias=False),
+            nn.InstanceNorm2d(16),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(16,32,5,1,2,bias=False),
+            nn.InstanceNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.AvgPool2d(2, 2))
+        self.proj = nn.Conv2d(96, 64, 1, 1, bias=False)
+
+       
+
+    def forward(self, x, img):
+        x_dino=self.conv0(x)
+        x_resized = F.interpolate(img, size=(32, 32), mode='bilinear', align_corners=False)
+        x_cnn = self.conv1(x_resized)
+        x_cat = torch.concat((x_dino, x_cnn), dim=1)
+        x_proj = self.proj(x_cat)
+        x=l2_normalize(x_proj,axis=1)  # [1,c,w//2, h//2]
+        return x
+
+class ExtractorWrapper(nn.Module):
+    def __init__(self,scale_num, rotation_num):
+        super(ExtractorWrapper, self).__init__()
+        self.extractor=VanillaLightCNN()
+        self.sn, self.rn = scale_num, rotation_num 
+       
+        dinov2_weights = torch.hub.load('facebookresearch/dinov2', "dinov2_vits14")
+        # torch.load("/media/Shen/Data/RingoData/WorldLoc/Code/dinov2_vits14_pretrain.pth")
+        from models.transformer import vit_small
+        vit_kwargs = dict(
+            patch_size= 14,
+            img_size=518,
+            init_values = 1.0,
+            ffn_layer = "mlp",
+            block_chunks = 0,
+        )
+
+        self.dinov2_vits14 = vit_small(**vit_kwargs).eval()
+        # self.dinov2_vits14.load_state_dict(dinov2_weights)
+
+    def forward(self,img_list,pts_list):
+        '''
+
+        :param img_list:  list of [b,3,h,w]
+        :param pts_list:  list of [b,n,2]
+        :return:gefeats [b,n,f,sn,rn]
+        '''
+        assert(len(img_list)==self.rn*self.sn)
+        gfeats_list = []
+        # feature extraction
+        
+        for img_index,img in enumerate(img_list):
+            # extract feature
+            
+            with torch.no_grad():
+                dinov2_features_16 = self.dinov2_vits14.forward_features(img)
+                B, _, H, W = img.shape
+                features_16 = dinov2_features_16['x_norm_patchtokens'].permute(0,2,1).reshape(B,-1,H//14, W//14)
+                
+
+            feats=self.extractor(features_16, img)
+            gfeats_list.append(interpolate_feats(img, pts_list[img_index], feats)[:,:,:,None])
+            
+        gfeats_list=torch.cat(gfeats_list,3)  # b,n,f,sn*rn
+        b,n,f,_=gfeats_list.shape
+        gfeats_list=gfeats_list.reshape(b,n,f,self.sn,self.rn)
+        
+        return gfeats_list
+
+class BilinearGCNN(nn.Module):
+    def __init__(self, scale_num, rotation_num):
+        super(BilinearGCNN, self).__init__()
+
+        self.r, self.s = rotation_num, scale_num
+
+        self.network1_embed1 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network1_embed1_short = nn.Conv2d(64, 64, 1, 1)
+        self.network1_embed1_relu = nn.ReLU(True)
+
+        self.network1_embed2 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network1_embed2_short = nn.Conv2d(64, 64, 1, 1)
+        self.network1_embed2_relu = nn.ReLU(True)
+
+        self.network1_embed3 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 16, 3, 1, 1),
+        )
+
+        ###########################
+        self.network2_embed1 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network2_embed1_short = nn.Conv2d(64, 64, 1, 1)
+        self.network2_embed1_relu = nn.ReLU(True)
+
+        self.network2_embed2 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+        )
+        self.network2_embed2_short = nn.Conv2d(64, 64, 1, 1)
+        self.network2_embed2_relu = nn.ReLU(True)
+
+        self.network2_embed3 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 16, 3, 1, 1),
+        )
+
+    def forward(self, x):
+        '''
+
+        :param x:  b,n,f,ssn,srn
+        :return:
+        '''
+        
+        b, n, f, ssn, srn = x.shape
+        # equal = x.reshape(b, n, f, ssn*srn)
+        # equ_features=torch.max(equal,dim=-1,keepdim=False)[0]
+        # x = l2_normalize(equ_features, axis=1)
+        assert (ssn == self.s and srn == self.r)
+        x = x.reshape(b * n, f, ssn, srn)
+
+        x1 = self.network1_embed1_relu(self.network1_embed1(x) + self.network1_embed1_short(x))
+        x1 = self.network1_embed2_relu(self.network1_embed2(x1) + self.network1_embed2_short(x1))
+        x1 = self.network1_embed3(x1)
+
+        x2 = self.network2_embed1_relu(self.network2_embed1(x) + self.network2_embed1_short(x))
+        x2 = self.network2_embed2_relu(self.network2_embed2(x2) + self.network2_embed2_short(x2))
+        x2 = self.network2_embed3(x2)
+
+        x1 = x1.reshape(b * n, 16, self.s * self.r)
+        x2 = x2.reshape(b * n, 16, self.s * self.r).permute(0, 2, 1)  # b*n,25,16
+        x = torch.bmm(x1, x2).reshape(b * n, 256)  # b*n,8,25
+        assert (x.shape[1] == 256)
+        x=x.reshape(b,n,256)
+        x=l2_normalize(x,axis=2)
+        return x
+
+class EmbedderWrapper(nn.Module):
+    def __init__(self, scale_num, rotation_num):
+        super(EmbedderWrapper, self).__init__()
+        self.embedder=BilinearGCNN(scale_num, rotation_num)
+
+    def forward(self, gfeats):
+        # group cnns
+        gefeats=self.embedder(gfeats) # b,n,f
+        return gefeats
+
+class GroupDinoNet(nn.Module):
+    def __init__(self, config=group_config):
+        super(GroupDinoNet, self).__init__()
+        self.scale_num = config["sample_scale_num"]
+        self.rotation_num = config["sample_rotate_num"]
+       
+
+        self.extractor=ExtractorWrapper(self.scale_num, self.rotation_num).cuda()
+        self.embedder=EmbedderWrapper(self.scale_num, self.rotation_num).cuda()
+
+    def forward(self, img_list, pts_list):
+        gfeats=self.extractor(dim_extend(img_list),dim_extend(pts_list))
+        efeats=self.embedder(gfeats)
+        return efeats, gfeats

GeoLoc-UAV-main/models/helper.py

@@ -0,0 +1,90 @@
+from models import group
+from models import aggregators
+from models import backbone
+
+def get_groupnet(groupnet_arch='groupnet', group_config={}):
+    
+    if  "groupnet" in groupnet_arch.lower():
+        return group.GroupNet(**group_config)
+
+def get_groupdinonet(groupnet_arch='groupdinonet', group_config={}):
+    
+    if  "groupdinonet" in groupnet_arch.lower():
+        return group.GroupDinoNet(**group_config)
+    
+def get_aggregator(agg_arch='ConvAP', agg_config={}):
+    """Helper function that returns the aggregation layer given its name.
+    If you happen to make your own aggregator, you might need to add a call
+    to this helper function.
+
+    Args:
+        agg_arch (str, optional): the name of the aggregator. Defaults to 'ConvAP'.
+        agg_config (dict, optional): this must contain all the arguments needed to instantiate the aggregator class. Defaults to {}.
+
+    Returns:
+        nn.Module: the aggregation layer
+    """
+    
+    if 'cosplace' in agg_arch.lower():
+        assert 'in_dim' in agg_config
+        assert 'out_dim' in agg_config
+        return aggregators.CosPlace(**agg_config)
+
+    elif 'gem' in agg_arch.lower():
+        if agg_config == {}:
+            agg_config['p'] = 3
+        else:
+            assert 'p' in agg_config
+        return aggregators.GeMPool(**agg_config)
+    
+    elif 'multiconvap' in agg_arch.lower():
+        assert 'in_channels' in agg_config
+        return aggregators.MulConvAP(**agg_config)
+    
+    elif 'convap' in agg_arch.lower():
+        assert 'in_channels' in agg_config
+        return aggregators.ConvAP(**agg_config)
+    
+    
+    elif 'mixvpr' in agg_arch.lower():
+        assert 'in_channels' in agg_config
+        assert 'out_channels' in agg_config
+        assert 'in_h' in agg_config
+        assert 'in_w' in agg_config
+        assert 'mix_depth' in agg_config
+        return aggregators.MixVPR(**agg_config)
+
+    elif 'salad' in agg_arch.lower():
+        assert 'num_channels' in agg_config
+        assert 'num_clusters' in agg_config
+        assert 'cluster_dim' in agg_config
+        assert 'token_dim' in agg_config
+        return aggregators.SALAD(**agg_config)
+
+    elif 'netvlad' in agg_arch.lower():
+        return aggregators.NetVLAD()
+
+def get_backbone(backbone_arch='resnet50',
+                 pretrained=True,
+                 layers_to_freeze=2,
+                 layers_to_crop=[],
+                 pretrain_flag=False):
+    """Helper function that returns the backbone given its name
+
+    Args:
+        backbone_arch (str, optional): . Defaults to 'resnet50'.
+        pretrained (bool, optional): . Defaults to True.
+        layers_to_freeze (int, optional): . Defaults to 2.
+        layers_to_crop (list, optional): This is mostly used with ResNet where we sometimes need to crop the last residual block (ex. [4]). Defaults to [].
+
+    Returns:
+        model: the backbone as a nn.Model object
+    """
+    if 'resnet' in backbone_arch.lower():
+        return backbone.ResNet(backbone_arch, pretrained, layers_to_freeze, layers_to_crop, pretrain_flag)
+    
+    elif 'dinov2' in backbone_arch.lower():
+        return backbone.DINOv2(model_name=backbone_arch,  num_trainable_blocks=4,
+            norm_layer=True,
+            return_token=True,
+            pretrain_flag=pretrain_flag)

GeoLoc-UAV-main/models/model.py

@@ -0,0 +1,122 @@
+import numpy as np
+import torch.nn as nn
+import torch
+
+from models import helper
+
+class GrounpDinoGlobal(nn.Module):
+
+    def __init__(self, 
+                groupnet_arch,
+                agg_arch, 
+                agg_config ):
+    
+        super(GrounpDinoGlobal, self).__init__()
+
+        self.groupnet = helper.get_groupdinonet(groupnet_arch)
+        self.aggregator = helper.get_aggregator(agg_arch, agg_config)
+        self.logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    
+    def forward(self, x, pts_list):
+
+        local_feature, gfeats_lists = self.groupnet(x, pts_list)
+        local_feature = local_feature.permute(0,2,1).unsqueeze(-1)
+        global_feature = self.aggregator(local_feature)
+        
+
+
+        # img_num = len(x)
+        # bs = x[0][0].shape[0]
+
+        # global_feature = torch.zeros(bs*len(x), 256, device='cuda')
+        # for i in range(img_num):
+        #     imgs, pts = x[i], pts_list[i]
+        #     local_feature = self.groupnet(imgs, pts)
+        #     local_feature = local_feature.permute(0,2,1).unsqueeze(-1)
+        #     des = self.aggregator(local_feature)
+        #     for j in range(len(des)):
+        #         global_feature[j*img_num+i,:] = des[j,:]
+        
+        return global_feature, gfeats_lists
+
+
+class GrounpGlobal(nn.Module):
+
+    def __init__(self, 
+                groupnet_arch,
+                agg_arch, 
+                agg_config ):
+    
+        super(GrounpGlobal, self).__init__()
+
+        self.groupnet = helper.get_groupnet(groupnet_arch)
+        self.aggregator = helper.get_aggregator(agg_arch, agg_config)
+        self.logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    
+    def forward(self, x, pts_list):
+
+        local_feature, gfeats_lists = self.groupnet(x, pts_list)
+        local_feature = local_feature.permute(0,2,1).unsqueeze(-1)
+        global_feature = self.aggregator(local_feature)
+
+
+        # img_num = len(x)
+        # bs = x[0][0].shape[0]
+
+        # global_feature = torch.zeros(bs*len(x), 256, device='cuda')
+        # for i in range(img_num):
+        #     imgs, pts = x[i], pts_list[i]
+        #     local_feature = self.groupnet(imgs, pts)
+        #     local_feature = local_feature.permute(0,2,1).unsqueeze(-1)
+        #     des = self.aggregator(local_feature)
+        #     for j in range(len(des)):
+        #         global_feature[j*img_num+i,:] = des[j,:]
+        
+        return global_feature, gfeats_lists
+    
+class BackboneGlobal(nn.Module):
+
+    def __init__(self, 
+                backbone_arch,
+                pretrain_flag,
+                agg_arch, 
+                agg_config ):
+    
+        super(BackboneGlobal, self).__init__()
+
+        self.backbone = helper.get_backbone(backbone_arch, pretrain_flag)
+        self.aggregator = helper.get_aggregator(agg_arch, agg_config)
+        self.logit_scale = torch.nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        if 'dinov2' in backbone_arch.lower():
+            self.FLAG = True 
+        else:
+            self.FLAG = False
+
+    
+    def forward(self, x):
+
+        local_feature = self.backbone(x)
+        
+        # dinov2
+        
+        if self.FLAG:
+            global_feature = self.aggregator(local_feature[0])
+        else:
+            global_feature = self.aggregator(local_feature)
+
+
+        # img_num = len(x)
+        # bs = x[0][0].shape[0]
+
+        # global_feature = torch.zeros(bs*len(x), 256, device='cuda')
+        # for i in range(img_num):
+        #     imgs, pts = x[i], pts_list[i]
+        #     local_feature = self.groupnet(imgs, pts)
+        #     local_feature = local_feature.permute(0,2,1).unsqueeze(-1)
+        #     des = self.aggregator(local_feature)
+        #     for j in range(len(des)):
+        #         global_feature[j*img_num+i,:] = des[j,:]
+        
+        return global_feature

GeoLoc-UAV-main/models/trainer.py

@@ -0,0 +1,231 @@
+import time
+import torch
+from tqdm import tqdm
+from utils import setting
+from torch.cuda.amp import autocast
+import torch.nn.functional as F
+
+def train(train_config, model, dataloader, loss_function, optimizer,scheduler=None, scaler=None, writer=None):
+
+    # set model train mode
+    model.train()
+    
+    losses = setting.AverageMeter()
+    
+    # wait before starting progress bar
+    time.sleep(0.1)
+    
+    # Zero gradients for first step
+    optimizer.zero_grad(set_to_none=True)
+    
+    step = 1
+    
+    if train_config.verbose:
+        bar = tqdm(dataloader, total=len(dataloader))
+    else:
+        bar = dataloader
+    
+    # for loop over one epoch
+    # 修改代码为带weight
+    # for query,query_pt, reference,  reference_pt, ids, weight in bar:
+    for query,query_pt, reference,  reference_pt, ids in bar:   
+        if scaler:
+            with autocast():
+            
+                # data (batches) to device   
+                query = query
+                reference = reference
+                query_pt = query_pt
+                reference_pt = reference_pt
+            
+                # Forward pass
+                features1, _ = model(query, query_pt)
+                features2, _ = model(reference, reference_pt)
+               
+                if torch.cuda.device_count() > 1 and len(train_config.gpu_ids) > 1: 
+                    loss = loss_function(features1, features2, model.module.logit_scale.exp())
+                    # loss = loss_function(features1, features2, model.module.logit_scale.exp(), weight)
+                else:
+                    # InfoNCE Loss
+                    loss = loss_function(features1, features2, model.logit_scale.exp()) 
+                    # loss = loss_function(features1, features2, model.logit_scale.exp(), weight)
+                    # SupCon Loss
+                    # feature = torch.cat((features1, features2), dim=0)
+                    # labels = torch.cat((ids, ids), dim=0)
+                    # loss = loss_function(feature, labels)
+
+                losses.update(loss.item())
+                
+                  
+            scaler.scale(loss).backward()
+            
+            # Gradient clipping 
+            if train_config.clip_grad:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_value_(model.parameters(), train_config.clip_grad) 
+            
+            # Update model parameters (weights)
+            scaler.step(optimizer)
+            scaler.update()
+
+            # Zero gradients for next step
+            optimizer.zero_grad()
+            
+            # Scheduler
+            if train_config.scheduler == "polynomial" or train_config.scheduler == "cosine" or train_config.scheduler ==  "constant":
+                scheduler.step()
+   
+        else:
+        
+            # data (batches) to device   
+            query = query.to(train_config.device)
+            reference = reference.to(train_config.device)
+
+            # Forward pass
+            features1, features2 = model(query, reference)
+            
+            if torch.cuda.device_count() > 1 and len(train_config.gpu_ids) > 1: 
+                # loss = loss_function(features1, features2, model.module.logit_scale.exp(), weight)
+                loss = loss_function(features1, features2, model.module.logit_scale.exp())
+            else:
+                loss = loss_function(features1, features2, model.logit_scale.exp())
+                # loss = loss_function(features1, features2, model.logit_scale.exp(), weight)
+            losses.update(loss.item())
+            
+            # Calculate gradient using backward pass
+            loss.backward()
+
+            
+            
+            # Gradient clipping 
+            if train_config.clip_grad:
+                torch.nn.utils.clip_grad_value_(model.parameters(), train_config.clip_grad)                  
+            
+            # Update model parameters (weights)
+            optimizer.step()
+            # Zero gradients for next step
+            optimizer.zero_grad()
+            
+            # Scheduler
+            if train_config.scheduler == "polynomial" or train_config.scheduler == "cosine" or train_config.scheduler ==  "constant":
+                scheduler.step()
+        
+        
+        
+        if train_config.verbose:
+            
+            monitor = {"loss": "{:.4f}".format(loss.item()),
+                       "loss_avg": "{:.4f}".format(losses.avg),
+                       "lr" : "{:.6f}".format(optimizer.param_groups[0]['lr'])}
+            
+            bar.set_postfix(ordered_dict=monitor)
+
+            writer.add_scalar('Loss/train', loss.item(), step)
+            writer.add_scalar('Loss/avg_loss', losses.avg, step)
+            writer.add_scalar('lr', optimizer.param_groups[0]['lr'], step)
+
+        
+        step += 1
+
+    if train_config.verbose:
+        bar.close()
+
+    return losses.avg
+
+
+def train_backbone(train_config, model, dataloader, loss_function, optimizer, scheduler=None, scaler=None, writer=None, LPN=False):
+
+   # set model train mode
+    model.train()
+    
+    losses = setting.AverageMeter()
+    
+    # wait before starting progress bar
+    time.sleep(0.1)
+    
+    # Zero gradients for first step
+    optimizer.zero_grad(set_to_none=True)
+    
+    step = 1
+    
+    if train_config.verbose:
+        bar = tqdm(dataloader, total=len(dataloader))
+    else:
+        bar = dataloader
+    
+    
+    
+    # for loop over one epoch
+    for query, reference,  ids in bar:
+        
+        loss = 0.0
+        query = query.to(train_config.device)
+        reference = reference.to(train_config.device)
+        
+            # Forward pass
+        features1 = model(query)
+        features2 = model(reference)
+        
+        if LPN == False:
+            if torch.cuda.device_count() > 1 and len(train_config.gpu_ids) > 1: 
+                loss = loss_function(features1, features2, model.module.logit_scale.exp())
+            else:
+                loss = loss_function(features1, features2, model.logit_scale.exp()) 
+        else:
+            for index in range(len(features1)):
+                feature1_one = features1[index]
+                feature2_one = features2[index]
+                if torch.cuda.device_count() > 1 and len(train_config.gpu_ids) > 1: 
+                    temp_loss = loss_function(feature1_one, feature2_one, model.module.logit_scale.exp())
+                else:
+                    temp_loss = loss_function(feature1_one, feature2_one, model.logit_scale.exp()) 
+                loss += temp_loss
+        
+        losses.update(loss.item())
+                
+                  
+      
+
+        # Zero gradients for next step
+        optimizer.zero_grad()
+        
+        # Scheduler
+        if train_config.scheduler == "polynomial" or train_config.scheduler == "cosine" or train_config.scheduler ==  "constant":
+            scheduler.step()
+
+
+        losses.update(loss.item())
+
+        # Calculate gradient using backward pass
+        loss.backward(retain_graph=True)
+            
+                 
+            
+        # Update model parameters (weights)
+        optimizer.step()
+        # Zero gradients for next step
+        optimizer.zero_grad()
+        
+        # Scheduler
+        if train_config.scheduler == "polynomial" or train_config.scheduler == "cosine" or train_config.scheduler ==  "constant":
+            scheduler.step()
+        
+        
+        
+        if train_config.verbose:
+            
+            monitor = {"loss": "{:.4f}".format(loss.item()),
+                       "loss_avg": "{:.4f}".format(losses.avg),
+                       "lr" : "{:.6f}".format(optimizer.param_groups[0]['lr'])}
+            
+            bar.set_postfix(ordered_dict=monitor)
+            writer.add_scalar('Loss/train', loss.item(), step)
+            writer.add_scalar('Loss/avg_loss', losses.avg, step)
+            writer.add_scalar('lr', optimizer.param_groups[0]['lr'], step)
+        
+        step += 1
+
+    if train_config.verbose:
+        bar.close()
+
+    return losses.avg

GeoLoc-UAV-main/models/transformer/__init__.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from romatch.utils.utils import get_grid, get_autocast_params
+from .layers.block import Block
+from .layers.attention import MemEffAttention
+from .dinov2 import vit_large, vit_small
+
+class TransformerDecoder(nn.Module):
+    def __init__(self, blocks, hidden_dim, out_dim, is_classifier = False, *args, 
+                 amp = False, pos_enc = True, learned_embeddings = False, embedding_dim = None, amp_dtype = torch.float16, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.blocks = blocks
+        self.to_out = nn.Linear(hidden_dim, out_dim)
+        self.hidden_dim = hidden_dim
+        self.out_dim = out_dim
+        self._scales = [16]
+        self.is_classifier = is_classifier
+        self.amp = amp
+        self.amp_dtype = amp_dtype
+        self.pos_enc = pos_enc
+        self.learned_embeddings = learned_embeddings
+        if self.learned_embeddings:
+            self.learned_pos_embeddings = nn.Parameter(nn.init.kaiming_normal_(torch.empty((1, hidden_dim, embedding_dim, embedding_dim))))
+
+    def scales(self):
+        return self._scales.copy()
+
+    def forward(self, gp_posterior, features, old_stuff, new_scale):
+        autocast_device, autocast_enabled, autocast_dtype = get_autocast_params(gp_posterior.device, enabled=self.amp, dtype=self.amp_dtype)
+        with torch.autocast(autocast_device, enabled=autocast_enabled, dtype = autocast_dtype):
+            B,C,H,W = gp_posterior.shape
+            x = torch.cat((gp_posterior, features), dim = 1)
+            B,C,H,W = x.shape
+            grid = get_grid(B, H, W, x.device).reshape(B,H*W,2)
+            if self.learned_embeddings:
+                pos_enc = F.interpolate(self.learned_pos_embeddings, size = (H,W), mode = 'bilinear', align_corners = False).permute(0,2,3,1).reshape(1,H*W,C)
+            else:
+                pos_enc = 0
+            tokens = x.reshape(B,C,H*W).permute(0,2,1) + pos_enc
+            z = self.blocks(tokens)
+            out = self.to_out(z)
+            out = out.permute(0,2,1).reshape(B, self.out_dim, H, W)
+            warp, certainty = out[:, :-1], out[:, -1:]
+            return warp, certainty, None
+
+

GeoLoc-UAV-main/models/transformer/dinov2.py

@@ -0,0 +1,360 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+
+from .layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        ffn_bias=True,
+        proj_bias=True,
+        drop_path_rate=0.0,
+        drop_path_uniform=False,
+        init_values=None,  # for layerscale: None or 0 => no layerscale
+        embed_layer=PatchEmbed,
+        act_layer=nn.GELU,
+        block_fn=Block,
+        ffn_layer="mlp",
+        block_chunks=1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+
+        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+        self.init_weights()
+        for param in self.parameters():
+            param.requires_grad = False
+    
+    @property
+    def device(self):
+        return self.cls_token.device
+
+    def init_weights(self):
+        trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.normal_(self.cls_token, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode="bicubic",
+        )
+
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_patchtokens": x_norm[:, 1:],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        # import pdb;pdb.set_trace()
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_patchtokens": x_norm[:, 1:],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm=True,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1:] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        **kwargs,
+    )
+    return model

GeoLoc-UAV-main/models/transformer/layers/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .dino_head import DINOHead
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention

GeoLoc-UAV-main/models/transformer/layers/attention.py

@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+
+from torch import Tensor
+from torch import nn
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import memory_efficient_attention, unbind, fmha
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    logger.warning("xFormers not available")
+    XFORMERS_AVAILABLE = False
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            assert attn_bias is None, "xFormers is required for nested tensors usage"
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

GeoLoc-UAV-main/models/transformer/layers/block.py

@@ -0,0 +1,252 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+from typing import Callable, List, Any, Tuple, Dict
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import fmha
+    from xformers.ops import scaled_index_add, index_select_cat
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    logger.warning("xFormers not available")
+    XFORMERS_AVAILABLE = False
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: Tensor) -> Tensor:
+        def attn_residual_func(x: Tensor) -> Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: Tensor) -> Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: Tensor,
+    residual_func: Callable[[Tensor], Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[Tensor],
+    residual_func: Callable[[Tensor, Any], Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError

GeoLoc-UAV-main/models/transformer/layers/dino_head.py

@@ -0,0 +1,59 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+from torch.nn.init import trunc_normal_
+from torch.nn.utils import weight_norm
+
+
+class DINOHead(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        use_bn=False,
+        nlayers=3,
+        hidden_dim=2048,
+        bottleneck_dim=256,
+        mlp_bias=True,
+    ):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias)
+        self.apply(self._init_weights)
+        self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        self.last_layer.weight_g.data.fill_(1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        eps = 1e-6 if x.dtype == torch.float16 else 1e-12
+        x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)
+        x = self.last_layer(x)
+        return x
+
+
+def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True):
+    if nlayers == 1:
+        return nn.Linear(in_dim, bottleneck_dim, bias=bias)
+    else:
+        layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]
+        if use_bn:
+            layers.append(nn.BatchNorm1d(hidden_dim))
+        layers.append(nn.GELU())
+        for _ in range(nlayers - 2):
+            layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+        layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))
+        return nn.Sequential(*layers)

GeoLoc-UAV-main/models/transformer/layers/drop_path.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

GeoLoc-UAV-main/models/transformer/layers/layer_scale.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

GeoLoc-UAV-main/models/transformer/layers/mlp.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

GeoLoc-UAV-main/models/transformer/layers/patch_embed.py

@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+from torch import Tensor
+import torch.nn as nn
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops

GeoLoc-UAV-main/models/transformer/layers/swiglu_ffn.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+try:
+    from xformers.ops import SwiGLU
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    SwiGLU = SwiGLUFFN
+    XFORMERS_AVAILABLE = False
+
+
+class SwiGLUFFNFused(SwiGLU):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )

GeoLoc-UAV-main/preprocess_data.py

@@ -0,0 +1,89 @@
+import os
+import json
+from glob import glob
+
+"""
+根据输入的.txt，生成相应的train_query.txt, train_db.txt
+"""
+
+def generate_img_list(root ,txt,  save_path, mode):
+
+    save_file_path_query = save_path + mode + '_query_all.txt'
+    save_file_path_db = save_path + mode + '_db_all.txt'
+    
+    label = 0
+     
+     # 处理query图像
+    with open(save_file_path_query, 'w') as f_query:
+        with open(txt, 'r') as f:
+            for line in f:
+                one_root = os.path.join(root, line.strip('\n'))
+                positive = json.load(open(one_root+'/positive.json'))
+                semi_positive = json.load(open(one_root+'/semi_positive.json'))
+                query_names = positive.keys()
+                query_dirs = os.listdir(one_root+'/query')
+                for query_name in query_names:
+                    for query_dir in query_dirs:
+                        # query路径
+                        if query_dir[0] != 'h':
+                            continue
+                        one_query_path = line.strip('\n') + '/query/' + query_dir + '/'  + 'footage/'+ query_dir + '_' + query_name + '.jpeg'
+                        f_query.write(one_query_path + ' ' + str(label) + ' ')
+                        pos_dbs = positive[query_name]
+                        try:
+                            FLAG = True
+                            semi_pos_dbs = semi_positive[query_name]
+                        except:
+                            FLAG = False
+                        # pos GT 路径
+                        for pos_db in pos_dbs:
+                            temp = line.strip('\n') + '/DB/' + 'img/' + pos_db
+                            f_query.write(temp + ' ')
+                        # semi GT 路径
+                        if FLAG:
+                            for semi_pos_db in semi_pos_dbs:
+                                temp = line.strip('\n') + '/DB/' + 'img/' + semi_pos_db
+                                f_query.write(temp + ' ')
+                        f_query.write('\n')
+                    label += 1
+    print('-------------------------finish-------------------------')
+
+    # 处理DB图像
+    with open(save_file_path_db, 'w') as f_db:
+        with open(txt, 'r') as f:
+            for line in f:
+                one_root = os.path.join(root, line.strip('\n'))
+                db_path = one_root + '/DB/' + 'img/'
+                db_imgs = glob(db_path + '*.png')
+                for db_img in db_imgs:
+                    temp_list = db_img.split('/')[6:]
+                    temp = ''
+                    for i in temp_list:
+                        temp += i
+                        if not i.endswith('.png'):
+                         temp += '/'
+                    f_db.write(temp + '\n')
+                  
+
+
+root = "/media/Shen/Data/RingoData/WorldLoc/"
+txt = "/media/Shen/Data/RingoData/WorldLoc/Index/train_all.txt"
+save_path = "/media/Shen/Data/RingoData/WorldLoc/Index/"
+mode = 'train'
+
+generate_img_list(root, txt, save_path, mode)
+
+
+# 验证代码
+# import cv2
+# txt = '/media/Shen/Data/RingoData/WorldLoc/Index/train_db.txt'
+# with open(txt, 'r') as f:
+#     for line in f:
+#         line_list = line.split(' ')
+#         query_img = line_list[0].strip('\n')
+#         root = '/media/Shen/Data/RingoData/WorldLoc'
+#         img_path = os.path.join(root, query_img)
+#         img = cv2.imread(img_path)
+#         print(img.shape)
+#         cv2.imshow('img', img)
+#         cv2.waitKey(0)

GeoLoc-UAV-main/train_group.py

@@ -0,0 +1,348 @@
+import os
+import time
+import numpy as np
+import math
+import shutil
+import sys
+import torch
+from dataclasses import dataclass,field
+from torch.cuda.amp import GradScaler
+from torch.utils.data import DataLoader
+from transformers import get_constant_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup, get_cosine_schedule_with_warmup
+from torchvision import transforms as T
+from torch.utils.tensorboard import SummaryWriter
+
+from dataset.World import WorldDatasetTrainGroup, WorldDatasetEvalGroup
+from models import model,trainer
+from utils import setting
+from utils import loss
+from eval import eval
+
+def default_group_config():
+
+    return {
+          "group_arch" : "groupnet", #group
+        "group_config": {
+            "none"
+        }
+    }
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "resnet18", 
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "multiconvap", #convap
+    "agg_config": {
+      "in_channels": 256, #256 #512
+      "out_channels": 256, #256
+      "s1": 1,
+      "s2": 1,
+      'LPN':False
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "group34"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world"
+
+    # model config
+    group:dict = field(default_factory=default_group_config)
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc"
+    train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/Index/train_query.txt"
+
+    # val_index
+    val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/val.txt"
+
+    # test_index
+    test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/test.txt"
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 if os.name == 'nt' else 4 
+    
+    # train on GPU if available
+    device: str = 'cuda:1' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 30
+    batch_size: int = 128                # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (1,2,3)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+if __name__ == '__main__':
+
+    model_path = "{}/{}/{}".format(config.model_path,
+                                       config.model,
+                                       time.strftime("%H%M%S"))
+
+    if not os.path.exists(model_path):
+        os.makedirs(model_path)
+    shutil.copyfile(os.path.basename(__file__), "{}/train.py".format(model_path))
+
+    # Redirect print to both console and log file
+    sys.stdout = setting.Logger(os.path.join(model_path, 'log.txt'))
+
+    setting.setup_system(seed=config.seed,
+                 cudnn_benchmark=config.cudnn_benchmark,
+                 cudnn_deterministic=config.cudnn_deterministic)
+    
+    #-----------------------------------------------------------------------------#
+    # Model                                                                       #
+    #-----------------------------------------------------------------------------#
+        
+    print("\nModel: {}".format(config.model))
+
+  
+    # backbone
+    model = model.GrounpGlobal(config.group['group_arch'],
+                                config.agg['agg_arch'],
+                                config.agg['agg_config'])
+    
+     # Load pretrained Checkpoint    
+    if config.checkpoint_start is not None:  
+        print("Start from:", config.checkpoint_start)
+        model_state_dict = torch.load(config.checkpoint_start)  
+        model.load_state_dict(model_state_dict, strict=False)   
+    
+       # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+        model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+
+      # Model to device   
+    model = model.to(config.device)
+
+    #------------------------setting dataset-------------------------------------------------#
+    IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                     'std': [0.229, 0.224, 0.225]}
+    train_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.RandAugment(num_ops=3, interpolation=T.InterpolationMode.BILINEAR),
+            T.AugMix(),
+            # T.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1,
+            #                         hue=0),
+            # T.RandomGrayscale(p=0.2),
+            # T.RandomPosterize(p=0.2, bits=4),
+            # T.GaussianBlur(kernel_size=(1, 5), sigma=(0.1, 5)),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+    
+    eval_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+
+    #-----------------------------------------------------------------------------#
+    # DataLoader                                                                  #
+    #-----------------------------------------------------------------------------#
+
+    train_dataset = WorldDatasetTrainGroup(data_dir=config.dataset_root_dir,
+                                      query_txt=config.train_query_txt,
+                                      transforms_query=train_transform,
+                                      transforms_db=train_transform,
+                                      shuffle_batch_size=config.batch_size)
+    
+    # train_dataloader = DataLoader(train_dataset,
+    #                               batch_size=config.batch_size,
+    #                               num_workers=config.num_workers,
+    #                               shuffle=not config.custom_sampling,
+    #                               pin_memory=True)
+    
+    train_dataloader = DataLoader(train_dataset,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=config.custom_sampling,
+                                  pin_memory=True)
+    
+    #-----------------------------------------------------------------------------#
+    # Loss                                                                        #
+    #-----------------------------------------------------------------------------#
+
+    # InfoNCE loss
+    loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
+    loss_function = loss.InfoNCE(loss_function=loss_fn,
+                            device=config.device,
+                            )
+    # Supervised Contrastive loss
+    # loss_function = loss.SupervisedContrastiveLoss(temperature = 0.07, device=config.device)
+
+    if config.mixed_precision:
+        scaler = GradScaler(init_scale=2.**10)
+    else:
+        scaler = None
+    
+    #-----------------------------------------------------------------------------#
+    # optimizer                                                                   #
+    #-----------------------------------------------------------------------------#
+        
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
+
+        #-----------------------------------------------------------------------------#
+    # Scheduler                                                                   #
+    #-----------------------------------------------------------------------------#
+
+    train_steps = len(train_dataloader) * config.epochs
+    warmup_steps = len(train_dataloader) * config.warmup_epochs
+       
+    if config.scheduler == "polynomial":
+        print("\nScheduler: polynomial - max LR: {} - end LR: {}".format(config.lr, config.lr_end))  
+        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer,
+                                                              num_training_steps=train_steps,
+                                                              lr_end = config.lr_end,
+                                                              power=1.5,
+                                                              num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "cosine":
+        print("\nScheduler: cosine - max LR: {}".format(config.lr))   
+        scheduler = get_cosine_schedule_with_warmup(optimizer,
+                                                    num_training_steps=train_steps,
+                                                    num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "constant":
+        print("\nScheduler: constant - max LR: {}".format(config.lr))   
+        scheduler =  get_constant_schedule_with_warmup(optimizer,
+                                                       num_warmup_steps=warmup_steps)
+           
+    else:
+        scheduler = None
+        
+    print("Warmup Epochs: {} - Warmup Steps: {}".format(str(config.warmup_epochs).ljust(2), warmup_steps))
+    print("Train Epochs:  {} - Train Steps:  {}".format(config.epochs, train_steps))
+
+
+    #-----------------------------------------------------------------------------#
+    # Shuffle                                                                     #
+    #-----------------------------------------------------------------------------#            
+    if config.custom_sampling:
+        train_dataloader.dataset.shuffle()
+    
+    #-----------------------------------------------------------------------------#
+    # Train                                                                       #
+    #-----------------------------------------------------------------------------#
+    start_epoch = 0   
+    best_score = 0
+
+    #-----------------------------------------------------------------------------#
+    # Writer
+    #-----------------------------------------------------------------------------#
+    # Writer
+    writer = SummaryWriter('world/' + config.model)
+    LPN_flag = config.agg['agg_config']['LPN']
+    
+
+    for epoch in range(1, config.epochs+1):
+        
+        print("\n{}[Epoch: {}]{}".format(30*"-", epoch, 30*"-"))
+        
+
+        train_loss = trainer.train(config,
+                           model,
+                           dataloader=train_dataloader,
+                           loss_function=loss_function,
+                           optimizer=optimizer,
+                           scheduler=scheduler,
+                           scaler=scaler,
+                           writer=writer)
+        
+        print("Epoch: {}, Train Loss = {:.3f}, Lr = {:.6f}".format(epoch,
+                                                                   train_loss,
+                                                                   optimizer.param_groups[0]['lr']))
+        
+        #------------------------------------------------------------Eval---------------------------------------------------------------------#
+        result_list = []
+        with open(config.val_index_txt,"r") as val_test:
+            for line in val_test:
+                eva_dataset_query = WorldDatasetEvalGroup(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='query',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalGroup(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='DB',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result,_ , _ = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='group', LPN=False)
+                print(line.strip('\n'), 'top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2))
+                result_list.append(result)
+                writer.add_scalar(line.strip('\n'), round(result[0]*100,2), epoch)
+                
+        
+        result_array = np.array(result_list)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        writer.add_scalar('Average/top1', round(average_result[0]*100,2), epoch)
+        writer.add_scalar('Average/top5', round(average_result[1]*100,2), epoch)
+        
+    #------------------------------------------------------------Save---------------------------------------------------------------------#
+        if average_result[0] > best_score:
+
+                best_score = average_result[0]
+
+                if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+                    torch.save(model.module.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                else:
+                    torch.save(model.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                
+
+        if config.custom_sampling:
+            train_dataloader.dataset.shuffle()

GeoLoc-UAV-main/train_group_dino.py

@@ -0,0 +1,342 @@
+import os
+import time
+import numpy as np
+import math
+import shutil
+import sys
+import torch
+from dataclasses import dataclass,field
+from torch.cuda.amp import GradScaler
+from torch.utils.data import DataLoader
+from transformers import get_constant_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup, get_cosine_schedule_with_warmup
+from torchvision import transforms as T
+from torch.utils.tensorboard import SummaryWriter
+
+from dataset.World import WorldDatasetTrainGroup, WorldDatasetEvalGroup
+from models import model,trainer
+from utils import setting
+from utils import loss
+from eval import eval
+
+def default_group_config():
+
+    return {
+          "group_arch" : "groupdinonet", #group
+        "group_config": {
+            "none"
+        }
+    }
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "groupdino", 
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "multiconvap", #convap
+    "agg_config": {
+      "in_channels": 256, #256 #512
+      "out_channels": 256, #256
+      "s1": 1,
+      "s2": 1,
+      'LPN':False
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "groupdino-new-city-s3r4"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world"
+
+    # model config
+    group:dict = field(default_factory=default_group_config)
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc"
+    train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/Index/train_query_country.txt"
+
+    # val_index
+    val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/val_country.txt"
+
+    # test_index
+    test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/test_country.txt"
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 #if os.name == 'nt' else 4 
+    
+    # train on GPU if available
+    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 10
+    batch_size: int = 128              # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (0,2,3)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+if __name__ == '__main__':
+
+    model_path = "{}/{}/{}".format(config.model_path,
+                                       config.model,
+                                       time.strftime("%H%M%S"))
+
+    if not os.path.exists(model_path):
+        os.makedirs(model_path)
+    shutil.copyfile(os.path.basename(__file__), "{}/train.py".format(model_path))
+
+    # Redirect print to both console and log file
+    sys.stdout = setting.Logger(os.path.join(model_path, 'log.txt'))
+
+    setting.setup_system(seed=config.seed,
+                 cudnn_benchmark=config.cudnn_benchmark,
+                 cudnn_deterministic=config.cudnn_deterministic)
+    
+    #-----------------------------------------------------------------------------#
+    # Model                                                                       #
+    #-----------------------------------------------------------------------------#
+        
+    print("\nModel: {}".format(config.model))
+
+  
+    # backbone
+    model = model.GrounpDinoGlobal(config.group['group_arch'],
+                               config.agg['agg_arch'],
+                               config.agg['agg_config'])
+    
+     # Load pretrained Checkpoint    
+    if config.checkpoint_start is not None:  
+        print("Start from:", config.checkpoint_start)
+        model_state_dict = torch.load(config.checkpoint_start)  
+        model.load_state_dict(model_state_dict, strict=False)   
+    
+       # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+        model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+
+      # Model to device   
+    model = model.to(config.device)
+
+    #------------------------setting dataset-------------------------------------------------#
+    IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                     'std': [0.229, 0.224, 0.225]}
+    train_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.RandAugment(num_ops=3, interpolation=T.InterpolationMode.BILINEAR),
+            T.AugMix(),
+            # T.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1,
+            #                         hue=0),
+            # T.RandomGrayscale(p=0.2),
+            # T.RandomPosterize(p=0.2, bits=4),
+            # T.GaussianBlur(kernel_size=(1, 5), sigma=(0.1, 5)),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+    
+    eval_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+
+    #-----------------------------------------------------------------------------#
+    # DataLoader                                                                  #
+    #-----------------------------------------------------------------------------#
+
+    train_dataset = WorldDatasetTrainGroup(data_dir=config.dataset_root_dir,
+                                      query_txt=config.train_query_txt,
+                                      transforms_query=train_transform,
+                                      transforms_db=train_transform,
+                                      shuffle_batch_size=config.batch_size)
+    
+    
+    train_dataloader = DataLoader(train_dataset,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=config.custom_sampling,
+                                  pin_memory=True)
+    
+    #-----------------------------------------------------------------------------#
+    # Loss                                                                        #
+    #-----------------------------------------------------------------------------#
+
+    # InfoNCE loss
+    loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
+    loss_function = loss.InfoNCE(loss_function=loss_fn,
+                            device=config.device,
+                            )
+    # Supervised Contrastive loss
+    # loss_function = loss.SupervisedContrastiveLoss(temperature = 0.07, device=config.device)
+
+    if config.mixed_precision:
+        scaler = GradScaler(init_scale=2.**10)
+    else:
+        scaler = None
+    
+    #-----------------------------------------------------------------------------#
+    # optimizer                                                                   #
+    #-----------------------------------------------------------------------------#
+        
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
+
+        #-----------------------------------------------------------------------------#
+    # Scheduler                                                                   #
+    #-----------------------------------------------------------------------------#
+
+    train_steps = len(train_dataloader) * config.epochs
+    warmup_steps = len(train_dataloader) * config.warmup_epochs
+       
+    if config.scheduler == "polynomial":
+        print("\nScheduler: polynomial - max LR: {} - end LR: {}".format(config.lr, config.lr_end))  
+        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer,
+                                                              num_training_steps=train_steps,
+                                                              lr_end = config.lr_end,
+                                                              power=1.5,
+                                                              num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "cosine":
+        print("\nScheduler: cosine - max LR: {}".format(config.lr))   
+        scheduler = get_cosine_schedule_with_warmup(optimizer,
+                                                    num_training_steps=train_steps,
+                                                    num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "constant":
+        print("\nScheduler: constant - max LR: {}".format(config.lr))   
+        scheduler =  get_constant_schedule_with_warmup(optimizer,
+                                                       num_warmup_steps=warmup_steps)
+           
+    else:
+        scheduler = None
+        
+    print("Warmup Epochs: {} - Warmup Steps: {}".format(str(config.warmup_epochs).ljust(2), warmup_steps))
+    print("Train Epochs:  {} - Train Steps:  {}".format(config.epochs, train_steps))
+
+
+    #-----------------------------------------------------------------------------#
+    # Shuffle                                                                     #
+    #-----------------------------------------------------------------------------#            
+    if config.custom_sampling:
+        train_dataloader.dataset.shuffle()
+    
+    #-----------------------------------------------------------------------------#
+    # Train                                                                       #
+    #-----------------------------------------------------------------------------#
+    start_epoch = 0   
+    best_score = 0
+
+    #-----------------------------------------------------------------------------#
+    # Writer
+    #-----------------------------------------------------------------------------#
+    # Writer
+    writer = SummaryWriter('world/' + config.model)
+    
+
+    for epoch in range(1, config.epochs+1):
+        
+        print("\n{}[Epoch: {}]{}".format(30*"-", epoch, 30*"-"))
+        
+
+        train_loss = trainer.train(config,
+                           model,
+                           dataloader=train_dataloader,
+                           loss_function=loss_function,
+                           optimizer=optimizer,
+                           scheduler=scheduler,
+                           scaler=scaler,
+                           writer=writer)
+        
+        print("Epoch: {}, Train Loss = {:.3f}, Lr = {:.6f}".format(epoch,
+                                                                   train_loss,
+                                                                   optimizer.param_groups[0]['lr']))
+        
+        #------------------------------------------------------------Eval---------------------------------------------------------------------#
+        result_list = []
+        with open(config.val_index_txt,"r") as val_test:
+            for line in val_test:
+                eva_dataset_query = WorldDatasetEvalGroup(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='query',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalGroup(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='DB',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result,_ , _ = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='group', LPN=False)
+                print(line.strip('\n'), 'top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2))
+                result_list.append(result)
+                writer.add_scalar(line.strip('\n'), round(result[0]*100,2), epoch)
+                
+        
+        result_array = np.array(result_list)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        writer.add_scalar('Average/top1', round(average_result[0]*100,2), epoch)
+        writer.add_scalar('Average/top5', round(average_result[1]*100,2), epoch)
+        
+    #------------------------------------------------------------Save---------------------------------------------------------------------#
+        if average_result[0] > best_score:
+
+                best_score = average_result[0]
+
+                if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+                    torch.save(model.module.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                else:
+                    torch.save(model.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                
+
+        if config.custom_sampling:
+            train_dataloader.dataset.shuffle()

GeoLoc-UAV-main/train_vanilia.py

@@ -0,0 +1,340 @@
+import os
+import time
+import numpy as np
+import math
+import shutil
+import sys
+import torch
+from dataclasses import dataclass,field
+from torch.cuda.amp import GradScaler
+from torch.utils.data import DataLoader
+from transformers import get_constant_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup, get_cosine_schedule_with_warmup
+from torchvision import transforms as T
+from torch.utils.tensorboard import SummaryWriter
+
+from dataset.World import WorldDatasetTrainVanilia, WorldDatasetEvalVanilia
+from models import model,trainer
+from utils import setting
+from utils import loss
+from eval import eval
+
+
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "resnet18", 
+          "pretrain_flag":True
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "multiconvap", #convap
+    "agg_config": {
+      "in_channels": 512, #256 #512
+      "out_channels": 512, #256
+      "s1": 1,
+      "s2": 1,
+      'LPN':False
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "resnet-new-all-frozen"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world_vanilia"
+
+    # model config
+    backbone:dict = field(default_factory=default_backbone_config)
+
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc"
+    train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/Index/train_query_all.txt"
+
+    # val_index
+    val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/val_all.txt"
+
+    # test_index
+    test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/test_country.txt"
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 if os.name == 'nt' else 4
+    
+    # train on GPU if available
+    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 10
+    batch_size: int = 128                # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (0,2,3)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+if __name__ == '__main__':
+
+    model_path = "{}/{}/{}".format(config.model_path,
+                                       config.model,
+                                       time.strftime("%H%M%S"))
+
+    if not os.path.exists(model_path):
+        os.makedirs(model_path)
+    shutil.copyfile(os.path.basename(__file__), "{}/train.py".format(model_path))
+
+    # Redirect print to both console and log file
+    sys.stdout = setting.Logger(os.path.join(model_path, 'log.txt'))
+
+    setting.setup_system(seed=config.seed,
+                 cudnn_benchmark=config.cudnn_benchmark,
+                 cudnn_deterministic=config.cudnn_deterministic)
+    
+    #-----------------------------------------------------------------------------#
+    # Model                                                                       #
+    #-----------------------------------------------------------------------------#
+        
+    print("\nModel: {}".format(config.model))
+
+  
+    # backbone
+    model = model.BackboneGlobal(config.backbone['backbone_arch'],
+                                 config.backbone['pretrain_flag'],
+                               config.agg['agg_arch'],
+                               config.agg['agg_config'])
+    
+     # Load pretrained Checkpoint    
+    if config.checkpoint_start is not None:  
+        print("Start from:", config.checkpoint_start)
+        model_state_dict = torch.load(config.checkpoint_start)  
+        model.load_state_dict(model_state_dict, strict=False)   
+    
+       # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+        model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+
+      # Model to device   
+    model = model.to(config.device)
+
+    #------------------------setting dataset-------------------------------------------------#
+    IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                     'std': [0.229, 0.224, 0.225]}
+    train_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.RandAugment(num_ops=3, interpolation=T.InterpolationMode.BILINEAR),
+            T.AugMix(),
+            # T.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1,
+            #                         hue=0),
+            # T.RandomGrayscale(p=0.2),
+            # T.RandomPosterize(p=0.2, bits=4),
+            # T.GaussianBlur(kernel_size=(1, 5), sigma=(0.1, 5)),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+    
+    eval_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+
+    #-----------------------------------------------------------------------------#
+    # DataLoader                                                                  #
+    #-----------------------------------------------------------------------------#
+
+    train_dataset = WorldDatasetTrainVanilia(data_dir=config.dataset_root_dir,
+                                      query_txt=config.train_query_txt,
+                                      transforms_query=train_transform,
+                                      transforms_db=train_transform,
+                                      shuffle_batch_size=config.batch_size)
+    
+    
+    train_dataloader = DataLoader(train_dataset,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=config.custom_sampling,
+                                  pin_memory=True)
+    
+    #-----------------------------------------------------------------------------#
+    # Loss                                                                        #
+    #-----------------------------------------------------------------------------#
+
+    # InfoNCE loss
+    loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
+    loss_function = loss.InfoNCE(loss_function=loss_fn,
+                            device=config.device,
+                            )
+    # Supervised Contrastive loss
+    # loss_function = loss.SupervisedContrastiveLoss(temperature = 0.07, device=config.device)
+
+    if config.mixed_precision:
+        scaler = GradScaler(init_scale=2.**10)
+    else:
+        scaler = None
+    
+    #-----------------------------------------------------------------------------#
+    # optimizer                                                                   #
+    #-----------------------------------------------------------------------------#
+        
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
+
+        #-----------------------------------------------------------------------------#
+    # Scheduler                                                                   #
+    #-----------------------------------------------------------------------------#
+
+    train_steps = len(train_dataloader) * config.epochs
+    warmup_steps = len(train_dataloader) * config.warmup_epochs
+       
+    if config.scheduler == "polynomial":
+        print("\nScheduler: polynomial - max LR: {} - end LR: {}".format(config.lr, config.lr_end))  
+        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer,
+                                                              num_training_steps=train_steps,
+                                                              lr_end = config.lr_end,
+                                                              power=1.5,
+                                                              num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "cosine":
+        print("\nScheduler: cosine - max LR: {}".format(config.lr))   
+        scheduler = get_cosine_schedule_with_warmup(optimizer,
+                                                    num_training_steps=train_steps,
+                                                    num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "constant":
+        print("\nScheduler: constant - max LR: {}".format(config.lr))   
+        scheduler =  get_constant_schedule_with_warmup(optimizer,
+                                                       num_warmup_steps=warmup_steps)
+           
+    else:
+        scheduler = None
+        
+    print("Warmup Epochs: {} - Warmup Steps: {}".format(str(config.warmup_epochs).ljust(2), warmup_steps))
+    print("Train Epochs:  {} - Train Steps:  {}".format(config.epochs, train_steps))
+
+
+    #-----------------------------------------------------------------------------#
+    # Shuffle                                                                     #
+    #-----------------------------------------------------------------------------#            
+    if config.custom_sampling:
+        train_dataloader.dataset.shuffle()
+    
+    #-----------------------------------------------------------------------------#
+    # Train                                                                       #
+    #-----------------------------------------------------------------------------#
+    start_epoch = 0   
+    best_score = 0
+
+    #-----------------------------------------------------------------------------#
+    # Writer
+    #-----------------------------------------------------------------------------#
+    # Writer
+    writer = SummaryWriter('world_vanillia/cnn' + config.model)
+    LPN_flag = config.agg['agg_config']['LPN']
+    
+
+    for epoch in range(1, config.epochs+1):
+        
+        print("\n{}[Epoch: {}]{}".format(30*"-", epoch, 30*"-"))
+        
+
+        train_loss = trainer.train_backbone(config,
+                           model,
+                           dataloader=train_dataloader,
+                           loss_function=loss_function,
+                           optimizer=optimizer,
+                           scheduler=scheduler,
+                           scaler=scaler,
+                           writer=writer,
+                           LPN=LPN_flag)
+        
+        print("Epoch: {}, Train Loss = {:.3f}, Lr = {:.6f}".format(epoch,
+                                                                   train_loss,
+                                                                   optimizer.param_groups[0]['lr']))
+        
+        #------------------------------------------------------------Eval---------------------------------------------------------------------#
+        result_list = []
+        with open(config.val_index_txt,"r") as val_test:
+            for line in val_test:
+                eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='query',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='DB',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result,_ , _ = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='vanilia',LPN=config.agg['agg_config']['LPN'])
+                print(line.strip('\n'), 'top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2))
+                result_list.append(result)
+                writer.add_scalar(line.strip('\n'), round(result[0]*100,2), epoch)
+                
+        
+        result_array = np.array(result_list)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        writer.add_scalar('Average/top1', round(average_result[0]*100,2), epoch)
+        writer.add_scalar('Average/top5', round(average_result[1]*100,2), epoch)
+        
+    #------------------------------------------------------------Save---------------------------------------------------------------------#
+        if average_result[0] > best_score:
+
+                best_score = average_result[0]
+
+                if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+                    torch.save(model.module.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                else:
+                    torch.save(model.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                
+
+        if config.custom_sampling:
+            train_dataloader.dataset.shuffle()

GeoLoc-UAV-main/train_vanilia_dino.py

@@ -0,0 +1,345 @@
+import os
+import time
+import numpy as np
+import math
+import shutil
+import sys
+import torch
+from dataclasses import dataclass,field
+from torch.cuda.amp import GradScaler
+from torch.utils.data import DataLoader
+from transformers import get_constant_schedule_with_warmup, get_polynomial_decay_schedule_with_warmup, get_cosine_schedule_with_warmup
+from torchvision import transforms as T
+from torch.utils.tensorboard import SummaryWriter
+
+from dataset.World import WorldDatasetTrainVanilia, WorldDatasetEvalVanilia
+from models import model,trainer
+from utils import setting
+from utils import loss
+from eval import eval
+
+
+
+def default_backbone_config():
+
+    return {
+          "backbone_arch" : "dinov2_vits14", 
+          "pretrain_flag":False
+    }
+
+def default_agg_config():
+
+    return {
+          "agg_arch": "multiconvap", #convap
+    "agg_config": {
+      "in_channels": 384, #256 #512,768
+      "out_channels": 384, #256
+      "s1": 1,
+      "s2": 1,
+      'LPN':True
+        }
+    }
+
+@dataclass
+class Configuration:
+
+    model: str = "dinos-newterrain-LPN"
+
+    # Savepath for model checkpoints
+    model_path: str = "./world_vanilia"
+
+    # model config
+    backbone:dict = field(default_factory=default_backbone_config)
+
+    agg:dict = field(default_factory=default_agg_config)
+
+    # dataset
+    dataset_root_dir: str = "/media/Shen/Data/RingoData/WorldLoc"
+    train_query_txt: str = "/media/Shen/Data/RingoData/WorldLoc/Index/train_query.txt" #train_query terrain train_query_country
+
+    # val_index
+    val_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/val.txt" #val.txt
+
+    # test_index
+    test_index_txt = "/media/Shen/Data/RingoData/WorldLoc/Index/test.txt" #test.txt
+
+
+      # Checkpoint to start from
+    checkpoint_start = None
+
+    # set num_workers to 0 if on Windows
+    num_workers: int = 0 if os.name == 'nt' else 8 
+    
+    # train on GPU if available
+    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' 
+
+        # for better performance
+    cudnn_benchmark: bool = True
+    
+    # make cudnn deterministic
+    cudnn_deterministic: bool = False
+
+    # trainning 
+    mixed_precision: bool = True
+    custom_sampling: bool = True         # use custom sampling instead of random
+    seed = 1
+    epochs: int = 10
+    batch_size: int = 128                # keep in mind real_batch_size = 2 * batch_size 128
+    verbose: bool = True
+    gpu_ids: tuple = (0,2,3)           # GPU ids for training
+
+      # Optimizer 
+    clip_grad = 100.                     # None | float
+    decay_exclue_bias: bool = False
+    grad_checkpointing: bool = False     # Gradient Checkpointing
+    
+    # Loss
+    label_smoothing: float = 0.1
+    
+    # Learning Rate
+    lr: float = 0.001                    # 1 * 10^-4 for ViT | 1 * 10^-1 for CNN
+    scheduler: str = "cosine"           # "polynomial" | "cosine" | "constant" | None
+    warmup_epochs: int = 0.1
+    lr_end: float = 0.0001               #  only for "polynomial"
+
+#-------------------------------------------------------------------------------------------#
+# Train Config
+#-------------------------------------------------------------------------------------------#
+config = Configuration() 
+
+
+if __name__ == '__main__':
+
+    model_path = "{}/{}/{}".format(config.model_path,
+                                       config.model,
+                                       time.strftime("%H%M%S"))
+
+    if not os.path.exists(model_path):
+        os.makedirs(model_path)
+    shutil.copyfile(os.path.basename(__file__), "{}/train.py".format(model_path))
+
+    # Redirect print to both console and log file
+    sys.stdout = setting.Logger(os.path.join(model_path, 'log.txt'))
+
+    setting.setup_system(seed=config.seed,
+                 cudnn_benchmark=config.cudnn_benchmark,
+                 cudnn_deterministic=config.cudnn_deterministic)
+    
+    #-----------------------------------------------------------------------------#
+    # Model                                                                       #
+    #-----------------------------------------------------------------------------#
+        
+    print("\nModel: {}".format(config.model))
+
+  
+    # backbone
+    model = model.BackboneGlobal(config.backbone['backbone_arch'],
+                                 config.backbone['pretrain_flag'],
+                               config.agg['agg_arch'],
+                               config.agg['agg_config'])
+    
+     # Load pretrained Checkpoint    
+    if config.checkpoint_start is not None:  
+        print("Start from:", config.checkpoint_start)
+        model_state_dict = torch.load(config.checkpoint_start)  
+        model.load_state_dict(model_state_dict, strict=False)   
+    
+       # Data parallel
+    print("GPUs available:", torch.cuda.device_count())  
+    if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+        model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
+
+      # Model to device   
+    model = model.to(config.device)
+
+    #------------------------setting dataset-------------------------------------------------#
+    IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406], 
+                     'std': [0.229, 0.224, 0.225]}
+    train_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.RandAugment(num_ops=3, interpolation=T.InterpolationMode.BILINEAR),
+            T.AugMix(),
+            # T.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1,
+            #                         hue=0),
+            # T.RandomGrayscale(p=0.2),
+            # T.RandomPosterize(p=0.2, bits=4),
+            # T.GaussianBlur(kernel_size=(1, 5), sigma=(0.1, 5)),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+    
+    eval_transform = T.Compose([
+            T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
+        ])
+
+    #-----------------------------------------------------------------------------#
+    # DataLoader                                                                  #
+    #-----------------------------------------------------------------------------#
+
+    train_dataset = WorldDatasetTrainVanilia(data_dir=config.dataset_root_dir,
+                                      query_txt=config.train_query_txt,
+                                      transforms_query=train_transform,
+                                      transforms_db=train_transform,
+                                      shuffle_batch_size=config.batch_size)
+    
+    # train_dataloader = DataLoader(train_dataset,
+    #                               batch_size=config.batch_size,
+    #                               num_workers=config.num_workers,
+    #                               shuffle=not config.custom_sampling,
+    #                               pin_memory=True)
+    
+    train_dataloader = DataLoader(train_dataset,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=config.custom_sampling,
+                                  pin_memory=True)
+    
+    #-----------------------------------------------------------------------------#
+    # Loss                                                                        #
+    #-----------------------------------------------------------------------------#
+
+    # InfoNCE loss
+    loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=config.label_smoothing)
+    loss_function = loss.InfoNCE(loss_function=loss_fn,
+                            device=config.device,
+                            )
+    # Supervised Contrastive loss
+    # loss_function = loss.SupervisedContrastiveLoss(temperature = 0.07, device=config.device)
+
+    if config.mixed_precision:
+        scaler = GradScaler(init_scale=2.**10)
+    else:
+        scaler = None
+    
+    #-----------------------------------------------------------------------------#
+    # optimizer                                                                   #
+    #-----------------------------------------------------------------------------#
+        
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
+
+        #-----------------------------------------------------------------------------#
+    # Scheduler                                                                   #
+    #-----------------------------------------------------------------------------#
+
+    train_steps = len(train_dataloader) * config.epochs
+    warmup_steps = len(train_dataloader) * config.warmup_epochs
+       
+    if config.scheduler == "polynomial":
+        print("\nScheduler: polynomial - max LR: {} - end LR: {}".format(config.lr, config.lr_end))  
+        scheduler = get_polynomial_decay_schedule_with_warmup(optimizer,
+                                                              num_training_steps=train_steps,
+                                                              lr_end = config.lr_end,
+                                                              power=1.5,
+                                                              num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "cosine":
+        print("\nScheduler: cosine - max LR: {}".format(config.lr))   
+        scheduler = get_cosine_schedule_with_warmup(optimizer,
+                                                    num_training_steps=train_steps,
+                                                    num_warmup_steps=warmup_steps)
+        
+    elif config.scheduler == "constant":
+        print("\nScheduler: constant - max LR: {}".format(config.lr))   
+        scheduler =  get_constant_schedule_with_warmup(optimizer,
+                                                       num_warmup_steps=warmup_steps)
+           
+    else:
+        scheduler = None
+        
+    print("Warmup Epochs: {} - Warmup Steps: {}".format(str(config.warmup_epochs).ljust(2), warmup_steps))
+    print("Train Epochs:  {} - Train Steps:  {}".format(config.epochs, train_steps))
+
+
+    #-----------------------------------------------------------------------------#
+    # Shuffle                                                                     #
+    #-----------------------------------------------------------------------------#            
+    if config.custom_sampling:
+        train_dataloader.dataset.shuffle()
+    
+    #-----------------------------------------------------------------------------#
+    # Train                                                                       #
+    #-----------------------------------------------------------------------------#
+    start_epoch = 0   
+    best_score = 0
+
+    #-----------------------------------------------------------------------------#
+    # Writer
+    #-----------------------------------------------------------------------------#
+    # Writer
+    writer = SummaryWriter('world_vanillia/dinov2/' + config.model)
+    LPN_flag = config.agg['agg_config']['LPN']
+    
+
+    for epoch in range(1, config.epochs+1):
+        
+        print("\n{}[Epoch: {}]{}".format(30*"-", epoch, 30*"-"))
+        
+
+        train_loss = trainer.train_backbone(config,
+                           model,
+                           dataloader=train_dataloader,
+                           loss_function=loss_function,
+                           optimizer=optimizer,
+                           scheduler=scheduler,
+                           scaler=scaler,
+                           writer=writer,
+                           LPN=LPN_flag)
+        
+        print("Epoch: {}, Train Loss = {:.3f}, Lr = {:.6f}".format(epoch,
+                                                                   train_loss,
+                                                                   optimizer.param_groups[0]['lr']))
+        
+        #------------------------------------------------------------Eval---------------------------------------------------------------------#
+        result_list = []
+        with open(config.val_index_txt,"r") as val_test:
+            for line in val_test:
+                eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='query',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_query = DataLoader(eva_dataset_query,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                eva_dataset_db = WorldDatasetEvalVanilia(data_dir=config.dataset_root_dir,
+                                      name=line.strip('\n'),
+                                      mode='DB',
+                                      transforms=eval_transform)
+    
+                eval_dataloader_db = DataLoader(eva_dataset_db,
+                                  batch_size=config.batch_size,
+                                  num_workers=config.num_workers,
+                                  shuffle=not config.custom_sampling,
+                                  pin_memory=True)
+                
+                pos_gt = eval_dataloader_db.dataset.get_gt()
+                result,_ , _ = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode='vanilia',LPN=config.agg['agg_config']['LPN'])
+                print(line.strip('\n'), 'top 1: ', round(result[0]*100,2),  'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2))
+                result_list.append(result)
+                writer.add_scalar(line.strip('\n'), round(result[0]*100,2), epoch)
+                
+        
+        result_array = np.array(result_list)
+        average_result = np.mean(result_array, axis=0)
+        print('Average', 'top 1: ', round(average_result[0]*100,2),  'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
+        writer.add_scalar('Average/top1', round(average_result[0]*100,2), epoch)
+        writer.add_scalar('Average/top5', round(average_result[1]*100,2), epoch)
+        
+    #------------------------------------------------------------Save---------------------------------------------------------------------#
+        if average_result[0] > best_score:
+
+                best_score = average_result[0]
+
+                if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
+                    torch.save(model.module.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                else:
+                    torch.save(model.state_dict(), '{}/weights_e{}_{:.4f}.pth'.format(model_path, epoch,  average_result[0]))
+                
+
+        if config.custom_sampling:
+            train_dataloader.dataset.shuffle()

GeoLoc-UAV-main/utils/evalution.py

@@ -0,0 +1,70 @@
+# 使用鞋带公式（也称为高斯面积公式）来计算多边形的面积
+# 这个示例假设四边形的顶点是按照顺时针或逆时针顺序提供的。如果顶点的顺序不正确，计算的面积可能会是负值
+import numpy as np
+
+def calculate_bbox(polygon):
+    return [
+        min(polygon, key=lambda x: x[0])[0],  # 最小经度
+        min(polygon, key=lambda x: x[1])[1],  # 最小纬度
+        max(polygon, key=lambda x: x[0])[0],  # 最大经度
+        max(polygon, key=lambda x: x[1])[1]   # 最大纬度
+    ]
+
+# 计算交集和并集的边界框
+def calculate_overlap_and_union(bbox1, bbox2):
+    overlap_bbox = [
+        max(bbox1[0], bbox2[0]),
+        max(bbox1[1], bbox2[1]),
+        min(bbox1[2], bbox2[2]),
+        min(bbox1[3], bbox2[3])
+    ]
+    union_bbox = [
+        min(bbox1[0], bbox2[0]),
+        min(bbox1[1], bbox2[1]),
+        max(bbox1[2], bbox2[2]),
+        max(bbox1[3], bbox2[3])
+    ]
+    return overlap_bbox, union_bbox
+
+# 计算面积
+def calculate_area(bbox):
+    return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+
+# 计算IoU
+def calculate_iou(polygon1, polygon2):
+    bbox1 = calculate_bbox(polygon1)
+    bbox2 = calculate_bbox(polygon2)
+    
+    overlap_bbox, union_bbox = calculate_overlap_and_union(bbox1, bbox2)
+    intersection_area = calculate_area(overlap_bbox) if overlap_bbox[0] < overlap_bbox[2] and overlap_bbox[1] < overlap_bbox[3] else 0
+    union_area = calculate_area(union_bbox)
+    
+    return intersection_area / union_area if union_area else 0
+
+def calculate_iou_dict(estimated_dict, real_dict, write_path):
+
+    # 计算多个估计值和真实值之间的IoU
+    ious = {}
+    all_temp = 0
+    with open(write_path, 'w') as f:
+        for key in estimated_dict.keys():
+            if key in real_dict:
+                if estimated_dict[key] != [None]*8:
+                    ious[key] = calculate_iou(estimated_dict[key], real_dict[key])
+                    all_temp += ious[key]
+                    info = key + ' ' + str(ious[key]) + '\n'
+                    f.write(info)
+                else:
+                    info = key + ' ' + str(0) + '\n'
+                    f.write(info)
+                    
+        
+    return ious, all_temp/len(real_dict)
+
+# # 示例：估计的四个点和真实的四个点，每个点是一个 (x, y) 坐标
+# estimated_polygon = [(10, 20), (30, 40), (50, 30), (10, 10)]  # 估计的四边形顶点
+# real_polygon = [(15, 25), (35, 45), (55, 35), (15, 15)]      # 真实的四边形顶点
+
+# # 计算IoU
+# iou = calculate_iou(estimated_polygon, real_polygon)
+# print(f"The IoU between the estimated and real polygons is: {iou:.2f}")

GeoLoc-UAV-main/utils/loss.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed.nn
+
+class InfoNCE(nn.Module):
+
+    def __init__(self, loss_function, device='cuda' if torch.cuda.is_available() else 'cpu'):
+        super().__init__()
+        
+        self.loss_function = loss_function
+        self.device = device
+
+    def forward(self, image_features1, image_features2, logit_scale):
+        
+        image_features1 = F.normalize(image_features1, dim=-1)
+        image_features2 = F.normalize(image_features2, dim=-1)
+        
+        logits_per_image1 = logit_scale * image_features1 @ image_features2.T
+        
+        logits_per_image2 = logits_per_image1.T
+        
+        labels = torch.arange(len(logits_per_image1), dtype=torch.long, device=self.device)
+        
+        loss = (self.loss_function(logits_per_image1, labels) + self.loss_function(logits_per_image2, labels))/2
+
+        return loss  
+ 
+
+class SupervisedContrastiveLoss(nn.Module):
+    def __init__(self, temperature=0.07, device='cuda' if torch.cuda.is_available() else 'cpu'):
+        super(SupervisedContrastiveLoss, self).__init__()
+        
+        self.temperature = temperature
+        self.device = device
+    
+    def forward(self, image_feature, labels):
+
+        dot_product = torch.mm(image_feature, image_feature.T) / self.temperature
+        exp_dot_product = torch.exp(dot_product - torch.max(dot_product, dim=1, keepdim=True)[0]) + 1e-5
+
+        mask_similar_class = (labels.unsqueeze(1).repeat(1, labels.shape[0]) == labels).to(self.device)
+        mask_anchor_out = (1 - torch.eye(exp_dot_product.shape[0])).to(self.device)
+        mask_combined = mask_similar_class * mask_anchor_out
+        per_sample = torch.sum(mask_combined, dim=1)
+
+        log_prob = -torch.log(exp_dot_product / (torch.sum(exp_dot_product * mask_anchor_out, dim=1, keepdim=True)))
+        supervised_loss_per_sample = torch.sum(log_prob * mask_combined, dim=1) / per_sample
+        supervised_loss = torch.mean(supervised_loss_per_sample)
+
+        return supervised_loss
+
+class WeightedInfoNCE(nn.Module):
+    def __init__(self, label_smoothing, k=-5, device='cuda' if torch.cuda.is_available() else 'cpu'):
+        super().__init__()
+        self.label_smoothing = label_smoothing
+        self.device = device
+        self.k = k
+
+    def loss(self, similarity_matrix, eps_all):
+        n = similarity_matrix.shape[0]
+        total_loss = 0.0
+        for i in range(n):
+            eps = eps_all[i]
+            total_loss += (1 - eps) * (-1. * similarity_matrix[i, i] + torch.logsumexp(similarity_matrix[i, :], dim=0))
+            total_loss += eps * (-1. / n * similarity_matrix[i, :].sum() + torch.logsumexp(similarity_matrix[i, :], dim=0))
+        total_loss /= n
+        return total_loss
+
+    def forward(self, image_features1, image_features2, logit_scale, positive_weights=None):
+        # Normalize the image features
+        image_features1 = F.normalize(image_features1, dim=-1)
+        image_features2 = F.normalize(image_features2, dim=-1)
+        
+        # Compute similarity logits
+        logits_per_image1 = logit_scale * image_features1 @ image_features2.T
+        
+        # Apply positive weights if provided
+        if positive_weights is not None:
+            eps = 1. - (1. - self.label_smoothing) / (1 + torch.exp(-self.k * positive_weights))
+        else:
+            eps = [self.label_smoothing for _ in range(image_features1.shape[0])]
+        
+        logits_per_image2 = logits_per_image1.T
+        
+        # Generate labels
+        # labels = torch.arange(len(logits_per_image1), dtype=torch.long, device=self.device)
+
+        loss1 = self.loss(logits_per_image1, eps)
+        loss2 = self.loss(logits_per_image2, eps)
+        # # Compute loss
+        # loss1 = self.loss_function(logits_per_image1, labels)
+        # loss2 = self.loss_function(logits_per_image2, labels)
+        loss = (loss1 + loss2) / 2
+
+        return loss

GeoLoc-UAV-main/utils/setting.py

@@ -0,0 +1,89 @@
+import os
+import sys
+import random
+import errno
+import time
+import torch
+import numpy as np
+from datetime import timedelta
+
+class AverageMeter:
+    """
+    Computes and stores the average and current value
+    """
+
+    def __init__(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val):
+        self.val = val
+        self.sum += val
+        self.count += 1
+        self.avg = self.sum / self.count
+
+def setup_system(seed, cudnn_benchmark=True, cudnn_deterministic=True) -> None:
+    '''
+    Set seeds for for reproducible training
+    '''
+    # python
+    random.seed(seed)
+    
+    # numpy
+    np.random.seed(seed)
+    
+    # pytorch
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    if torch.cuda.is_available():
+        torch.backends.cudnn_benchmark_enabled = cudnn_benchmark
+        torch.backends.cudnn.deterministic = cudnn_deterministic
+      
+        
+def mkdir_if_missing(dir_path):
+    try:
+        os.makedirs(dir_path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+class Logger(object):
+    def __init__(self, fpath=None):
+        self.console = sys.stdout
+        self.file = None
+        if fpath is not None:
+            mkdir_if_missing(os.path.dirname(fpath))
+            self.file = open(fpath, 'w')
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):
+        self.close()
+
+    def write(self, msg):
+        self.console.write(msg)
+        if self.file is not None:
+            self.file.write(msg)
+
+    def flush(self):
+        self.console.flush()
+        if self.file is not None:
+            self.file.flush()
+            os.fsync(self.file.fileno())
+
+    def close(self):
+        self.console.close()
+        if self.file is not None:
+            self.file.close()

GeoLoc-UAV-main/utils/utils.py

@@ -0,0 +1,156 @@
+
+import numpy as np
+import torch
+
+import cv2
+import torch.nn.functional as F
+
+def read_db_pose(txt):
+    db_pose = {}
+    with open(txt, 'r') as f:
+        for line in f:
+            name = line.split(' ')[0]
+            pose = np.asarray(line.split(' ')[1:])
+            db_pose[name] = pose
+    return db_pose
+
+def read_rerank_pose(txt):
+
+    gt_rerank_pose = {}
+    with open(txt, 'r') as f:
+        for line in f:
+            type_name = line.split(' ')[0].split('/')[2]
+            if type_name not in gt_rerank_pose.keys():
+                gt_rerank_pose[type_name] = {}
+            name = line.split(' ')[0].split('/')[-1]
+            left_top = [eval(line.split(' ')[4]), eval(line.split(' ')[5])]
+            right_top = [eval(line.split(' ')[6]), eval(line.split(' ')[7])]
+            right_bottom = [eval(line.split(' ')[8]), eval(line.split(' ')[9])]
+            left_bottom = [eval(line.split(' ')[10]), eval(line.split(' ')[11])]
+            gt_rerank_pose[type_name][name] = [left_top, right_top, right_bottom, left_bottom]
+    return gt_rerank_pose
+
+def dim_extend(data_list):
+    results = []
+    for i, tensor in enumerate(data_list):
+        # 修改
+        if tensor.device is not "cuda":
+            tensor = tensor.cuda()
+        results.append(tensor)#tensor[None,...])
+    return results
+
+def interpolate_feats(img,pts,feats):
+    # compute location on the feature map (due to pooling)
+    _, _, h, w = feats.shape
+    pool_num = img.shape[-1] // feats.shape[-1]
+    pts_warp=(pts+0.5)/pool_num-0.5
+    pts_norm=normalize_coordinates(pts_warp,h,w)
+    pts_norm=torch.unsqueeze(pts_norm, 1)  # b,1,n,2
+
+    # interpolation
+    pfeats=F.grid_sample(feats, pts_norm, 'bilinear',align_corners=False)[:, :, 0, :]  # b,f,n
+    pfeats=pfeats.permute(0,2,1) # b,n,f
+    return pfeats
+
+
+def l2_normalize(x,ratio=1.0,axis=1):
+    norm=torch.unsqueeze(torch.clamp(torch.norm(x,2,axis),min=1e-6),axis)
+    x=x/norm*ratio
+    return x
+
+def normalize_coordinates(coords, h, w):
+    h=h-1
+    w=w-1
+    coords=coords.clone().detach()
+    coords[:, :, 0]-= w / 2
+    coords[:, :, 1]-= h / 2
+    coords[:, :, 0]/= w / 2
+    coords[:, :, 1]/= h / 2
+    return coords
+
+def get_rot_m(angle):
+    return np.asarray([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]], np.float32)
+
+def normalize_image(img, mask=None):
+    if mask is not None: img[np.logical_not(mask.astype(np.bool))]=127
+    img=(img.transpose([2,0,1]).astype(np.float32)-127.0)/128.0
+    return torch.tensor(img,dtype=torch.float32)
+
+class TransformerCV:
+    def __init__(self, cfg):
+        ssb = cfg['sample_scale_begin']
+        ssi = cfg['sample_scale_inter']
+        ssn = cfg['sample_scale_num']
+
+        srb = cfg['sample_rotate_begin'] / 180 * np.pi
+        sri = cfg['sample_rotate_inter'] / 180 * np.pi
+        srn = cfg['sample_rotate_num']
+
+        self.scales = [ssi ** (si + ssb) for si in range(ssn)]
+        self.rotations = [sri * ri + srb for ri in range(srn)]
+
+        self.ssi=ssi
+
+        self.ssn=ssn
+        self.srn=srn
+
+        self.SRs=[]
+        for scale in self.scales:
+            Rs=[]
+            for rotation in self.rotations:
+                Rs.append(scale*get_rot_m(rotation))
+            self.SRs.append(Rs)
+    
+    def transform(self, img, pts=None):
+        '''
+
+        :param img:
+        :return: img_list
+        '''
+        h,w,_=img.shape
+        pts0=np.asarray([[0,0],[0,h],[w,h],[w,0]],np.float32)
+        center = np.mean(pts0, 0)        
+
+        pts_warps, img_warps, grid_warps = [], [], []
+        img_cur=img.copy()
+        for si,Rs in enumerate(self.SRs):
+            if si>0:
+                if self.ssi<0.6:
+                    img_cur=cv2.GaussianBlur(img_cur,(5,5),1.5)
+                else:
+                    img_cur=cv2.GaussianBlur(img_cur,(3,3),0.75)
+            for M in Rs:
+                pts1 = (pts0 - center[None, :]) @ M.transpose()
+                min_pts1 = np.min(pts1, 0)
+                tw, th = np.round(np.max(pts1 - min_pts1[None, :], 0)).astype(np.int32)
+
+                # compute A
+                offset = - M @ center - min_pts1
+                A = np.concatenate([M, offset[:, None]], 1)
+                # note!!!! the border type is constant 127!!!! because in the subsequent processing, we will subtract 127
+                img_warp=cv2.warpAffine(img_cur,A,(tw,th),flags=cv2.INTER_LINEAR,borderMode=cv2.BORDER_CONSTANT,borderValue=(127,127,127))
+
+                # for dino
+                img_warp = cv2.resize(img_warp, (224,224))
+
+                img_warps.append(img_warp[:,:,:3])
+                if pts is not None:
+                    pts_warp = pts @ M.transpose() + offset[None, :]
+                    pts_warps.append(pts_warp)
+                
+        outputs={'img':img_warps}
+        if pts is not None: outputs['pts']=pts_warps
+        
+
+        return outputs
+
+
+
+    @staticmethod
+    def postprocess_transformed_imgs(results):
+        img_list,pts_list=[],[]
+        for img_id, img in enumerate(results['img']):
+            img_list.append(normalize_image(img))
+            pts_list.append(torch.tensor(results['pts'][img_id],dtype=torch.float32))
+
+        return img_list, pts_list