first commit

app/services/raster_base.py

@@ -0,0 +1,121 @@
+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)