rnd-v1

bgtopo_poc/georef.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Tuple
+
+import numpy as np
+import rasterio
+from rasterio.errors import RasterioIOError
+from rasterio.transform import rowcol, xy
+from rasterio.windows import Window
+
+LOG = logging.getLogger(__name__)
+
+
+@dataclass
+class RasterHandle:
+    path: Path
+    dataset: rasterio.io.DatasetReader
+    used_map_file: bool
+
+    @property
+    def width(self) -> int:
+        return self.dataset.width
+
+    @property
+    def height(self) -> int:
+        return self.dataset.height
+
+    @property
+    def crs(self):
+        return self.dataset.crs
+
+    @property
+    def transform(self):
+        return self.dataset.transform
+
+    def close(self) -> None:
+        self.dataset.close()
+
+
+def open_georaster(map_path: Optional[str | Path] = None, tif_path: Optional[str | Path] = None) -> RasterHandle:
+    """Open .map if possible, otherwise .tif. .map gives georeferencing through GDAL's MAP driver."""
+    last_err = None
+    if map_path:
+        try:
+            p = Path(map_path)
+            ds = rasterio.open(p)
+            LOG.info("Opened georeferenced MAP dataset: %s", p)
+            return RasterHandle(p, ds, used_map_file=True)
+        except Exception as e:  # noqa: BLE001
+            last_err = e
+            LOG.warning("Could not open MAP dataset %s: %s", map_path, e)
+    if tif_path:
+        try:
+            p = Path(tif_path)
+            ds = rasterio.open(p)
+            LOG.info("Opened TIFF dataset: %s", p)
+            return RasterHandle(p, ds, used_map_file=False)
+        except RasterioIOError as e:
+            last_err = e
+    raise RuntimeError(f"Could not open raster. Last error: {last_err}")
+
+
+def read_window_rgb(ds: rasterio.io.DatasetReader, window: Window) -> np.ndarray:
+    """Read a raster window as uint8 RGB HxWx3."""
+    arr = ds.read(window=window, boundless=True, fill_value=255)
+    if arr.ndim != 3:
+        raise ValueError(f"Expected band-first array, got shape={arr.shape}")
+    if arr.shape[0] >= 3:
+        arr = arr[:3]
+    elif arr.shape[0] == 1:
+        arr = np.repeat(arr, 3, axis=0)
+    arr = np.moveaxis(arr, 0, -1)
+    if arr.dtype != np.uint8:
+        arr = np.clip(arr, 0, 255).astype(np.uint8)
+    return arr
+
+
+def iter_windows(width: int, height: int, tile_size: int, overlap: int):
+    step = max(1, tile_size - overlap)
+    y = 0
+    while y < height:
+        x = 0
+        h = min(tile_size, height - y)
+        while x < width:
+            w = min(tile_size, width - x)
+            yield Window(x, y, w, h)
+            if x + tile_size >= width:
+                break
+            x += step
+        if y + tile_size >= height:
+            break
+        y += step
+
+
+def lonlat_to_pixel(ds: rasterio.io.DatasetReader, lon: float, lat: float) -> Tuple[int, int]:
+    """Convert lon/lat to row/col. Assumes the raster CRS accepts lon/lat or GDAL handles geographic transform.
+
+    For a production version, reproject coordinates into ds.crs with pyproj first. The PoC does that in coordinates.py.
+    """
+    row, col = rowcol(ds.transform, lon, lat)
+    return int(row), int(col)
+
+
+def pixel_to_lonlat(ds: rasterio.io.DatasetReader, row: int, col: int) -> Tuple[float, float]:
+    x, y = xy(ds.transform, row, col, offset="center")
+    return float(x), float(y)
+
+
+def has_real_georef(ds: rasterio.io.DatasetReader) -> bool:
+    try:
+        return ds.crs is not None and not ds.transform.is_identity
+    except Exception:
+        return False