122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
import rasterio
|
|
from cachetools import LRUCache
|
|
from rasterio.io import DatasetReader
|
|
from rasterio.transform import rowcol
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TileInfo:
|
|
path: Path
|
|
bounds: tuple[float, float, float, float] # left, bottom, right, top
|
|
nodata: float | None
|
|
|
|
|
|
class RasterService:
|
|
"""Base GeoTIFF reader with tile index and LRU cache."""
|
|
|
|
def __init__(self, data_dir: str | Path, cache_size: int = 16) -> None:
|
|
self.data_dir = Path(data_dir)
|
|
self._tiles: list[TileInfo] = []
|
|
self._cache: LRUCache[str, DatasetReader] = LRUCache(maxsize=cache_size)
|
|
self._scan_tiles()
|
|
|
|
def _scan_tiles(self) -> None:
|
|
if not self.data_dir.exists():
|
|
return
|
|
|
|
for path in sorted(self.data_dir.glob("*.tif")) + sorted(self.data_dir.glob("*.tiff")):
|
|
try:
|
|
with rasterio.open(path) as dataset:
|
|
bounds = dataset.bounds
|
|
nodata = dataset.nodata
|
|
except rasterio.errors.RasterioIOError:
|
|
continue
|
|
|
|
self._tiles.append(
|
|
TileInfo(
|
|
path=path,
|
|
bounds=(bounds.left, bounds.bottom, bounds.right, bounds.top),
|
|
nodata=nodata,
|
|
)
|
|
)
|
|
|
|
def is_available(self) -> bool:
|
|
return len(self._tiles) > 0
|
|
|
|
def _find_tile(self, lng: float, lat: float) -> TileInfo | None:
|
|
for tile in self._tiles:
|
|
left, bottom, right, top = tile.bounds
|
|
if left <= lng <= right and bottom <= lat <= top:
|
|
return tile
|
|
return None
|
|
|
|
def _open_dataset(self, tile: TileInfo) -> DatasetReader:
|
|
key = str(tile.path)
|
|
if key not in self._cache:
|
|
self._cache[key] = rasterio.open(tile.path)
|
|
return self._cache[key]
|
|
|
|
def close(self) -> None:
|
|
for dataset in self._cache.values():
|
|
dataset.close()
|
|
self._cache.clear()
|
|
|
|
def _normalize_value(self, value: float, nodata: float | None, default: float) -> float:
|
|
if not math.isfinite(value):
|
|
return default
|
|
if nodata is not None and value == nodata:
|
|
return default
|
|
return float(value)
|
|
|
|
def get_value(self, lat: float, lng: float, default: float = 0.0) -> float:
|
|
tile = self._find_tile(lng, lat)
|
|
if tile is None:
|
|
return default
|
|
|
|
dataset = self._open_dataset(tile)
|
|
row, col = rowcol(dataset.transform, lng, lat)
|
|
if row < 0 or col < 0 or row >= dataset.height or col >= dataset.width:
|
|
return default
|
|
|
|
value = float(dataset.read(1, window=((row, row + 1), (col, col + 1)))[0, 0])
|
|
return self._normalize_value(value, tile.nodata, default)
|
|
|
|
def get_values(self, coordinates: Iterable[tuple[float, float]], default: float = 0.0) -> list[float]:
|
|
coords = list(coordinates)
|
|
if not coords:
|
|
return []
|
|
|
|
results = [default] * len(coords)
|
|
by_tile: dict[str, list[tuple[int, float, float]]] = {}
|
|
|
|
for index, (lat, lng) in enumerate(coords):
|
|
tile = self._find_tile(lng, lat)
|
|
if tile is None:
|
|
continue
|
|
by_tile.setdefault(str(tile.path), []).append((index, lat, lng))
|
|
|
|
for tile_path, items in by_tile.items():
|
|
tile = next(t for t in self._tiles if str(t.path) == tile_path)
|
|
dataset = self._open_dataset(tile)
|
|
lats = [lat for _, lat, _ in items]
|
|
lngs = [lng for _, _, lng in items]
|
|
rows, cols = rowcol(dataset.transform, lngs, lats)
|
|
|
|
for (index, lat, lng), row, col in zip(items, rows, cols, strict=True):
|
|
if row < 0 or col < 0 or row >= dataset.height or col >= dataset.width:
|
|
continue
|
|
value = float(dataset.read(1, window=((row, row + 1), (col, col + 1)))[0, 0])
|
|
results[index] = self._normalize_value(value, tile.nodata, default)
|
|
|
|
return results
|
|
|
|
def get_value_at_distance(self, lat: float, lng: float) -> float:
|
|
return self.get_value(lat, lng, default=0.0)
|