localisation/VPR_eval-fol.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 utils.utils import get_all_corr_files
from FoL.reranking import run_rerank
from natsort import natsorted, index_natsorted
import torch
from tqdm import tqdm
from glob import glob
from math import floor

################## 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 = [
    'FoL',
]


img_calib_file = f"./camera_calib.txt"

dist_tolerance = 10 # metres
# qry_idx = 4
slice_len = 1000

# 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))

        if slice_len is None:

            # 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")
            ref_local_ftr = np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/qry_local_feats.npy")
            if first:
                ref_ftrs = ref_ftr[ref_name_sort_idx]
                ref_local_ftrs = ref_local_ftr[ref_name_sort_idx]
                first = False
            else:
                ref_ftrs = np.vstack((ref_ftrs, ref_ftr[ref_name_sort_idx]))
                ref_local_ftrs = np.vstack((ref_local_ftrs, ref_local_ftr[ref_name_sort_idx]))
        
        else:
            num_slices = floor(len(ref_img_filenames)/slice_len)
            if len(ref_img_filenames) % slice_len > 0:
                num_slices += 1

            for idx in tqdm(range(num_slices)):
                if idx == 0:
                    ref_ftrs = np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/sliced/queries_descriptors_slice_{idx:05d}.npy")
                    ref_local_ftrs = np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/sliced/qry_local_feats_slice_{idx:05d}.npy")
                else:
                    ref_ftrs = np.vstack((ref_ftrs, np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/sliced/queries_descriptors_slice_{idx:05d}.npy")))
                    ref_local_ftrs = np.vstack((ref_local_ftrs, np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/sliced/qry_local_feats_slice_{idx:05d}.npy")))

            
            print(f"Loaded ref ftr slices: {len(ref_ftrs)}")
            print(f"Loaded ref local ftr slices: {len(ref_local_ftrs)}")


    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))

        # if slice_len is None:
        #     # 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 = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/queries_descriptors.npy")
        #     qry_local_ftrs = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/qry_local_feats.npy")
        #     qry_ftrs = qry_ftrs[qry_name_sort_idx]
        #     qry_local_ftrs = qry_local_ftrs[qry_name_sort_idx]

        # mInds, dMat = getMatchIndsGPU(ref_ftrs,qry_ftrs,topK=1)
        # mInds = mInds.cpu().numpy()
        if slice_len is None:
            mInds = run_rerank(qry_ftrs, ref_ftrs, qry_local_ftrs, ref_local_ftrs, recall_values=[1, 5, 10, 20])[:,0] # 5, 10, 20
        else:
            print(f"Performing VPR on slices")
            num_slices = floor(len(qry_timestamps)/slice_len)
            if len(qry_timestamps) % slice_len > 0:
                num_slices += 1

            for idx in tqdm(range(num_slices)):
                qry_ftrs = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/sliced/queries_descriptors_slice_{idx:05d}.npy")
                qry_local_ftrs = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/sliced/qry_local_feats_slice_{idx:05d}.npy")
                if idx == 0:
                    mInds = run_rerank(qry_ftrs, ref_ftrs, qry_local_ftrs, ref_local_ftrs, recall_values=[1, 5, 10, 20])[:,0]
                else:
                    mInds_slice = run_rerank(qry_ftrs, ref_ftrs, qry_local_ftrs, ref_local_ftrs, recall_values=[1, 5, 10, 20])[:,0]
                    mInds = np.vstack((np.expand_dims(mInds, axis=1), np.expand_dims(mInds_slice, axis=1))).squeeze()

                del qry_ftrs
                del qry_local_ftrs

            print(f"VPR on query slices: {len(mInds)}")

        np.save(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/mInds.npy", mInds)

        
        in_tol = []
        dists = []
        valid_qry = 0

        qry_utm_timestamps, qry_utm_idxs = get_all_corr_files(qry_timestamps, [qry_utm_dir,])
        ref_utm_timestamp, ref_utm_idxs = 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)

    # else:
    #     num_slices = floor(len(ref_img_filenames)/slice_len)
    #     if len(ref_img_filenames) % slice_len > 0:
    #         num_slices += 1

    #     for idx in range(num_slices):
    #         if idx == 0:
    #             ref_ftrs = np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/sliced/queries_descriptors_slice_{idx:05d}.npy")
    #         else:
    #             ref_ftrs = np.vstack((ref_ftrs, np.load(f"{ref_vpr_root}/{ref_set}/{vpr_desc}/sliced/queries_descriptors_slice_{idx:05d}.npy")))

        
    #     print(f"Loaded ref ftr slices: {len(ref_ftrs)}")




    #     if slice_len is None:
    #         mInds, dMat = getMatchIndsGPU(ref_ftrs,qry_ftrs,topK=1)
    #         mInds = mInds.cpu().numpy()
    #     else:
    #         print(f"Performing VPR on slices")
    #         num_slices = floor(len(qry_timestamps)/slice_len)
    #         if len(qry_timestamps) % slice_len > 0:
    #             num_slices += 1

    #         for idx in tqdm(range(num_slices)):
    #             qry_ftrs = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/sliced/queries_descriptors_slice_{idx:05d}.npy")
    #             if idx == 0:
    #                 mInds, dMat = getMatchIndsGPU(ref_ftrs,qry_ftrs,topK=1)
    #                 mInds = mInds.cpu().numpy()
    #             else:
    #                 mInds_slice, dMat = getMatchIndsGPU(ref_ftrs,qry_ftrs,topK=1)
    #                 mInds_slice = mInds_slice.cpu().numpy()
    #                 mInds = np.vstack((np.expand_dims(mInds, axis=1), np.expand_dims(mInds_slice, axis=1))).squeeze()

    #         print(f"VPR on query slices: {len(mInds)}")