Spaces:
Running
Running
File size: 2,951 Bytes
e4b6a37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os
import sys
sys.path.append(os.path.abspath(".")) # one level up
import numpy as np
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 ###################
slice_size = 1000
# qry_set = '20210909_124816_v2'
qry_set = '20230509_115540_v2'
vpr_desc = 'FoL'
img_calib_file = f"./camera_calib.txt"
# User parameters
location = 'dalby-to-brigalow'
################ 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/"
save_dir = f"../../Datasets/dalby/{location}/vpr_ftrs/{qry_set}/{vpr_desc}/sliced/"
os.makedirs(save_dir, exist_ok=True)
qry_timestamps = [filename.split('.png')[0] for filename in natsorted(os.listdir(qry_image_dir)) if os.path.isfile(qry_image_dir+filename)]
# 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)
print(f"Loading query features")
qry_ftrs = np.load(f"{qry_vpr_root}/{qry_set}/{vpr_desc}/queries_descriptors.npy")
print(f"Loading query local features")
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]
assert qry_ftrs.shape[0] == qry_local_ftrs.shape[0], f"There should be equal number of global ({qry_ftrs.shape[0]}) and local ({qry_local_ftrs.shape[0]}) features"
check_len_ftrs = 0
check_len_local_ftrs = 0
slice_num = 0
# f"{42:05d}"
print(f"Starting slice n dice")
for idx in tqdm(range(0, qry_ftrs.shape[0], slice_size)):
end_idx = idx+min(slice_size, qry_ftrs.shape[0]-idx)
qry_ftrs_slice = qry_ftrs[idx:end_idx]
qry_local_ftrs_slice = qry_local_ftrs[idx:end_idx]
np.save(f"{save_dir}/queries_descriptors_slice_{slice_num:05d}.npy", qry_ftrs_slice)
np.save(f"{save_dir}/qry_local_feats_slice_{slice_num:05d}.npy", qry_local_ftrs_slice)
check_len_ftrs += qry_ftrs_slice.shape[0]
check_len_local_ftrs += qry_local_ftrs_slice.shape[0]
slice_num += 1
print(f"Query descriptors: {qry_ftrs.shape[0]}, Slice query descriptors: {check_len_ftrs}")
print(f"Query local descriptors: {qry_local_ftrs.shape[0]}, Slice query local descriptors: {check_len_local_ftrs}")
print(f"Number of slices: {slice_num}")
|