Merge pull request #15 from Dinghye/6-tessera-crs-problem

6 tessera crs problem

Author: Dinghye

src/rs_embed/embedders/precomputed_tessera.py

@@ -87,6 +87,56 @@ def _tile_bounds(transform, w: int, h: int) -> Tuple[float, float, float, float]
     return left, bottom, right, top
 
 
+def _reproject_tile(
+    hwc: np.ndarray,
+    src_transform: Any,
+    src_crs: str,
+    dst_crs: str,
+    target_res: Optional[Tuple[float, float]] = None,
+) -> Tuple[np.ndarray, Any]:
+    """Reproject an HWC embedding tile to *dst_crs* via nearest-neighbour.
+
+    If *target_res* is given as ``(pixel_width, pixel_height)`` (both
+    positive), the output is snapped to that resolution so tiles can be
+    mosaicked without sub-pixel drift.
+    """
+    try:
+        from rasterio.warp import reproject, Resampling, calculate_default_transform
+        from rasterio.transform import array_bounds
+    except ImportError as exc:
+        raise ModelError(
+            "Mixed-CRS mosaic requires rasterio.  Install: pip install rasterio"
+        ) from exc
+
+    h, w, d = hwc.shape
+    src_bounds = array_bounds(h, w, src_transform)
+
+    kwargs: Dict[str, Any] = {}
+    if target_res is not None:
+        kwargs["resolution"] = (abs(target_res[0]), abs(target_res[1]))
+
+    dst_transform, dst_w, dst_h = calculate_default_transform(
+        src_crs, dst_crs, w, h, *src_bounds, **kwargs,
+    )
+
+    dst_hwc = np.zeros((dst_h, dst_w, d), dtype=np.float32)
+    for i in range(d):
+        src_band = np.ascontiguousarray(hwc[:, :, i])
+        dst_band = np.zeros((dst_h, dst_w), dtype=np.float32)
+        reproject(
+            source=src_band,
+            destination=dst_band,
+            src_transform=src_transform,
+            src_crs=src_crs,
+            dst_transform=dst_transform,
+            dst_crs=dst_crs,
+            resampling=Resampling.nearest,
+        )
+        dst_hwc[:, :, i] = dst_band
+
+    return dst_hwc, dst_transform
+
+
 def _reproject_bbox_4326_to(tile_crs_str: str, bbox: BBox) -> Tuple[float, float, float, float]:
     # returns (xmin, ymin, xmax, ymax) in tile CRS
     if str(tile_crs_str).upper() in ("EPSG:4326", "WGS84", "CRS:84"):
