calc_distance_cvusa.py

import pandas as pd
from sklearn.metrics import DistanceMetric
import torch
import pickle

TOP_K = 128

df_train = pd.read_csv(r'/media/xiapanwang/主数据盘/xiapanwang/Codes/python/New_Geolocalization/0_Datasets/CVUSA/splits/train-19zl.csv', header=None)

df_train = df_train.rename(columns={0: "sat", 1: "ground", 2: "ground_anno"})

df_train["idx"] = df_train.sat.map(lambda x: int(x.split("/")[-1].split(".")[0]))

train_sat_ids = df_train["idx"].values

print("Length Train Ids:", len(train_sat_ids))

df_gps = pd.read_csv(r'/media/xiapanwang/主数据盘/xiapanwang/Codes/python/New_Geolocalization/0_Datasets/CVUSA/all.csv', header=None)

df_gps = df_gps.rename(columns={0: "sat_lat", 1: "sat_long", 2: "ground_lat", 3: "ground_long", 4: "i_dont_know"})

gps_coords = {}
gps_coords_list = []

for idx in train_sat_ids:
    row = df_gps.iloc[idx - 1]

    coordinates = (float(row["ground_lat"]), float(row["ground_long"]))

    gps_coords[idx] = coordinates
    gps_coords_list.append(coordinates)

print("Length of gps coords : " + str(len(gps_coords_list)))
print("Calculation...")

dist = DistanceMetric.get_metric('haversine')
dm = dist.pairwise(gps_coords_list, gps_coords_list)
print("Distance Matrix:", dm.shape)

dm_torch = torch.from_numpy(dm)
dm_torch = dm_torch.fill_diagonal_(dm.max())

values, ids = torch.topk(dm_torch, k=TOP_K, dim=1, largest=False)

values_near_numpy = values.numpy()
ids_near_numpy = ids.numpy()

near_neighbors = dict()

for i, idx in enumerate(train_sat_ids):
    near_neighbors[idx] = train_sat_ids[ids_near_numpy[i]].tolist()

print("Saving...")
with open(r"/media/xiapanwang/主数据盘/xiapanwang/Codes/python/New_Geolocalization/0_Datasets/CVUSA/gps_dict.pkl", "wb") as f:
    pickle.dump(near_neighbors, f)