localisation/VPR_eval-all-v2.py

import os
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import sys
sys.path.append(os.path.abspath("."))   # one level up
import numpy as np
import cv2
import open3d as o3d
from scipy.spatial.transform import Rotation
from utils.lidar import PointCloud
from utils.camera import ImageData
import utils.utils as utils
from natsort import natsorted, index_natsorted
import torch
from tqdm import tqdm
from glob import glob

################## set device based on cuda availability #################
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print('CUDA availability: ' + str(torch.cuda.is_available()))

####################### Functions for matching using numpy on CPU or Pytorch on GPU ###################
def getMatchIndsCPU(ft_ref,ft_qry,topK=20,metric='cosine'):
    """
    metric: 'euclidean' or 'cosine'
    """
    # dMat = cdist(ft_ref,ft_qry,metric)

    ft_qry_norm = ft_qry / np.linalg.norm(ft_qry, axis=1, keepdims=True)  # Shape (M, N)
    ft_ref_norm = ft_ref / np.linalg.norm(ft_ref, axis=1, keepdims=True)  # Shape (C, N)

    # Step 2: Compute cosine similarity
    dMat = 1 - (ft_ref_norm @ ft_qry_norm.T)
    mInds = np.argsort(dMat,axis=0)[:topK].squeeze() # shape: K x ft_qry.shape[0]
    return mInds, dMat


def getMatchIndsGPU(ft_ref, ft_qry,topK=20, metric='cosine'):
    # metric: 'euclidean' or 'cosine'
    ft_qry_tensor = torch.Tensor(ft_qry).to(device)
    ft_ref_tensor = torch.Tensor(ft_ref).to(device)

    if metric == 'euclidean':
        # Use torch's cdist for Euclidean distance
        dMat = torch.cdist(ft_ref, ft_qry)
    
    elif metric == 'cosine':
        # # Normalize both the query and reference tensors
        ft_qry_norm = ft_qry_tensor / ft_qry_tensor.norm(dim=1, keepdim=True)
        ft_ref_norm = ft_ref_tensor / ft_ref_tensor.norm(dim=1, keepdim=True)
        # Compute cosine similarity (1 - cosine similarity for distance)
        dMat = 1 - ft_ref_norm @ ft_qry_norm.t()

    # Get the indices of the top 5 closest matches
    mInds = torch.argsort(dMat.cpu(), dim=0)[:topK].squeeze()
    
    return mInds, dMat

qry_sets = [
    '20210909_124816_v2',
]

ref_sets = [
    '20230509_115540_v2',
]

# vpr_descs = [
#     'cosplace',
#     'boq',
#     'clique-mining',
#     'cricavpr',
#     'eigenplaces',
#     'mixvpr',
#     'megaloc',
#     'salad',
#     'supervlad',
# ]
vpr_descs = [
    'FoL',
]


img_calib_file = f"./camera_calib.txt"

dist_tolerance = 10 # metres
# qry_idx = 4

# User parameters
location = 'dalby-to-brigalow'

################ Reference filenames and directories #################################
ref_condition = ''
ref_camera_pos = 'front'

ref_timestamps = []
ref_utms = []
ref_img_filenames = []
ref_utm_filenames = []

for ref_set in ref_sets:
    print(f"Loading {ref_set}")
    
    ref_root_directory = f"../../Datasets/dalby/{location}"
    ref_vpr_root = f"../../Datasets/dalby/{location}/vpr_ftrs/"
    ref_image_dir = f"{ref_root_directory}/{ref_set}/{ref_camera_pos}-imgs/"
    ref_utm_dir = f"{ref_root_directory}/{ref_set}/utm/"


    this_ref_timestamp = [filename.split('.png')[0] for filename in natsorted(os.listdir(ref_image_dir)) if os.path.isfile(ref_image_dir+filename)]
    ref_utms = ref_utms+[np.loadtxt(ref_utm_dir+filename) for filename in natsorted(os.listdir(ref_utm_dir)) if os.path.isfile(ref_utm_dir+filename)][55::]
    ref_img_filenames = [filename for filename in natsorted(os.listdir(ref_image_dir)) if os.path.isfile(ref_image_dir+filename)]
    ref_utm_filenames = np.array([filename for filename in natsorted(os.listdir(ref_utm_dir)) if os.path.isfile(ref_utm_dir+filename)])[:len(os.listdir(ref_utm_dir))-55]
    ref_timestamps = ref_timestamps+this_ref_timestamp

ref_utms = np.array(ref_utms)