@@ -111,34 +161,55 @@ def _mosaic_and_crop_strict_roi(
     tiles_rows: list of (year, tile_lon, tile_lat, embedding_array, crs, transform)
     Return cropped CHW + meta.
     """
+    if not tiles_rows:
+        raise ModelError("No tiles fetched; cannot mosaic.")
+
+    # --- detect mixed CRS and choose the most-common one as target ---
+    crs_counts: Dict[str, int] = {}
+    for _, _, _, _, crs, _ in tiles_rows:
+        key = str(crs)
+        crs_counts[key] = crs_counts.get(key, 0) + 1
+    target_crs = max(crs_counts, key=lambda k: crs_counts[k])
+    mixed_crs = len(crs_counts) > 1
+
+    # when CRS differ, lock the target pixel size from a native tile
+    target_res: Optional[Tuple[float, float]] = None
+    if mixed_crs:
+        for _, _, _, _, crs, transform in tiles_rows:
+            if str(crs) == target_crs:
+                target_res = (abs(float(transform.a)), abs(float(transform.e)))
+                break
+
     # normalize + collect tile meta
     hwc_list = []
-    crs0 = None
+    crs0 = target_crs
     a0 = e0 = None
 
     bounds_list = []
     for year, tlon, tlat, emb, crs, transform in tiles_rows:
         _assert_north_up(transform)
         hwc = _to_hwc(emb)
+
+        # reproject tile to target CRS when CRS differ
+        if str(crs) != target_crs:
+            hwc, transform = _reproject_tile(
+                hwc, transform, str(crs), target_crs, target_res=target_res,
+            )
+
         h, w, d = hwc.shape
         left, bottom, right, top = _tile_bounds(transform, w, h)
 
-        if crs0 is None:
-            crs0 = crs
+        if a0 is None:
             a0 = float(transform.a)
             e0 = float(transform.e)
         else:
-            if str(crs) != str(crs0):
-                raise ModelError("Tiles have different CRS; cannot mosaic.")
-            if abs(float(transform.a) - a0) > 1e-12 or abs(float(transform.e) - e0) > 1e-12:
+            # reprojected tiles may have sub-pixel rounding; use 1e-6 tolerance
+            if abs(float(transform.a) - a0) > 1e-6 or abs(float(transform.e) - e0) > 1e-6:
                 raise ModelError("Tiles have different resolution; cannot mosaic without resampling.")
 
         hwc_list.append((hwc, transform, (left, bottom, right, top)))
         bounds_list.append((left, bottom, right, top))
 
-    if crs0 is None:
-        raise ModelError("No tiles fetched; cannot mosaic.")
-
     # global mosaic bounds
     left = min(b[0] for b in bounds_list)
     bottom = min(b[1] for b in bounds_list)

tests/test_tessera_crs.py

@@ -0,0 +1,145 @@
+"""Tests for mixed-CRS mosaic in precomputed_tessera (GitHub issue #6).
+
+When an ROI sits on a UTM zone boundary (e.g. lon = 120°) the tile
+store can return tiles in *two* EPSG codes (e.g. 32650 / 32651).  The
+mosaic helper must reproject them to a common CRS instead of rejecting
+the operation.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from affine import Affine
+from pyproj import Transformer
+
+from rs_embed.core.specs import BBox
+from rs_embed.core.errors import ModelError
+from rs_embed.embedders.precomputed_tessera import (
+    _mosaic_and_crop_strict_roi,
+    _reproject_tile,
+)
+
+rasterio = pytest.importorskip("rasterio", reason="rasterio needed for mixed-CRS tests")
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+D = 64  # embedding dimension (must be in _to_hwc's allowlist)
+TILE_HW = 8  # small tiles for fast tests
+PX_M = 500.0  # 500 m pixels
+
+
+def _make_tile(lon: float, lat: float, epsg: str, value: float = 1.0):
+    """Create a synthetic (year, lon, lat, embedding, crs, transform) tuple."""
+    tfm = Transformer.from_crs("EPSG:4326", epsg, always_xy=True)
+    x, y = tfm.transform(lon, lat)
+
+    # north-up transform: origin at top-left of tile
+    transform = Affine(PX_M, 0.0, x, 0.0, -PX_M, y + TILE_HW * PX_M)
+    emb = np.full((TILE_HW, TILE_HW, D), value, dtype=np.float32)
+    return (2021, lon, lat, emb, epsg, transform)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestMixedCRSMosaic:
+    """Regression tests for issue #6 – UTM zone boundary."""
+
+    def test_same_crs_still_works(self):
+        """Sanity: single-CRS mosaic must still succeed."""
+        tile_a = _make_tile(119.5, 30.0, "EPSG:32650", value=1.0)
+        tile_b = _make_tile(119.6, 30.0, "EPSG:32650", value=2.0)
+
+        bbox = BBox(minlon=119.45, minlat=29.95, maxlon=119.65, maxlat=30.05)
+        chw, meta = _mosaic_and_crop_strict_roi([tile_a, tile_b], bbox_4326=bbox)
+
+        assert chw.ndim == 3
+        assert chw.shape[0] == D
+        assert meta["tile_crs"] == "EPSG:32650"
+
+    def test_mixed_crs_at_utm_boundary(self):
+        """Two tiles in EPSG:32650 / 32651 at lon ≈ 120° must mosaic."""
+        tile_a = _make_tile(119.95, 30.0, "EPSG:32650", value=1.0)
+        tile_b = _make_tile(120.05, 30.0, "EPSG:32651", value=2.0)
+
+        bbox = BBox(minlon=119.90, minlat=29.95, maxlon=120.10, maxlat=30.05)
+        chw, meta = _mosaic_and_crop_strict_roi([tile_a, tile_b], bbox_4326=bbox)
+
+        assert chw.ndim == 3
+        assert chw.shape[0] == D
+        # target CRS should be one of the two (whichever has more tiles, or
+        # first in case of tie)
+        assert meta["tile_crs"] in ("EPSG:32650", "EPSG:32651")
+
+    def test_mixed_crs_majority_wins(self):
+        """Target CRS is the most-common one among tiles."""
+        tile_a = _make_tile(119.8, 30.0, "EPSG:32650", value=1.0)
+        tile_b = _make_tile(119.9, 30.0, "EPSG:32650", value=1.5)
+        tile_c = _make_tile(120.1, 30.0, "EPSG:32651", value=2.0)
+
+        bbox = BBox(minlon=119.75, minlat=29.95, maxlon=120.15, maxlat=30.05)
+        chw, meta = _mosaic_and_crop_strict_roi(
+            [tile_a, tile_b, tile_c], bbox_4326=bbox,
+        )
+
+        assert chw.ndim == 3
+        assert chw.shape[0] == D
+        # 2 tiles in 32650, 1 in 32651 → target should be 32650
+        assert meta["tile_crs"] == "EPSG:32650"
+
+    def test_empty_tiles_raises(self):
+        """No tiles at all should still raise."""
+        bbox = BBox(minlon=119.0, minlat=29.0, maxlon=120.0, maxlat=30.0)
+        with pytest.raises(ModelError, match="No tiles"):
+            _mosaic_and_crop_strict_roi([], bbox_4326=bbox)
+
+
+class TestReprojectTile:
+    """Unit tests for _reproject_tile helper."""
+
+    def test_identity_reproject(self):
+        """Reprojecting to the same CRS should return data of similar shape."""
+        hwc = np.ones((4, 4, D), dtype=np.float32)
+        transform = Affine(PX_M, 0.0, 500_000.0, 0.0, -PX_M, 3_500_000.0)
+
+        out_hwc, out_tf = _reproject_tile(
+            hwc, transform, "EPSG:32650", "EPSG:32650",
+        )
+
+        assert out_hwc.ndim == 3
+        assert out_hwc.shape[-1] == D
+        np.testing.assert_allclose(out_hwc, 1.0, atol=1e-5)
+
+    def test_cross_zone_reproject(self):
+        """Reproject from zone 50 to zone 51 preserves data values."""
+        hwc = np.full((4, 4, D), 42.0, dtype=np.float32)
+        transform = Affine(PX_M, 0.0, 800_000.0, 0.0, -PX_M, 3_500_000.0)
+
+        out_hwc, out_tf = _reproject_tile(
+            hwc, transform, "EPSG:32650", "EPSG:32651",
+        )
+
+        assert out_hwc.ndim == 3
+        assert out_hwc.shape[-1] == D
+        # nearest-neighbour: non-zero pixels should carry the original value
+        nonzero = out_hwc[out_hwc != 0]
+        if nonzero.size:
+            np.testing.assert_allclose(nonzero, 42.0, atol=1e-5)
+
+    def test_target_res_snaps_resolution(self):
+        """Providing target_res should lock the output pixel size."""
+        hwc = np.ones((4, 4, D), dtype=np.float32)
+        transform = Affine(PX_M, 0.0, 500_000.0, 0.0, -PX_M, 3_500_000.0)
+
+        target_res = (PX_M, PX_M)
+        out_hwc, out_tf = _reproject_tile(
+            hwc, transform, "EPSG:32650", "EPSG:32651", target_res=target_res,
+        )
+
+        assert abs(abs(float(out_tf.a)) - PX_M) < 1e-3
+        assert abs(abs(float(out_tf.e)) - PX_M) < 1e-3