Spaces:
Sleeping
Sleeping
| 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 | |
| # Toggle the following boolean to False if not using HuggingFace App | |
| hf_app = True | |
| if hf_app: | |
| from huggingface_hub import snapshot_download | |
| cmap = plt.get_cmap("jet") | |
| # User parameters | |
| location = 'Cambogan' | |
| ################ Query filenames and directories ################################# | |
| qry_sequence = '20250811_113017' | |
| qry_condition = 'flooded' | |
| qry_camera_pos = 'front' | |
| qry_root_directory = f"/data/FRED/{qry_condition}/KITTI-style" | |
| if (not os.path.exists(qry_root_directory)) and (hf_app): | |
| snapshot_download( | |
| repo_id="CMalone-Jupiter/FRED", | |
| repo_type="dataset", | |
| local_dir="/data/FRED", | |
| allow_patterns=f"{qry_condition}/KITTI-style/{location}_{qry_sequence}/**", | |
| token=os.environ.get("HF_TOKEN") | |
| ) | |
| qry_image_dir = f"{qry_root_directory}/{location}_{qry_sequence}/{qry_camera_pos}-imgs/" | |
| qry_utm_dir = f"{qry_root_directory}/{location}_{qry_sequence}/utm/" | |
| qry_timestamps = [filename.split('.png')[0] for filename in natsorted(os.listdir(qry_image_dir)) if os.path.isfile(qry_image_dir+filename)] | |
| ################ Reference filenames and directories ################################# | |
| ref_sequence = '20250812_122339' | |
| ref_condition = 'dry' | |
| ref_camera_pos = 'front' | |
| ref_root_directory = f"/data/FRED/{ref_condition}/KITTI-style" | |
| if (not os.path.exists(ref_root_directory)) and (hf_app): | |
| snapshot_download( | |
| repo_id="CMalone-Jupiter/FRED", | |
| repo_type="dataset", | |
| local_dir="/data/FRED", | |
| allow_patterns=f"{ref_condition}/KITTI-style/{location}_{ref_sequence}/**", | |
| token=os.environ.get("HF_TOKEN") | |
| ) | |
| ref_image_dir = f"{ref_root_directory}/{location}_{ref_sequence}/{ref_camera_pos}-imgs/" | |
| ref_utm_dir = f"{ref_root_directory}/{location}_{ref_sequence}/utm/" | |
| ref_utms = np.array([np.loadtxt(ref_utm_dir+filename) for filename in natsorted(os.listdir(ref_utm_dir)) if os.path.isfile(ref_utm_dir+filename)]) | |
| 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)]) | |
| img_calib_file = f"./camera_calib.txt" | |
| dist_tolerance = 10 # metres | |
| fig, ax = plt.subplots(1, 2, figsize=(19.4, 6)) | |
| idx = [130]- | |
| def show_image(i): | |
| ax[0].clear() | |
| ax[1].clear() | |
| if i >= len(qry_timestamps): | |
| plt.close(fig) | |
| return | |
| qry_image_timestamp = qry_timestamps[i] | |
| qry_image_filename = f"{qry_image_dir}/{qry_image_timestamp}.png" | |
| qry_utm_timestamp = utils.get_corr_files(qry_image_timestamp, [qry_utm_dir,]) | |
| qry_utm = np.loadtxt(qry_utm_timestamp) | |
| diffs = ref_utms - qry_utm # shape (N, 2) | |
| dists = np.linalg.norm(diffs, axis=1) # shape (N,) | |
| closest_idx = np.argmin(dists) | |
| closest_dist = dists[closest_idx] | |
| qry_image = ImageData(qry_image_filename, img_calib_file) | |
| ax[0].imshow(qry_image.image[:, :, ::-1]) | |
| ax[0].set_title(f"{qry_image_timestamp}.png") | |
| ax[0].axis("off") | |
| if closest_dist <= dist_tolerance: | |
| # Show matching reference image | |
| ref_img_timestamp = utils.get_corr_files(ref_utm_filenames[closest_idx].split('.txt')[0], [ref_image_dir,]) | |
| ref_image = ImageData(ref_img_timestamp, img_calib_file) | |
| ax[1].imshow(ref_image.image[:, :, ::-1]) | |
| ax[1].set_title(f"{ref_img_timestamp.split('/')[-1]}\nDist={closest_dist:.2f}m") | |
| else: | |
| # Show black image with message | |
| black_img = np.zeros_like(qry_image.image) | |
| ax[1].imshow(black_img) | |
| ax[1].text( | |
| 0.5, 0.5, "No reference image found\nwithin distance tolerance", | |
| color="white", fontsize=16, ha="center", va="center", transform=ax[1].transAxes | |
| ) | |
| ax[1].set_title(f"No Match (min dist={closest_dist:.2f}m)") | |
| ax[1].axis("off") | |
| fig.canvas.draw() | |
| def on_key(event): | |
| if event.key in [' ', 'right']: # space or right arrow | |
| idx[0] += 1 | |
| show_image(idx[0]) | |
| elif event.key in [' ', 'left']: # space or right arrow | |
| if idx[0] > 0: | |
| idx[0] -= 1 | |
| show_image(idx[0]) | |
| elif event.key in ['q', 'escape']: # q or Esc → quit | |
| plt.close(fig) | |
| fig.canvas.mpl_connect('key_press_event', on_key) | |
| show_image(idx[0]) | |
| plt.show() |