height_predictor/get_height_gba.py

import math
import ee


def gba_tile_id(lon: float, lat: float, step_deg: int = 5) -> str:
    lon0 = math.floor(lon / step_deg) * step_deg
    lat0 = math.floor(lat / step_deg) * step_deg
    lon1 = lon0 + step_deg
    lat1 = lat0 + step_deg

    def fmt_lon(x: int) -> str:
        return ("e" if x >= 0 else "w") + f"{abs(x):03d}"

    def fmt_lat(x: int) -> str:
        return ("n" if x >= 0 else "s") + f"{abs(x):02d}"

    return f"{fmt_lon(lon0)}_{fmt_lat(lat1)}_{fmt_lon(lon1)}_{fmt_lat(lat0)}"


def gba_fc_path(tile_id: str) -> str:
    return f"projects/sat-io/open-datasets/GLOBAL_BUILDING_ATLAS/{tile_id}"


class GlobalBuildingAtlasHeight:
    def __init__(self) -> None:
        self._fc_cache = {}

    def _get_fc(self, tile_id: str) -> ee.FeatureCollection:
        if tile_id not in self._fc_cache:
            self._fc_cache[tile_id] = ee.FeatureCollection(gba_fc_path(tile_id))
        return self._fc_cache[tile_id]

    def get_height_m(self, lat: float, lon: float, buffer_m: float = 20.0, centroid_scale_m: float = 1.0):
        tile_id = gba_tile_id(lon, lat)
        fc = self._get_fc(tile_id)

        pt = ee.Geometry.Point([lon, lat])
        search_geom = pt.buffer(buffer_m) if buffer_m and buffer_m > 0 else pt

        candidates = (
            fc.filterBounds(search_geom)
            .filter(ee.Filter.gt("height", 0))
        )

        def add_dist(f):
            f = ee.Feature(f)
            d = f.geometry().centroid(centroid_scale_m).distance(pt)
            return f.set("dist_m", d)

        best = ee.Feature(candidates.map(add_dist).sort("dist_m").first())

        count = candidates.size().getInfo()
        if count == 0:
            return {
                "status": "not_found",
                "tile_id": tile_id,
                "buffer_m": buffer_m,
                "predicted_height": None,
            }

        props = best.toDictionary(["height", "dist_m"]).getInfo()

        return {
            "status": "success",
            "tile_id": tile_id,
            "buffer_m": buffer_m,
            "predicted_height": props.get("height"),
            "distance_m": props.get("dist_m"),
        }