for vpr_desc in vpr_descs:

    all_results = []

    first = True

    print(f"Loading references")

    for ref_set in ref_sets:
        print(f"Loading {ref_set} {vpr_desc} descriptors")
        ref_root_directory = f"../../Datasets/dalby/{location}"
        ref_vpr_root = f"../../Datasets/dalby/{location}/vpr_ftrs/"

        ref_image_dir = f"{ref_root_directory}/{ref_set}/{ref_camera_pos}-imgs/"

        # ref_name_sort_idx = index_natsorted(os.listdir(ref_image_dir))
        # Get the two orderings
        glob_sorted_paths = sorted(glob(f"{ref_image_dir}/*.png"))
        glob_sorted_filenames = [os.path.basename(p) for p in glob_sorted_paths]

        # Get the indices that would sort glob_sorted_filenames into natsorted order
        ref_name_sort_idx = index_natsorted(glob_sorted_filenames)

        ref_ftr = np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/queries_descriptors.npy")
        if first:
            ref_ftrs = ref_ftr[ref_name_sort_idx]
            first = False
        else:
            ref_ftrs = np.vstack((ref_ftrs, ref_ftr[ref_name_sort_idx])) # [ref_name_sort_idx]


    for qry_set in qry_sets:

        ################ Query filenames and directories #################################
        qry_condition = ''
        qry_camera_pos = 'front'

        qry_root_directory = f"../../Datasets/dalby/{location}"
        qry_vpr_root = f"../../Datasets/dalby/{location}/vpr_ftrs/"
        qry_image_dir = f"{qry_root_directory}/{qry_set}/{qry_camera_pos}-imgs/"
        qry_utm_dir = f"{qry_root_directory}/{qry_set}/utm/"


        qry_timestamps = [filename.split('.png')[0] for filename in natsorted(os.listdir(qry_image_dir)) if os.path.isfile(qry_image_dir+filename)]
        qry_utms = np.array([np.loadtxt(qry_utm_dir+filename) for filename in natsorted(os.listdir(qry_utm_dir)) if os.path.isfile(qry_utm_dir+filename)])
        # qry_name_sort_idx = index_natsorted(os.listdir(qry_image_dir))
        qry_ftrs = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/queries_descriptors.npy")

        # Get the two orderings
        glob_sorted_paths = sorted(glob(f"{qry_image_dir}/*.png"))
        glob_sorted_filenames = [os.path.basename(p) for p in glob_sorted_paths]

        # Get the indices that would sort glob_sorted_filenames into natsorted order
        qry_name_sort_idx = index_natsorted(glob_sorted_filenames)

        qry_ftrs = qry_ftrs[qry_name_sort_idx]

        # paths_sorted = sorted(glob(f"{qry_image_dir}/**/*", recursive=True))
        # paths_natsorted = natsorted(os.listdir(qry_image_dir))
        # paths_natsorted_idx = index_natsorted(os.listdir(qry_image_dir))

        # # Compare just the filenames
        # print(paths_sorted[:3])
        # print((np.array(paths_sorted)[paths_natsorted_idx])[:3])
        # print(paths_natsorted[:3])


        # print(torch.cuda.memory_allocated()/1e9)
        # print(torch.cuda.memory_reserved()/1e9)
        mInds, dMat = getMatchIndsGPU(ref_ftrs,qry_ftrs,topK=1)
        mInds = mInds.numpy() # .cpu()
        in_tol = []
        dists = []
        valid_qry = 0

        qry_utm_timestamps, qry_utm_idxs = utils.get_all_corr_files(qry_timestamps, [qry_utm_dir,])
        ref_utm_timestamp, ref_utm_idxs = utils.get_all_corr_files(ref_timestamps, [ref_utm_dir,])

        for qry_idx in tqdm(range(len(qry_timestamps))):

            qry_image_timestamp = qry_timestamps[qry_idx]
            qry_image_filename = f"{qry_image_dir}/{qry_image_timestamp}.png"
            qry_utm = qry_utms[qry_utm_idxs[qry_idx]]


            diffs = ref_utms - qry_utm           # shape (N, 2)
            qry_dists = np.linalg.norm(diffs, axis=1)   # shape (N,)
            if qry_dists.min() > dist_tolerance:
                continue
            else:
                valid_qry += 1

            ref_utm = ref_utms[ref_utm_idxs[int(mInds[qry_idx])]]

            diff = ref_utm - qry_utm           # shape (N, 2)
            dist = np.linalg.norm(diff)   # shape (N,)
            dists.append(dist)
            if dist < dist_tolerance:
                in_tol.append(1)
            else:
                in_tol.append(0)

            # qry_image = ImageData(qry_image_filename, img_calib_file)

            # fig, ax = plt.subplots(1, 2, figsize=(19.4, 6))
            # ax[0].clear()
            # ax[1].clear()

            # ax[0].imshow(qry_image.image[:, :, ::-1])
            # ax[0].set_title(f"{qry_image_timestamp}.png")
            # ax[0].axis("off")

            # # Show matching reference image
            # # ref_img_timestamp = utils.get_corr_files(ref_timestamps[int(mInds[qry_idx])], [ref_image_dir,])
            # ref_image = ImageData(f"{ref_image_dir}/{ref_timestamps[int(mInds[qry_idx])]}.png", img_calib_file)
            # ax[1].imshow(ref_image.image[:, :, ::-1])
            # ax[1].set_title(f"{ref_timestamps[int(mInds[qry_idx])]}\nDist={dist:.2f}m")

            # ax[1].axis("off")
            # fig.canvas.draw()

        print(f"Recall for {qry_set} using {vpr_desc}: {np.sum(np.array(in_tol))/valid_qry:.02%}")
        all_results.append(np.sum(np.array(in_tol))/valid_qry)
        # plt.figure()
        # plt.plot(np.clip(dists, 0, 30))
        # plt.ylim((0,35))
    
    print(f"All {vpr_desc} results:")
    print(all_results)