import gradio as gr
import torch
import faiss
import numpy as np
import pandas as pd
import folium
from PIL import Image
from pathlib import Path
import torchvision.transforms as tfm
from torchvision.transforms import functional as F
import logging
import sys
import io
import base64
import random
import ast
import webdataset as wds
import os
import pickle
from functools import lru_cache
import tarfile
from huggingface_hub import hf_hub_download, snapshot_download
from models.apl_model_dinov2 import DINOv2FeatureExtractor
sys.path.append(str(Path("image-matching-models")))
sys.path.append(str(Path("image-matching-models/matching/third_party")))
import util_matching
from matching import get_matcher
logging.basicConfig(level=logging.INFO)
HF_TOKEN = os.getenv("HF_TOKEN") # Set this as a secret in your Space
def ensure_files_exist():
Path("./faiss_index").mkdir(exist_ok=True)
Path("./data/webdataset_shards").mkdir(parents=True, exist_ok=True)
"""Check for required files and download if missing"""
# 1. Check FAISS index
if not Path("faiss_index/faiss_index_2021.bin").exists():
print("Downloading FAISS index...")
hf_hub_download(
repo_id='pawlo2013/EarthLoc2_FAISS',
filename="faiss_index.bin",
local_dir="./faiss_index",
token=HF_TOKEN,
repo_type="dataset"
)
# 2. Check WebDataset shards
shard_dir = Path("data/webdataset_shards")
required_shards = [f"shard-{i:06d}.tar" for i in range(11)] # Adjust range as needed
required_indices = [f"{s}.index" for s in required_shards]
missing_files = [
f for f in required_shards + required_indices
if not (shard_dir / f).exists()
]
if missing_files:
print(f"Downloading {len(missing_files)} missing shard files...")
snapshot_download(
repo_id="pawlo2013/EarthLoc_2021_Database",
local_dir=shard_dir,
allow_patterns="*.tar*", # Gets both .tar and .tar.index
token=HF_TOKEN,
repo_type="dataset"
)
# --- Integration Point ---
# Call this BEFORE loading any models or datasets
ensure_files_exist()
# --- Paths and device ---
MODEL_CHECKPOINT_PATH = Path("weights/best_model_95.6.torch")
FAISS_INDEX_PATH = Path("faiss_index/faiss_index.bin")
CSV_MAPPING_PATH = Path("faiss_index/faiss_index_webdataset.csv") # Updated CSV with shards
DEVICE = "cpu"
MATCHING_IMG_SIZE = 512
logging.info(f"Using device: {DEVICE}")
for path, desc in [
(MODEL_CHECKPOINT_PATH, "Model checkpoint"),
(FAISS_INDEX_PATH, "FAISS index"),
(CSV_MAPPING_PATH, "Path mapping CSV"),
]:
if not path.exists():
raise FileNotFoundError(f"{desc} not found at: {path}")
MODEL_NAME = "xfeat_steerers"
matcher = get_matcher(MODEL_NAME, device=DEVICE, max_num_keypoints=2048)
if MODEL_NAME == "xfeat_steerers":
matcher.model.dev = DEVICE
# Load mapping CSV with keys and shard paths
mapping_df = pd.read_csv(CSV_MAPPING_PATH, index_col="faiss_index")
parsed = mapping_df["key"].str.extract(r"@(?P<z>\d{1,2})_(?P<r>\d{1,5})_(?P<c>\d{1,5})@").astype("int32")
mapping_df = mapping_df.join(parsed)
logging.info(f"Loaded mapping CSV with {len(mapping_df)} entries.")
# Cache for opened shards: {shard_path: WebDataset object}
shard_cache = {}
def get_shard_dataset(shard_path):
"""Load or get cached WebDataset for a shard path."""
if shard_path not in shard_cache:
shard_cache[shard_path] = wds.WebDataset(shard_path, handler=wds.warn_and_continue)
return shard_cache[shard_path]
@lru_cache(maxsize=100)
def load_index(index_path):
with open(index_path, "rb") as f:
return pickle.load(f)
def load_image_from_shard(key):
row = mapping_df[mapping_df["key"] == key]
if row.empty:
return None
shard_path = row.iloc[0]["shard_path"]
index_path = shard_path + ".index"
if not os.path.exists(index_path):
return _load_via_linear_scan(shard_path, key) # Fallback
try:
index = load_index(index_path)
offset = index.get(key)
if offset is None:
return None
with open(shard_path, "rb") as f:
f.seek(offset)
with tarfile.open(fileobj=f) as tar:
member = tar.next()
if member and member.name.startswith(key):
jpg_file = tar.extractfile(member)
return Image.open(io.BytesIO(jpg_file.read())).convert("RGB")
return None
except Exception as e:
logging.error(f"Error loading {key}: {str(e)}")
return _load_via_linear_scan(shard_path, key) # Fallback on error
# Fallback linear scan (original method)
def _load_via_linear_scan(shard_path, key):
dataset = get_shard_dataset(shard_path)
for sample in dataset:
if sample["__key__"] == key:
if img_bytes := sample.get("jpg"):
return Image.open(io.BytesIO(img_bytes)).convert("RGB")
return None
def pil_to_base64(image):
"""Convert a PIL image to a base64-encoded string for HTML embedding."""
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
return f"data:image/png;base64,{img_str}"
def create_map(
final_footprint, candidate_num, filename, inliers, query_footprint=None
):
"""
Create and return a Folium map (as HTML string) showing the final footprint (blue)
and optionally the query's ground truth footprint (orange).
"""
if final_footprint:
center = final_footprint[0]
zoom = 10
elif query_footprint:
center = query_footprint[0]
zoom = 10
else:
center = [0, 0]
zoom = 2
m = folium.Map(location=center, zoom_start=zoom)
if query_footprint:
folium.Polygon(
locations=query_footprint,
popup="Ground Truth Query Footprint",
color="orange",
fill=True,
fill_color="orange",
fill_opacity=0.4,
).add_to(m)
if final_footprint:
footprint_text = "\n".join(
[f"({lat:.4f}, {lon:.4f})" for lat, lon in final_footprint]
)
popup_text = (
f"Predicted Footprint:<br>{footprint_text}<br><br>"
f"Candidate: {candidate_num}<br>Inliers: {inliers}"
)
folium.Polygon(
locations=final_footprint,
popup=popup_text,
color="blue",
fill=True,
fill_color="cyan",
fill_opacity=0.5,
).add_to(m)
folium.Marker(
location=[final_footprint[0][0], final_footprint[0][1]],
popup=f"Footprint Coordinates:<br>{footprint_text}",
icon=folium.Icon(color="blue"),
).add_to(m)
elif not query_footprint:
folium.Marker(
location=[0, 0],
popup="No valid location found.",
icon=folium.Icon(color="red"),
).add_to(m)
return m._repr_html_()
def parse_zoom_row_col_from_key(key: str):
"""
Extract zoom, row, col from the complex key string.
The key format is like:
2021_10_30_90_@34,30714@92,81250@35,46067@92,81250@35,46067@94,21875@34,30714@94,21875@10_0404_0776@2021@34,88391@93,51562@16489@0@
The zoom, row, col are expected to be in the last underscore-separated fields,
specifically the last three fields before the final '@' or at the end.
This function tries to extract zoom, row, col as integers.
"""
try:
# Split by underscore
parts = util_matching.get_image_metadata_from_path(key)
image_id_str = parts[9] # 9th field: image_id
parts = image_id_str.split("_")
zoom = int(parts[0])
row = int(parts[1])
col = int(parts[2])
return zoom, row, col
except Exception as e:
raise ValueError(f"Failed to parse zoom,row,col from key: {key}") from e
def get_surrounding_tiles_sharded(candidate_key, zoom):
"""
Given a candidate key, find all keys in mapping_df with the same zoom,
and row/col within ±4 offsets, then load images from shards.
Return list of (img, (row, col), key) sorted by row, col.
"""
try:
zoom, row, col = parse_zoom_row_col_from_key(candidate_key)
except Exception as e:
logging.warning(f"Failed to parse candidate key {candidate_key}: {e}")
return []
row_offsets = [-4, 0, 4]
col_offsets = [-4, 0, 4]
desired_rows = {row + r for r in row_offsets}
desired_cols = {col + c for c in col_offsets}
# ── 2. Vectorised filter ──────────────────────────────────────────────────────
mask = (
(mapping_df["z"] == zoom) &
(mapping_df["r"].isin(desired_rows)) &
(mapping_df["c"].isin(desired_cols))
)
matched_rows = mapping_df[mask]
tiles = []
seen_positions = set() # Track (row, col) to avoid duplicates
for _, row_data in matched_rows.iterrows():
k = row_data["key"]
try:
_, r, c = parse_zoom_row_col_from_key(k)
except Exception:
continue
if (r, c) in seen_positions:
continue # Skip duplicate position
img = load_image_from_shard(k)
if img is not None:
tiles.append((img, (r, c), k))
seen_positions.add((r, c))
tiles.sort(key=lambda t: (t[1][0], t[1][1]))
return tiles
def compose_tiles_ordered_sharded(tiles, tile_size, candidate_indices):
"""
Compose a 3x3 grid image from tiles loaded from shards.
Missing tiles replaced with blank.
"""
candidate_row, candidate_col = candidate_indices
grid_img = Image.new("RGB", (tile_size[0] * 3, tile_size[1] * 3))
blank = Image.new("RGB", tile_size, color=(0, 0, 0))
tile_dict = {(rc[0], rc[1]): img for img, rc, key in tiles if img is not None}
for i, row_offset in enumerate([-4, 0, 4]):
for j, col_offset in enumerate([-4, 0, 4]):
desired_row = candidate_row + row_offset
desired_col = candidate_col + col_offset
img = tile_dict.get((desired_row, desired_col), blank)
if img.mode != "RGB":
img = img.convert("RGB")
img_resized = tfm.Resize(tile_size, antialias=True)(img).copy()
grid_img.paste(img_resized, (j * tile_size[0], i * tile_size[1]))
return grid_img
def run_matching(query_image, candidate_image, base_footprint):
local_fm = None
viz_params = None
for iteration in range(4):
(
num_inliers,
local_fm,
predicted_footprint,
pretty_footprint,
) = util_matching.estimate_footprint(
local_fm,
query_image,
candidate_image,
matcher,
base_footprint,
HW=MATCHING_IMG_SIZE,
viz_params=viz_params,
)
if num_inliers == -1 or num_inliers is None:
return -1, []
if hasattr(predicted_footprint, "tolist"):
best_footprint = predicted_footprint.tolist()
else:
best_footprint = predicted_footprint
return num_inliers, best_footprint
# --- Load assets ---
logging.info("Loading assets. This may take a moment...")
try:
model = DINOv2FeatureExtractor(
model_type="vit_base_patch14_reg4_dinov2.lvd142m",
num_of_layers_to_unfreeze=0,
desc_dim=768,
aggregator_type="SALAD",
)
logging.info(f"Loading model checkpoint from {MODEL_CHECKPOINT_PATH}...")
model_state_dict = torch.load(MODEL_CHECKPOINT_PATH, map_location=DEVICE)
model.load_state_dict(model_state_dict)
model = model.to(DEVICE)
model.eval()
logging.info("DINOv2 model and checkpoint loaded successfully.")
except Exception as e:
logging.error(f"Failed to load the model: {e}")
raise
faiss_index = faiss.read_index(str(FAISS_INDEX_PATH))
num_db_images = faiss_index.ntotal // 4
logging.info(
f"FAISS index loaded. Contains {faiss_index.ntotal} vectors for {num_db_images} unique images."
)
image_transform = tfm.Compose(
[
tfm.Resize((model.image_size, model.image_size), antialias=True),
tfm.ToTensor(),
tfm.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
),
]
)
logging.info("Assets loaded. Gradio app is ready.")
# --- Core app logic ---
def search_and_retrieve(
query_image: Image.Image, query_footprint_str: str, num_results: int
):
progress = gr.Progress()
query_footprint = None
if query_footprint_str:
try:
print(query_footprint_str)
query_footprint = ast.literal_eval(query_footprint_str)
query_footprint = [list(coord) for coord in query_footprint]
except (ValueError, SyntaxError):
logging.warning("Could not parse query footprint string.")
query_footprint = None
if query_image is None:
yield create_map(None, None, None, None), None
return
progress(0.1, desc="Preprocessing query")
if query_image.mode == "RGBA":
query_image = query_image.convert("RGB")
image_tensor = image_transform(query_image).to(DEVICE)
with torch.no_grad():
descriptor = model(image_tensor.unsqueeze(0))
descriptor_np = descriptor.cpu().numpy()
progress(0.2, desc=f"Searching database for {num_results} neighbors")
distances, indices = faiss_index.search(descriptor_np, num_results)
flat_indices = indices.flatten()
global_best_inliers = -1
global_best_footprint = None
global_candidate_num = None
global_filename = None
global_best_display_image = None
candidate_infos = []
processed_image_indices = set()
query_tensor = tfm.ToTensor()(
tfm.Resize((MATCHING_IMG_SIZE, MATCHING_IMG_SIZE), antialias=True)(query_image)
)
progress(0.4, desc="Processing candidates")
for faiss_idx in flat_indices:
image_index = faiss_idx % num_db_images
best_rotation_index = faiss_idx // num_db_images
query_tensor = F.rotate(query_tensor, [0, -90, -180, -270][best_rotation_index] )
if image_index in processed_image_indices:
continue
processed_image_indices.add(image_index)
candidate_num = len(candidate_infos) + 1
try:
candidate_row = mapping_df.loc[int(image_index)]
except Exception as e:
logging.warning(f"Failed to get candidate info for index {image_index}: {e}")
continue
candidate_key = candidate_row["key"]
candidate_path_str = candidate_row["local_path"]
shard_path = candidate_row['shard_path']
base_footprint = util_matching.path_to_footprint(Path(candidate_path_str))
try:
parts = util_matching.get_image_metadata_from_path(candidate_key)
image_id_str = parts[9] # 9th field: image_id
parts = image_id_str.split("_")
zoom = int(parts[0])
candidate_row_idx = int(parts[1])
candidate_col_idx = int(parts[2])
except Exception as e:
logging.warning(f"Failed to parse candidate key {candidate_key}: {e}")
continue
debug_dir = Path("debug_tiles")
debug_dir.mkdir(exist_ok=True)
tiles = get_surrounding_tiles_sharded(candidate_key, zoom)
composite_img = compose_tiles_ordered_sharded(
tiles, (1024, 1024), (candidate_row_idx, candidate_col_idx)
)
display_img = F.rotate(
composite_img, [0, 90, 180, 270][best_rotation_index]
)
candidate_img_tensor = tfm.ToTensor()(composite_img)
candidate_img_tensor = tfm.Resize(
(MATCHING_IMG_SIZE * 3, MATCHING_IMG_SIZE * 3), antialias=True
)(candidate_img_tensor)
candidate_img_tensor = candidate_img_tensor.to(DEVICE)
progress(
0.5 + len(candidate_infos) / num_results * 0.4,
desc=f"Running matching for candidate {candidate_num}",
)
best_inliers, best_footprint = run_matching(
query_tensor, candidate_img_tensor, base_footprint
)
if best_inliers > -1:
candidate_infos.append(
{
"candidate_num": candidate_num,
"filename": Path(candidate_path_str).name,
"inliers": best_inliers,
"display_image": display_img,
"footprint": best_footprint,
}
)
if best_inliers > global_best_inliers:
global_best_inliers = best_inliers
global_best_footprint = best_footprint
global_candidate_num = candidate_num
global_filename = Path(candidate_path_str).name
global_best_display_image = display_img
progress(0.9, desc="Finalizing results")
folium_map_html = create_map(
global_best_footprint,
global_candidate_num,
global_filename,
global_best_inliers,
query_footprint=query_footprint,
)
progress(1, desc="Done")
yield folium_map_html, None
yield folium_map_html, global_best_display_image
# --- Gradio app setup ---
if __name__ == "__main__":
example_list = []
google_examples = []
queries_folder = Path("./data/queries")
if queries_folder.exists() and queries_folder.is_dir():
image_extensions = ["*.jpg", "*.jpeg", "*.png"]
image_files = []
for ext in image_extensions:
image_files.extend(queries_folder.glob(ext))
if image_files:
num_examples = min(10, len(image_files))
random_examples = random.sample(image_files, num_examples)
example_list = [
[str(p), str(util_matching.get_footprint_from_path(p))]
for p in random_examples
]
logging.info(
f"Loaded {len(example_list)} examples for Gradio with footprints."
)
else:
logging.warning(
f"No images found in the examples folder: {queries_folder}"
)
else:
logging.warning(f"Examples folder not found: {queries_folder}")
google_folder = Path("./data/google_maps_queries")
if google_folder.exists() and google_folder.is_dir():
image_extensions = ["*.jpg", "*.jpeg", "*.png"]
google_files = []
for ext in image_extensions:
google_files.extend(google_folder.glob(ext))
if google_files:
num_google = min(10, len(google_files))
google_examples = [
[str(p), str(p.stem).split("_")[0]] # Empty footprint for Google Maps
for p in random.sample(google_files, num_google)
]
model_description = """
## Model Details
This is a public API for inference of the EarthLoc2 model, which implements the amazing works of:
- EarthLoc (https://earthloc-and-earthmatch.github.io/)
- EarthMatch (https://earthloc-and-earthmatch.github.io/)
- AstroLoc (https://astro-loc.github.io/)
### Architecture
- DINOv2 base with SALAD aggregator out dim = 3072
- FAISS index ~ 8gb, indexes 161496 * 4 images (4 rotated versions) from 2021
### Training
- Trained on the original EarthLoc dataset (zooms 9,10,11), in range -60,60 latitude, polar regions not supported
- Training included additional queries which were not part of the test/val sets
- 5000 iterations with a batch size of 96
### Performance
- Achieves R@10 = 90.6 on the original EarthLoc test and val sets (when retrieving against whole db as is)
- Overall performance is around 10% worse than AstroLoc (https://9d214e4bc329a5c3f9.gradio.live/)
- Works well on satelite images between 1000 sq.km and 50000 sq.km, smaller or higher areas will not produce good results.
### Matching
- Uses the Xfeat_steerers matcher with 2048 maximal number of keypoints, we recommend Master with 2048 if you have access to GPU (we are too poor for it).
"""
theme = gr.themes.Soft(
primary_hue=gr.themes.colors.blue,
font=gr.themes.GoogleFont("Inter"),
).set(
button_primary_background_fill="*primary_900",
button_primary_background_fill_hover="*primary_700",
)
with gr.Blocks(theme=theme) as demo:
gr.Markdown("# Aerial Photography Locator ")
with gr.Row():
with gr.Column(scale=2):
image_input = gr.Image(
type="pil",
label="Aerial Photos of Earth",
height=400,
)
hidden_footprint_text = gr.Textbox(
visible=False, label="Query Footprint"
)
slider = gr.Slider(
minimum=1,
maximum=20,
step=1,
value=10,
label="Number of Candidates to Process",
info=(
"The higher this number to more likely the model is to find a match, "
"however it takes longer to find it. Expect around 5 second more compute per candidate."
),
)
submit_btn = gr.Button("Localize Image", variant="primary")
with gr.Row():
with gr.Column():
if example_list:
gr.Markdown("### ISS Example Queries")
gr.Examples(
examples=example_list,
inputs=[image_input, hidden_footprint_text],
examples_per_page=5,
cache_examples=False,
)
with gr.Column():
if google_examples:
gr.Markdown("### Google Maps Example Queries")
gr.Examples(
examples=google_examples,
inputs=[image_input, hidden_footprint_text],
examples_per_page=5,
cache_examples=False,
)
with gr.Column(scale=2):
map_output = gr.HTML(label="Final Footprint Map")
image_output = gr.Image(
type="pil",
label="Best Matching Candidate",
height=400,
show_download_button=True,
)
gr.Markdown(model_description)
submit_btn.click(
fn=search_and_retrieve,
inputs=[image_input, hidden_footprint_text, slider],
outputs=[map_output, image_output],
)
demo.launch(share=True)