SatFetch / src /geo /spatial.py
karansharmaworkspace's picture
Upload 68 files
f343f06 verified
Raw
History Blame Contribute Delete
3.85 kB
"""
H3 spatial indexing for geo-filtered satellite image retrieval.
Uses Uber's H3 hexagonal grid system for efficient spatial queries.
"""
import h3
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
@dataclass
class GeoBox:
"""Bounding box for spatial queries."""
lat_min: float
lat_max: float
lon_min: float
lon_max: float
class SpatialIndex:
"""H3-based spatial index for satellite images."""
def __init__(self, resolution: int = 7):
"""
Initialize spatial index.
Args:
resolution: H3 resolution (0-15). Level 7 ≈ 1.2km cells.
"""
self.resolution = resolution
self.h3_to_indices: Dict[str, List[int]] = {}
self.index_to_geo: Dict[int, Tuple[float, float]] = {}
def add_image(self, index: int, lat: float, lon: float) -> str:
"""Add image coordinates to spatial index. Returns H3 cell."""
h3_cell = h3.latlng_to_cell(lat, lon, self.resolution)
if h3_cell not in self.h3_to_indices:
self.h3_to_indices[h3_cell] = []
self.h3_to_indices[h3_cell].append(index)
self.index_to_geo[index] = (lat, lon)
return h3_cell
def query_radius(self, lat: float, lon: float, radius_km: float = 10.0) -> List[int]:
"""
Find all images within radius of a point.
Args:
lat: Center latitude
lon: Center longitude
radius_km: Search radius in kilometers
Returns:
List of image indices within radius
"""
# Get H3 cells within radius
center_cell = h3.latlng_to_cell(lat, lon, self.resolution)
ring = h3.grid_disk(center_cell, k=max(1, int(radius_km / 5)))
# Collect all image indices in matching cells
results = []
for cell in ring:
if cell in self.h3_to_indices:
results.extend(self.h3_to_indices[cell])
return results
def query_bbox(self, bbox: GeoBox) -> List[int]:
"""
Find all images within a bounding box.
Args:
bbox: Bounding box with lat/lon bounds
Returns:
List of image indices within bbox
"""
results = []
for index, (lat, lon) in self.index_to_geo.items():
if (bbox.lat_min <= lat <= bbox.lat_max and
bbox.lon_min <= lon <= bbox.lon_max):
results.append(index)
return results
def get_neighbors(self, index: int, k: int = 1) -> List[int]:
"""Get neighboring image indices."""
if index not in self.index_to_geo:
return []
lat, lon = self.index_to_geo[index]
center_cell = h3.latlng_to_cell(lat, lon, self.resolution)
ring = h3.grid_disk(center_cell, k=k)
results = []
for cell in ring:
if cell in self.h3_to_indices:
results.extend(self.h3_to_indices[cell])
return [i for i in results if i != index]
def get_cell_count(self) -> int:
"""Number of occupied H3 cells."""
return len(self.h3_to_indices)
def get_image_count(self) -> int:
"""Total number of indexed images."""
return len(self.index_to_geo)
def lat_lon_to_h3(lat: float, lon: float, resolution: int = 7) -> str:
"""Convert lat/lon to H3 cell index."""
return h3.latlng_to_cell(lat, lon, resolution)
def h3_to_lat_lon(h3_cell: str) -> Tuple[float, float]:
"""Convert H3 cell to center lat/lon."""
return h3.cell_to_latlng(h3_cell)
def get_h3_neighbors(h3_cell: str, k: int = 1) -> List[str]:
"""Get neighboring H3 cells."""
return list(h3.grid_disk(h3_cell, k=k))