diff --git a/.github/workflows/k8s-integration.yml b/.github/workflows/k8s-integration.yml index a2c25b8..5833fba 100644 --- a/.github/workflows/k8s-integration.yml +++ b/.github/workflows/k8s-integration.yml @@ -34,6 +34,8 @@ jobs: enable-cache: true - uses: extractions/setup-just@v2 + with: + just-version: "1.39.0" - uses: helm/kind-action@v1 with: diff --git a/.github/workflows/validate-stac.yml b/.github/workflows/validate-stac.yml index 899e71d..ace9332 100644 --- a/.github/workflows/validate-stac.yml +++ b/.github/workflows/validate-stac.yml @@ -36,6 +36,8 @@ jobs: enable-cache: true - name: Install just uses: extractions/setup-just@v2 + with: + just-version: "1.39.0" - name: Install dependencies run: uv sync --group dev - name: Validate STAC items diff --git a/.gitignore b/.gitignore index 4329b6c..2880253 100644 --- a/.gitignore +++ b/.gitignore @@ -225,11 +225,19 @@ implementation.md # data data/sample/predict/predictions +data/sample/ramp_work +data/sample/yolo_work stac_catalog # hatch-vcs generated version file fair/_version.py + +# runs +runs/ + +# weights +yolov8s_v2-seg.pt # dev infra infra/dev/.port-forward.pids infra/dev/.pf-logs/ diff --git a/models/ramp/Dockerfile b/models/ramp/Dockerfile new file mode 100644 index 0000000..338fd42 --- /dev/null +++ b/models/ramp/Dockerfile @@ -0,0 +1,86 @@ +# syntax=docker/dockerfile:1.7 + +# Base image: ghcr.io/hotosm/fair-utilities-ramp:cpu-latest (or :gpu-latest via BASE_IMAGE build arg). +# Build from the fAIr-models repo root: +# docker build -f models/ramp/Dockerfile --target test -t ramp:test . +# docker build -f models/ramp/Dockerfile --target runtime -t ramp:runtime . +# docker build -f models/ramp/Dockerfile --target inference -t ramp:inference . +ARG BASE_IMAGE=ghcr.io/hotosm/fair-utilities-ramp:cpu-latest + +# --------------------------------------------------------------------------- +# Builder stage: install model-pack deps into the base image venv +# --------------------------------------------------------------------------- +FROM ${BASE_IMAGE} AS builder + +ENV UV_LINK_MODE=copy + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /usr/local/bin/ + +WORKDIR /workspace + +COPY pyproject.toml README.md fair_zenml_patch.pth /tmp/fair-src/ +COPY fair /tmp/fair-src/fair + +# The fair-utilities-ramp image already provides hot_fair_utilities, ramp-fair, +# segmentation_models, and TensorFlow under /app/.venv. Add the fAIr model-pack +# deps (fair itself, STAC validation, fairpredictor, onnx tooling) into that +# same venv so training, tests, and inference share one interpreter. +RUN --mount=type=cache,target=/root/.cache/uv \ + SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 \ + uv pip install --python /app/.venv/bin/python \ + "pystac[validation]>=1.14.3" \ + "universal-pathlib" \ + "fairpredictor>=0.5.1" \ + "onnx>=1.16" \ + "tf2onnx>=1.16" \ + /tmp/fair-src[k8s] + +# --------------------------------------------------------------------------- +# Runtime stage: production image (no test deps) +# --------------------------------------------------------------------------- +FROM ${BASE_IMAGE} AS runtime + +WORKDIR /workspace + +ENV PATH="/app/.venv/bin:$PATH" \ + MPLBACKEND=Agg \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + RAMP_HOME=/app \ + SM_FRAMEWORK=tf.keras + +COPY --from=builder /app/.venv /app/.venv + +COPY models/ramp models/ramp +COPY models/conftest.py models/conftest.py +COPY models/test_integration.py models/test_integration.py + +ENTRYPOINT ["/app/.venv/bin/python"] + +# --------------------------------------------------------------------------- +# Test stage: add pytest + ZenML server deps for step & integration tests +# --------------------------------------------------------------------------- +FROM runtime AS test + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv +COPY --from=builder /tmp/fair-src /tmp/fair-src + +RUN --mount=type=cache,target=/root/.cache/uv \ + SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 \ + uv pip install --python /app/.venv/bin/python /tmp/fair-src[test] + +# --------------------------------------------------------------------------- +# Inference stage: runtime + serving deps for the smoke test / live API +# --------------------------------------------------------------------------- +FROM runtime AS inference + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv +COPY --from=builder /tmp/fair-src /tmp/fair-src + +RUN --mount=type=cache,target=/root/.cache/uv \ + SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 \ + uv pip install --python /app/.venv/bin/python /tmp/fair-src[serve] + +ENV PYTHONPATH=/workspace +EXPOSE 8080 +CMD ["/app/.venv/bin/uvicorn", "fair.serve.base:create_app", "--factory", "--host", "0.0.0.0", "--port", "8080"] diff --git a/models/ramp/README.md b/models/ramp/README.md new file mode 100644 index 0000000..e55d189 --- /dev/null +++ b/models/ramp/README.md @@ -0,0 +1,60 @@ +# RAMP EfficientNetB0 + U-Net Building Segmentation + +Semantic segmentation model for building footprint extraction from RGB aerial imagery, derived from the RAMP (Replicable AI for Microplanning) project. + +## Architecture + +- **Model**: EfficientNetB0 encoder + U-Net decoder (`EffUnet`) +- **Framework**: TensorFlow 2.15 / `tf.keras` (via `segmentation_models` with `SM_FRAMEWORK=tf.keras`) +- **Task**: Semantic segmentation (sparse categorical crossentropy) +- **Input**: RGB chips (256x256, float32, channels-last) +- **Classes**: 4 (background=0, building=1, boundary=2, contact=3) + +The boundary (class 2) and contact (class 3) channels help the model cleanly separate adjacent buildings at inference time, even when they share a wall. The `predict()` helper collapses the 4-class softmax to a binary building mask before vectorization. + +## Pretrained Source + +Baseline RAMP weights (TensorFlow SavedModel) hosted by HOTOSM: + +- Checkpoint: `https://api-prod.fair.hotosm.org/api/v1/workspace/download/ramp/baseline.zip` +- ONNX model: `https://api-prod.fair.hotosm.org/api/v1/workspace/download/ramp/ramp-v1.onnx` + +## Pipeline + +Training pipeline steps (ZenML) defined in `pipeline.py`: + +- `split_dataset` - preprocesses chips + labels into 4-class multimasks and produces a seeded random train/validation split via `hot_fair_utilities.split_training_2_validation` +- `train_model` - fine-tunes RAMP on the split, returning the best SavedModel serialized as a zipped byte stream +- `evaluate_model` - computes `fair:accuracy`, `fair:mean_iou` (building IoU), `fair:precision`, and `fair:recall` on the validation split +- `export_onnx` - converts the trained SavedModel to an ONNX byte stream via `tf2onnx` + +Inference is served through `fair.serve.base`, which calls the module-level `predict(session, input_images, params) -> FeatureCollection`: each chip is preprocessed, run through an `onnxruntime` session, decoded to a binary building mask, and vectorized to georeferenced polygons. + +## Base Image + +Training, test, and inference Docker stages all build on +`ghcr.io/hotosm/fair-utilities-ramp:cpu-latest` (or `:gpu-latest` via the +`BASE_IMAGE` build arg), which provides TensorFlow, GDAL, `hot_fair_utilities` +RAMP extras, and the RAMP runtime under `/app/.venv`. + +```bash +# Build targets +docker build -f models/ramp/Dockerfile --target runtime -t ramp:runtime . +docker build -f models/ramp/Dockerfile --target test -t ramp:test . +docker build -f models/ramp/Dockerfile --target inference -t ramp:inference . +``` + +## Limitations and Bias + +- Training data and baseline weights are derived from the RAMP corpus (primarily humanitarian-mapping contexts); performance on dense urban scenes with complex roof structures may be lower than on sparser rural settlements. +- The model is sensitive to imagery with strong color casts, motion blur, or significant off-nadir angle; preprocess inputs to approximately nadir RGB at the target zoom before inference. +- Binary building output from `predict()` discards the boundary/contact auxiliary classes after decoding; downstream polygonization may still merge neighbouring buildings that share a footprint edge. + +## Citation + +RAMP - Replicable AI for Microplanning. Upstream source: https://github.com/radiantearth/ramp-code + +## License + +- Model weights and code: Apache-2.0 +- Training data: ODbL-1.0 (OpenStreetMap-derived labels) diff --git a/models/ramp/pipeline.py b/models/ramp/pipeline.py new file mode 100644 index 0000000..7e0d37b --- /dev/null +++ b/models/ramp/pipeline.py @@ -0,0 +1,841 @@ +"""ZenML pipeline for RAMP (EfficientNetB0 + U-Net) building semantic segmentation.""" + +import json +import os +import re +import shutil +import tempfile +import zipfile +from pathlib import Path +from typing import Annotated, Any +from urllib.request import urlretrieve + +from zenml import log_metadata, pipeline, step + +from fair.zenml.materializers import CheckpointBytesMaterializer, ONNXMaterializer + +_DEFAULT_MODEL_CACHE = Path("/workspace/.ramp_model_cache") +_QUBVEL_EFFICIENTNET_RELEASE = "https://github.com/qubvel/efficientnet/releases/download/v0.0.1/" + + +def _to_local_path(path_value: str, purpose: str) -> Path: + """Resolve a path with UPath and ensure local filesystem semantics.""" + from upath import UPath + + upath_obj = UPath(path_value) + protocol = getattr(upath_obj, "protocol", "") or "" + if protocol not in ("", "file"): + raise NotImplementedError( + f"{purpose} requires a local filesystem path. Received protocol={protocol!r} for {path_value!r}." + ) + return Path(str(upath_obj)) + + +def _resolve_input_directory(path_value: str, purpose: str) -> Path: + """Resolve local/remote dataset directories to a local path.""" + from fair.utils.data import resolve_directory + + if "://" in str(path_value): + return resolve_directory(path_value, pattern="*") + return _to_local_path(path_value, purpose) + + +def _resolve_input_file(path_value: str, purpose: str) -> Path: + """Resolve local/remote file paths to a local path.""" + from fair.utils.data import resolve_path + + if "://" in str(path_value): + return resolve_path(path_value) + return _to_local_path(path_value, purpose) + + +def _download_and_extract_zip(zip_url: str, dest_dir: Path) -> None: + """Download a ZIP URL and extract in dest_dir.""" + dest_dir.mkdir(parents=True, exist_ok=True) + zip_name = Path(zip_url.split("/")[-1]).name or "ramp_v1.zip" + zip_path = dest_dir / zip_name + urlretrieve(zip_url, zip_path) + with zipfile.ZipFile(zip_path, "r") as zf: + zf.extractall(dest_dir) + zip_path.unlink(missing_ok=True) + + +def _extract_zip(zip_path: Path, dest_dir: Path) -> None: + dest_dir.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(zip_path, "r") as zf: + zf.extractall(dest_dir) + + +def resolve_model_href(model_uri: str, cache_dir: Path | None = None) -> str: + """Resolve .onnx, .keras, or .zip model URIs to local paths. + - .onnx -> local .onnx file path + - .keras -> local .keras file path + - .zip -> extracted SavedModel directory path (folder containing saved_model.pb) + """ + if not isinstance(model_uri, str): + raise TypeError("model_uri must be a string") + if cache_dir is not None and not isinstance(cache_dir, Path): + raise TypeError("cache_dir must be a pathlib.Path or None") + cache_dir = cache_dir or _DEFAULT_MODEL_CACHE + cache_dir.mkdir(parents=True, exist_ok=True) + is_http = model_uri.startswith(("http://", "https://")) + clean_uri = model_uri.split("?", 1)[0] + suffix = Path(clean_uri).suffix.lower() + # ONNX path + if suffix == ".onnx": + if is_http: + dest = cache_dir / (Path(clean_uri).name or "model.onnx") + if not dest.is_file(): + urlretrieve(model_uri, dest) + return str(dest) + # local onnx path + p = Path(model_uri).resolve() + if p.is_file(): + return str(p) + raise FileNotFoundError(f"ONNX model not found: {p}") + # Keras single-file checkpoint + if suffix == ".keras": + if is_http: + dest = cache_dir / (Path(clean_uri).name or "model.keras") + if not dest.is_file(): + urlretrieve(model_uri, dest) + return str(dest) + p = Path(model_uri).resolve() + if p.is_file(): + return str(p) + raise FileNotFoundError(f"Keras model not found: {p}") + + # ZIP path -> must contain saved_model.pb somewhere inside + if suffix == ".zip": + stem = Path(clean_uri).stem or "ramp_model" + dest_dir = cache_dir / stem + # cache hit + for existing in dest_dir.rglob("saved_model.pb"): + return str(existing.parent) + if is_http: + zip_path = cache_dir / (Path(clean_uri).name or "model.zip") + if not zip_path.is_file(): + urlretrieve(model_uri, zip_path) + else: + zip_path = Path(model_uri).resolve() + if not zip_path.is_file(): + raise FileNotFoundError(f"ZIP model not found: {zip_path}") + dest_dir.mkdir(parents=True, exist_ok=True) + import zipfile + + with zipfile.ZipFile(zip_path, "r") as zf: + zf.extractall(dest_dir) + for existing in dest_dir.rglob("saved_model.pb"): + return str(existing.parent) + raise RuntimeError(f"Zip from {model_uri} did not contain a valid SavedModel") + raise ValueError(f"Unsupported model format for {model_uri!r}. Only .onnx, .keras and .zip are accepted.") + + +def _ensure_ramp_baseline(base_model_weights: str, data_base_path: str | Path) -> Path: + """Return a local SavedModel directory for fine-tuning, downloading if necessary. + + ``base_model_weights`` may be an HTTP(S) .zip URL (preferred), a local .zip, + or a SavedModel directory. A pre-provisioned baseline under + ``/app/ramp-data/baseline`` (from the RAMP utilities Docker image) is used + when present and no explicit URL is provided. If the pre-provisioned baseline + is missing, callers must provide `base_model_weights` (typically from STAC). + """ + image_ck = Path("/app/ramp-data/baseline") + if not base_model_weights and (image_ck / "saved_model.pb").exists(): + return image_ck + + if not base_model_weights: + raise ValueError("RAMP baseline weights are required but were not provided. ") + + return Path(resolve_model_href(base_model_weights, cache_dir=Path(data_base_path) / ".baseline_cache")) + + +def _patch_keras_get_file_for_efficientnet_weights() -> None: + """Redirect broken Callidior EfficientNet weight URLs to qubvel's GitHub release assets. + + The ``efficientnet`` package (via ``segmentation_models``) downloads encoder weights from + ``github.com/Callidior/keras-applications/releases/...``; those assets now return **404**. + qubvel hosts compatible ``*_imagenet_1000_notop.h5`` files on the same model's releases. + """ + import tensorflow as tf + + ku = tf.keras.utils + if hasattr(ku.get_file, "_ramp_efficientnet_mirror"): + return + + original = ku.get_file + + def _get_file(fname, origin, *args, **kwargs): + if isinstance(origin, str) and "Callidior" in origin: + m = re.match( + r"^(efficientnet-b\d+)_weights_tf_dim_ordering_tf_kernels_autoaugment_notop\.h5$", + fname or "", + ) + if m: + alt = f"{m.group(1)}_imagenet_1000_notop.h5" + origin = f"{_QUBVEL_EFFICIENTNET_RELEASE}{alt}" + kwargs = dict(kwargs) + kwargs["file_hash"] = None + return original(fname, origin, *args, **kwargs) + + _get_file._ramp_efficientnet_mirror = True # type: ignore[attr-defined] + ku.get_file = _get_file + + +def _materialize_training_input(dataset_chips: str, dataset_labels: str, work_dir: Path) -> Path: + """Create the preprocess input folder with PNG chips and a single labels.geojson. + + Accepts .tif/.tiff/.png chip files on disk; TIFFs are converted to 3-band PNGs while + preserving filename stem (e.g. OAM-{x}-{y}-{z}.png) so hot_fair_utilities label clipping + can parse tile ids. + """ + chips_dir = _resolve_input_directory(dataset_chips, "dataset_chips") + labels_path = _resolve_input_file(dataset_labels, "dataset_labels") + + input_dir = work_dir / "input" + if input_dir.exists(): + shutil.rmtree(input_dir) + input_dir.mkdir(parents=True, exist_ok=True) + + tif_paths = sorted(list(chips_dir.glob("*.tif")) + list(chips_dir.glob("*.tiff"))) + png_paths = sorted(chips_dir.glob("*.png")) + + if tif_paths: + import numpy as np + import rasterio + from PIL import Image + + for tif_path in tif_paths: + png_path = input_dir / (tif_path.stem + ".png") + with rasterio.open(tif_path) as src: + data = src.read() + if data.shape[0] < 3: + continue + rgb = np.transpose(data[:3], (1, 2, 0)) + rgb = (rgb * 255).astype(np.uint8) if rgb.max() <= 1.0 else np.clip(rgb, 0, 255).astype(np.uint8) + Image.fromarray(rgb).save(png_path) + + for png_path in png_paths: + shutil.copy2(png_path, input_dir / png_path.name) + + if not list(input_dir.glob("*.png")): + raise FileNotFoundError(f"No train chips (.tif/.tiff/.png) found in {chips_dir}") + + if not labels_path.is_file(): + raise FileNotFoundError(f"dataset_labels file not found: {labels_path}") + shutil.copy2(labels_path, input_dir / "labels.geojson") + return input_dir + + +def preprocess( + input_path: str, + output_path: str, + boundary_width: int = 3, + contact_spacing: int = 8, +) -> str: + """Preprocess OAM PNG chips + labels into RAMP 4-class multimasks. + + Emits: + - preprocessed/chips/*.tif (georeferenced RGB chips, EPSG:3857) + - preprocessed/labels/*.geojson (per-chip labels, reprojected + clipped) + - preprocessed/multimasks/*.mask.tif (4-class sparse categorical masks) + """ + from hot_fair_utilities import preprocess as _preprocess + + local_input = _resolve_input_directory(input_path, "input_path") + _preprocess( + input_path=str(local_input), + output_path=output_path, + rasterize=True, + rasterize_options=["binary"], + georeference_images=True, + multimasks=True, + input_boundary_width=boundary_width, + input_contact_spacing=contact_spacing, + ) + return output_path + + +def postprocess(prediction_masks_dir: str, output_dir: str) -> dict[str, Any]: + """Merge prediction TIFF tiles into a building-footprint GeoJSON (EPSG:4326).""" + from geomltoolkits.geometry.validate import validate_polygon_geometries + from geomltoolkits.raster.merge import merge_rasters + from geomltoolkits.raster.morphology import morphological_cleaning + from geomltoolkits.raster.vectorize import vectorize_mask + + pred_dir = _to_local_path(prediction_masks_dir, "prediction_masks_dir") + out_dir = _to_local_path(output_dir, "output_dir") + out_dir.mkdir(parents=True, exist_ok=True) + + pred_tifs = sorted(pred_dir.glob("*.tif")) + if not pred_tifs: + return {"type": "FeatureCollection", "features": []} + + merged_mask_path = out_dir / "merged_prediction_mask.tif" + merged_geojson_path = out_dir / "predictions.geojson" + + merge_rasters(str(pred_dir), str(merged_mask_path)) + morphological_cleaning(str(merged_mask_path)) + gdf = vectorize_mask( + input_tiff=str(merged_mask_path), + output_geojson=str(merged_geojson_path), + simplify_tolerance=0.5, + min_area=3, + orthogonalize=True, + ortho_skew_tolerance_deg=15, + ortho_max_angle_change_deg=15, + ) + + geojson_dict = json.loads(gdf.to_json()) + if not geojson_dict.get("features"): + if not merged_geojson_path.is_file(): + merged_geojson_path.write_text(json.dumps(geojson_dict), encoding="utf-8") + return geojson_dict + + if gdf.crs and gdf.crs != "EPSG:4326": + gdf = gdf.to_crs("EPSG:4326") + elif not gdf.crs: + gdf.set_crs("EPSG:3857", inplace=True) + gdf = gdf.to_crs("EPSG:4326") + + import pandas as pd + + for col in gdf.columns: + if col == "geometry": + continue + if pd.api.types.is_extension_array_dtype(gdf[col].dtype): + gdf[col] = gdf[col].astype(object).where(gdf[col].notna(), None) + gdf.to_file(merged_geojson_path, driver="GeoJSON") + + validated_geojson = validate_polygon_geometries( + geojson_dict, + output_path=str(merged_geojson_path), + ) + if isinstance(validated_geojson, str) and Path(validated_geojson).is_file(): + if Path(validated_geojson) != merged_geojson_path: + shutil.copy2(validated_geojson, merged_geojson_path) + return json.loads(merged_geojson_path.read_text(encoding="utf-8")) + + if isinstance(validated_geojson, dict): + return validated_geojson + return json.loads(gdf.to_json()) + + +def _prepare_training_split( + dataset_chips: str, + dataset_labels: str, + hyperparameters: dict[str, Any], + force_rebuild: bool = False, +) -> dict[str, Any]: + """Preprocess + split chips/labels into RAMP train/val layout; return split_info.""" + from hot_fair_utilities.training.ramp.prepare_data import split_training_2_validation + + val_fraction = float( + hyperparameters.get( + "training.val_ratio", + hyperparameters.get("val_fraction", hyperparameters.get("val_ratio", 0.15)), + ) + ) + if not 0.0 < val_fraction < 1.0: + raise ValueError("val_fraction must be in (0.0, 1.0)") + boundary_width = int(hyperparameters.get("training.boundary_width", hyperparameters.get("boundary_width", 3))) + contact_spacing = int(hyperparameters.get("training.contact_spacing", hyperparameters.get("contact_spacing", 8))) + seed = int(hyperparameters.get("training.split_seed", hyperparameters.get("split_seed", 42))) + + work_dir = Path(tempfile.mkdtemp(prefix="ramp_training_")) + preprocessed_dir = work_dir / "preprocessed" + ramp_train_dir = work_dir / "ramp_training_work" + + work_dir.mkdir(parents=True, exist_ok=True) + input_dir = _materialize_training_input(dataset_chips, dataset_labels, work_dir) + preprocess(str(input_dir), str(preprocessed_dir), boundary_width, contact_spacing) + split_training_2_validation(str(preprocessed_dir), str(ramp_train_dir), multimasks=True) + + train_count = len(list((ramp_train_dir / "chips").glob("*.tif"))) + val_count = len(list((ramp_train_dir / "val-chips").glob("*.tif"))) + + return { + "strategy": "random", + "val_ratio": val_fraction, + "seed": seed, + "train_count": train_count, + "val_count": val_count, + "description": "Preprocess chips+labels into 4-class multimasks, then random train/val split.", + "_work_dir": str(work_dir), + "_preprocessed_dir": str(preprocessed_dir), + "_ramp_train_dir": str(ramp_train_dir), + } + + +def train_ramp_model( + ramp_train_dir: str, + base_model_weights: str, + hyperparameters: dict[str, Any], + data_base_path: str | None = None, +) -> Path: + """Fine-tune EfficientNetB0 + U-Net and return the best SavedModel directory path.""" + # run_training reads RAMP_HOME at import time; set it first so saved_model lookups resolve. + data_base_path = str(Path(data_base_path).resolve()) if data_base_path else str(Path(ramp_train_dir).resolve()) + image_baseline_ck = Path("/app/ramp-data/baseline/saved_model.pb") + if image_baseline_ck.exists(): + os.environ["RAMP_HOME"] = "/app" + else: + os.environ["RAMP_HOME"] = data_base_path + + os.environ.setdefault("SM_FRAMEWORK", "tf.keras") + + _patch_keras_get_file_for_efficientnet_weights() + import segmentation_models as sm + from hot_fair_utilities.training.ramp.cleanup import extract_highest_accuracy_model + from hot_fair_utilities.training.ramp.config import RAMP_CONFIG + from hot_fair_utilities.training.ramp.run_training import ( + manage_fine_tuning_config, + run_main_train_code, + ) + + sm.set_framework("tf.keras") + + epochs = int(hyperparameters.get("training.epochs", hyperparameters.get("epochs", RAMP_CONFIG["num_epochs"]))) + batch_size = int( + hyperparameters.get("training.batch_size", hyperparameters.get("batch_size", RAMP_CONFIG["batch_size"])) + ) + backbone = str( + hyperparameters.get( + "training.backbone", + hyperparameters.get("backbone", RAMP_CONFIG["model"]["model_fn_parms"]["backbone"]), + ) + ) + learning_rate = float( + hyperparameters.get( + "training.learning_rate", + hyperparameters.get("learning_rate", RAMP_CONFIG["optimizer"]["optimizer_fn_parms"]["learning_rate"]), + ) + ) + patience = int( + hyperparameters.get( + "training.early_stopping_patience", + hyperparameters.get( + "early_stopping_patience", RAMP_CONFIG["early_stopping"]["early_stopping_parms"]["patience"] + ), + ) + ) + + cfg = manage_fine_tuning_config(ramp_train_dir, epochs, batch_size, freeze_layers=False, multimasks=True) + cfg["model"]["model_fn_parms"]["backbone"] = backbone + cfg["optimizer"]["optimizer_fn_parms"]["learning_rate"] = learning_rate + cfg["early_stopping"]["early_stopping_parms"]["patience"] = patience + + if cfg["saved_model"]["use_saved_model"]: + baseline_dir = _ensure_ramp_baseline(base_model_weights, data_base_path) + cfg["saved_model"]["use_saved_model"] = (baseline_dir / "saved_model.pb").exists() + + run_main_train_code(cfg) + _final_accuracy, final_model_path = extract_highest_accuracy_model(ramp_train_dir) + return Path(final_model_path) + + +def _materialize_keras_bytes(blob: bytes) -> Path: + """Write Keras `.keras` bytes to a temp file and return its path.""" + dest = Path(tempfile.mkdtemp(prefix="ramp_keras_")) / "model.keras" + dest.write_bytes(blob) + return dest + + +def _restore_checkpoint(trained_model: Any) -> Path: + """Restore a trained RAMP Keras checkpoint as a local `.keras` file path.""" + if isinstance(trained_model, bytes): + return _materialize_keras_bytes(trained_model) + if isinstance(trained_model, (str, Path)): + p = Path(str(trained_model)) + if p.is_file() and p.suffix.lower() == ".keras": + return p + return Path(resolve_model_href(str(p))) + raise TypeError(f"Cannot restore RAMP checkpoint from {type(trained_model).__name__}") + + +def _convert_savedmodel_to_onnx_bytes(saved_model_dir: Path, opset: int = 13) -> bytes: + """Convert a TF SavedModel directory to ONNX bytes via tf2onnx.""" + import tensorflow as tf + import tf2onnx + + with tempfile.TemporaryDirectory() as tmp: + onnx_path = Path(tmp) / "model.onnx" + from_saved_model = getattr(tf2onnx.convert, "from_saved_model", None) + if callable(from_saved_model): + from_saved_model( + str(saved_model_dir), + output_path=str(onnx_path), + opset=opset, + ) + else: + model = tf.keras.models.load_model(str(saved_model_dir), compile=False) + tf2onnx.convert.from_keras( + model, + opset=opset, + output_path=str(onnx_path), + ) + return onnx_path.read_bytes() + + +def _build_feature_collection(features: list[dict[str, Any]]) -> dict[str, Any]: + return {"type": "FeatureCollection", "features": features} + + +def _extract_ramp_shapes(session: Any) -> tuple[str, int, int]: + """Return (input_name, input_height, input_width) for a RAMP ONNX session (NHWC).""" + input_meta = session.get_inputs()[0] + shape = input_meta.shape + if len(shape) != 4: + raise RuntimeError(f"Unexpected ONNX input shape: {shape}") + # RAMP ONNX is [batch, H, W, bands] (channels last) — keep this order. + height = int(shape[1]) if isinstance(shape[1], int) and shape[1] > 0 else 256 + width = int(shape[2]) if isinstance(shape[2], int) and shape[2] > 0 else 256 + return input_meta.name, height, width + + +def _prepare_onnx_image(img_path: Path, input_height: int, input_width: int) -> tuple[Any, Any, Any]: + """Load an RGB chip and return (batch, transform, meta) for ONNX inference.""" + import numpy as np + import rasterio + from PIL import Image + + with rasterio.open(img_path) as src: + arr = src.read([1, 2, 3]).astype(np.float32) / 255.0 + transform = src.transform + crs = src.crs + src_height = src.height + src_width = src.width + + resized = [ + np.asarray(Image.fromarray(arr[c]).resize((input_width, input_height), Image.Resampling.BILINEAR)) + for c in range(arr.shape[0]) + ] + hwc = np.stack(resized, axis=-1).astype(np.float32) # NHWC + batch = hwc[np.newaxis, ...] + return batch, transform, (src_width, src_height, input_width, input_height, crs) + + +def _decode_ramp_building_mask( + output: Any, + input_height: int, + input_width: int, + src_height: int, + src_width: int, + min_class_value: int = 1, +) -> Any: + """Decode a RAMP ONNX output tensor into a (src_height, src_width) uint8 building mask. + + Accepts either a 4-class softmax/logits tensor [B,H,W,C] or a [B,H,W,1] class-index tensor. + Pixels with class >= ``min_class_value`` (default 1=building) become 1; others 0. + """ + import numpy as np + from PIL import Image + + arr = np.asarray(output) + if arr.ndim == 5: + arr = arr[0] + if arr.ndim == 4 and arr.shape[0] == 1: + arr = arr[0] + if arr.ndim == 3 and arr.shape[-1] > 1: + class_idx = arr.argmax(axis=-1) + elif arr.ndim == 3 and arr.shape[-1] == 1: + class_idx = arr[..., 0] + elif arr.ndim == 2: + class_idx = arr + else: + raise RuntimeError(f"Unexpected RAMP ONNX output shape: {arr.shape}") + + class_idx = np.asarray(class_idx).astype(np.int32) + # Collapse multiclass (background=0, building=1, boundary=2, contact=3) to binary building. + binary = (class_idx == min_class_value).astype(np.uint8) + if binary.shape != (src_height, src_width): + resized = Image.fromarray(binary * 255).resize((src_width, src_height), Image.Resampling.NEAREST) + binary = (np.asarray(resized) > 127).astype(np.uint8) + return binary + + +def _vectorize_binary_mask(mask: Any, transform: Any, crs: Any, confidence: float) -> list[dict[str, Any]]: + import numpy as np + import rasterio.features + from pyproj import Transformer + + mask_uint8 = np.asarray(mask).astype(np.uint8) + transformer = Transformer.from_crs(crs, "EPSG:4326", always_xy=True) if crs and str(crs) != "EPSG:4326" else None + + features: list[dict[str, Any]] = [] + for geom, value in rasterio.features.shapes(mask_uint8, transform=transform): + if int(value) < 1: + continue + if transformer is not None: + coords = geom["coordinates"] + geom["coordinates"] = [[list(transformer.transform(x, y)) for x, y in ring] for ring in coords] + features.append( + { + "type": "Feature", + "properties": {"class": 1, "confidence": round(confidence, 4)}, + "geometry": geom, + } + ) + return features + + +def predict(session: Any, input_images: str, params: dict[str, Any]) -> dict[str, Any]: + """Run RAMP ONNX inference and return a FeatureCollection of building polygons. + + Required by ``fair.serve.base``. ``session`` is an ``onnxruntime.InferenceSession`` built by + the serving layer from the STAC ``assets.model`` ONNX artifact. + """ + from fair.utils.data import resolve_directory + + confidence_threshold = float(params.get("confidence_threshold", 0.5)) + min_class_value = int(params.get("min_class_value", 1)) + + input_name, input_height, input_width = _extract_ramp_shapes(session) + input_dir = resolve_directory(input_images) + patterns = ("*.png", "*.tif", "*.tiff", "*.jpg") + img_paths = sorted(p for pat in patterns for p in input_dir.glob(pat)) + if not img_paths: + raise FileNotFoundError(f"No input images found in {input_dir}") + + features: list[dict[str, Any]] = [] + for img_path in img_paths: + batch, transform, meta = _prepare_onnx_image(img_path, input_height, input_width) + src_width, src_height, _iw, _ih, crs = meta + outputs = session.run(None, {input_name: batch}) + if not outputs: + continue + mask = _decode_ramp_building_mask( + outputs[0], + input_height=input_height, + input_width=input_width, + src_height=src_height, + src_width=src_width, + min_class_value=min_class_value, + ) + features.extend(_vectorize_binary_mask(mask, transform, crs, confidence_threshold)) + return _build_feature_collection(features) + + +@step +def split_dataset( + dataset_chips: str, + dataset_labels: str, + hyperparameters: dict[str, Any], +) -> Annotated[dict[str, Any], "split_info"]: + """Preprocess chips+labels and create the RAMP train/val layout.""" + split_info = _prepare_training_split(dataset_chips, dataset_labels, hyperparameters) + log_metadata(metadata={"fair/split": {k: v for k, v in split_info.items() if not k.startswith("_")}}) + return split_info + + +@step(output_materializers={"trained_model": CheckpointBytesMaterializer}) +def train_model( + dataset_chips: str, + dataset_labels: str, + base_model_weights: str, + hyperparameters: dict[str, Any], + split_info: dict[str, Any], + num_classes: int = 4, + model_name: str | None = None, + base_model_id: str | None = None, + dataset_id: str | None = None, +) -> Annotated[bytes, "trained_model"]: + """Fine-tune RAMP EfficientNetB0 U-Net; return the best model as `.keras` bytes.""" + _ = (num_classes, model_name, base_model_id, dataset_id) + + ramp_train_dir = Path(split_info["_ramp_train_dir"]) + if not ramp_train_dir.exists(): + split_info = _prepare_training_split(dataset_chips, dataset_labels, hyperparameters, force_rebuild=True) + ramp_train_dir = Path(split_info["_ramp_train_dir"]) + + work_dir = split_info.get("_work_dir") or str(ramp_train_dir.parent) + final_model_path = train_ramp_model( + ramp_train_dir=str(ramp_train_dir), + base_model_weights=base_model_weights, + hyperparameters=hyperparameters, + data_base_path=work_dir, + ) + import tensorflow as tf + + model = tf.keras.models.load_model(str(final_model_path), compile=False) + exported_path = Path(tempfile.mkdtemp(prefix="ramp_keras_export_")) / "model.keras" + model.save(str(exported_path)) + blob = exported_path.read_bytes() + log_metadata(metadata={"keras_path": str(exported_path), "checkpoint_bytes": len(blob)}) + return blob + + +@step +def evaluate_model( + trained_model: Any, + dataset_chips: str, + dataset_labels: str, + hyperparameters: dict[str, Any], + split_info: dict[str, Any], + class_names: list[str] | None = None, +) -> Annotated[dict[str, Any], "metrics"]: + """Compute per-pixel building-class metrics on the validation split.""" + _ = class_names + + import numpy as np + import rasterio + + ramp_train_dir = Path(split_info.get("_ramp_train_dir", "")) + if not ramp_train_dir.exists(): + split_info = _prepare_training_split(dataset_chips, dataset_labels, hyperparameters, force_rebuild=True) + ramp_train_dir = Path(split_info["_ramp_train_dir"]) + + val_chips_dir = ramp_train_dir / "val-chips" + val_masks_dir = ramp_train_dir / "val-multimasks" + pairs: list[tuple[Path, Path]] = [] + for chip in sorted(val_chips_dir.glob("*.tif")): + mask = val_masks_dir / f"{chip.stem}.mask.tif" + if mask.is_file(): + pairs.append((chip, mask)) + + restored = _restore_checkpoint(trained_model) + + if not pairs: + # No val data to evaluate against (e.g. CI mocks); return zeroed metrics with the + # required fair:* keys so downstream validation still sees the expected schema. + zero_metrics: dict[str, Any] = { + "fair:accuracy": 0.0, + "fair:mean_iou": 0.0, + "fair:precision": 0.0, + "fair:recall": 0.0, + } + log_metadata(metadata=zero_metrics) + return zero_metrics + + import tensorflow as tf + + model = tf.keras.models.load_model(str(restored), compile=False) + + tp = fp = fn = 0 + correct = total = 0 + for chip_path, mask_path in pairs: + with rasterio.open(chip_path) as src: + chip = src.read([1, 2, 3]).astype(np.float32) / 255.0 + with rasterio.open(mask_path) as src: + gt = src.read(1).astype(np.int32) + batch = np.transpose(chip, (1, 2, 0))[np.newaxis, ...] + pred = model.predict(batch, verbose=0) + pred_arr = np.asarray(pred) + if pred_arr.ndim == 4 and pred_arr.shape[-1] > 1: + pred_idx = pred_arr[0].argmax(axis=-1) + elif pred_arr.ndim == 4 and pred_arr.shape[-1] == 1: + pred_idx = pred_arr[0, ..., 0].astype(np.int32) + else: + pred_idx = pred_arr[0].astype(np.int32) + + gt_bin = (gt == 1).astype(np.uint8) + pr_bin = (pred_idx == 1).astype(np.uint8) + tp += int(((gt_bin == 1) & (pr_bin == 1)).sum()) + fp += int(((gt_bin == 0) & (pr_bin == 1)).sum()) + fn += int(((gt_bin == 1) & (pr_bin == 0)).sum()) + correct += int((pred_idx == gt).sum()) + total += int(gt.size) + + precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 + iou = tp / (tp + fp + fn) if (tp + fp + fn) > 0 else 0.0 + accuracy = correct / total if total > 0 else 0.0 + + metrics_dict: dict[str, Any] = { + "fair:accuracy": float(accuracy), + "fair:mean_iou": float(iou), + "fair:precision": float(precision), + "fair:recall": float(recall), + } + log_metadata(metadata=metrics_dict) + return metrics_dict + + +@step(output_materializers={"onnx_model": ONNXMaterializer}) +def export_onnx(trained_model: Any) -> Annotated[bytes, "onnx_model"]: + """Convert the trained RAMP Keras checkpoint to ONNX bytes and validate.""" + import onnx + import tensorflow as tf + import tf2onnx + + restored = _restore_checkpoint(trained_model) + model = tf.keras.models.load_model(str(restored), compile=False) + with tempfile.TemporaryDirectory() as tmp: + onnx_path = Path(tmp) / "model.onnx" + tf2onnx.convert.from_keras(model, opset=13, output_path=str(onnx_path)) + onnx_bytes = onnx_path.read_bytes() + onnx.checker.check_model(onnx.load_from_string(onnx_bytes)) + log_metadata(metadata={"onnx_bytes": len(onnx_bytes)}) + return onnx_bytes + + +@step +def run_inference( + model_uri: str, + input_images: str, + inference_params: dict[str, Any], +) -> Annotated[dict[str, Any], "predictions"]: + """RAMP batch inference using the promoted ONNX `assets.model`.""" + from fair.serve.base import load_session + + session = load_session(model_uri) + return predict(session, input_images, inference_params) + + +@step +def run_preprocessing( + input_path: str, + output_path: str, + boundary_width: int = 3, + contact_spacing: int = 8, +) -> str: + """STAC entrypoint wrapper for RAMP preprocessing.""" + return preprocess(input_path, output_path, boundary_width, contact_spacing) + + +@step +def run_postprocessing(prediction_path: str, output_dir: str) -> dict[str, Any]: + """STAC entrypoint wrapper for RAMP postprocessing.""" + return postprocess(prediction_path, output_dir) + + +@pipeline +def training_pipeline( + base_model_weights: str, + dataset_chips: str, + dataset_labels: str, + num_classes: int, + hyperparameters: dict[str, Any], +) -> None: + """RAMP training pipeline: split → train → evaluate → export ONNX.""" + split_info = split_dataset( + dataset_chips=dataset_chips, + dataset_labels=dataset_labels, + hyperparameters=hyperparameters, + ) + trained_model = train_model( + dataset_chips=dataset_chips, + dataset_labels=dataset_labels, + base_model_weights=base_model_weights, + hyperparameters=hyperparameters, + split_info=split_info, + num_classes=num_classes, + ) + evaluate_model( + trained_model=trained_model, + dataset_chips=dataset_chips, + dataset_labels=dataset_labels, + hyperparameters=hyperparameters, + split_info=split_info, + ) + export_onnx(trained_model=trained_model) + + +@pipeline +def inference_pipeline( + model_uri: str, + input_images: str, + inference_params: dict[str, Any], +) -> None: + """RAMP inference pipeline: load ONNX session → predict → FeatureCollection.""" + run_inference(model_uri=model_uri, input_images=input_images, inference_params=inference_params) diff --git a/models/ramp/stac-item.json b/models/ramp/stac-item.json new file mode 100644 index 0000000..16fed4b --- /dev/null +++ b/models/ramp/stac-item.json @@ -0,0 +1,374 @@ +{ + "type": "Feature", + "stac_version": "1.1.0", + "stac_extensions": [ + "https://stac-extensions.github.io/mlm/v1.5.1/schema.json", + "https://stac-extensions.github.io/version/v1.2.0/schema.json", + "https://stac-extensions.github.io/classification/v2.0.0/schema.json", + "https://stac-extensions.github.io/file/v2.1.0/schema.json", + "https://stac-extensions.github.io/raster/v1.1.0/schema.json", + "https://hotosm.github.io/fAIr-models/schemas/v1.0.0/base-model/schema.json" + ], + "id": "ramp-v1", + "geometry": { + "type": "Polygon", + "coordinates": [ + [ + [-180, -90], + [180, -90], + [180, 90], + [-180, 90], + [-180, -90] + ] + ] + }, + "bbox": [-180, -90, 180, 90], + "properties": { + "datetime": "2024-01-01T00:00:00Z", + "created": "2024-01-01T00:00:00Z", + "updated": "2024-01-01T00:00:00Z", + "title": "RAMP EfficientNetB0 + U-Net Building Segmentation", + "description": "RAMP semantic-segmentation base model: EfficientNetB0 encoder with U-Net decoder outputs 4-class sparse-categorical masks (background, building, boundary, contact) from RGB aerial imagery. Packaged for fAIr finetuning and ONNX-based inference.", + "mlm:name": "ramp-v1", + "mlm:architecture": "EffUnet", + "mlm:tasks": [ + "semantic-segmentation" + ], + "mlm:framework": "TensorFlow", + "mlm:framework_version": "2.15.1", + "mlm:pretrained": true, + "mlm:pretrained_source": "https://huggingface.co/hotosm/ramp/resolve/74daea54694f2e4924f1222520c614c7f5c029fe/v1-baseline.zip", + "mlm:accelerator": "amd64", + "mlm:accelerator_constrained": false, + "mlm:total_parameters": 7500000, + "mlm:memory_size": 536870912, + "keywords": [ + "building", + "semantic-segmentation", + "polygon" + ], + "license": "Apache-2.0", + "version": "1", + "deprecated": false, + "providers": [ + { + "name": "HOTOSM", + "roles": [ + "producer", + "host" + ], + "url": "https://www.hotosm.org", + "description": "Humanitarian OpenStreetMap Team" + }, + { + "name": "Development Seed", + "roles": [ + "producer" + ], + "url": "https://developmentseed.org", + "description": "Original authors of the RAMP model." + } + ], + "fair:metrics_spec": [ + { + "key": "fair:accuracy", + "name": "Pixel accuracy", + "description": "Per-pixel classification accuracy across all 4 classes on the validation split." + }, + { + "key": "fair:mean_iou", + "name": "Building IoU", + "description": "Intersection-over-Union of the building class (class 1) on the validation split." + }, + { + "key": "fair:precision", + "name": "Building precision", + "description": "Per-pixel precision of the building class (class 1) on the validation split." + }, + { + "key": "fair:recall", + "name": "Building recall", + "description": "Per-pixel recall of the building class (class 1) on the validation split." + } + ], + "fair:split_spec": { + "strategy": "random", + "default_ratio": 0.15, + "seed": 42, + "description": "Preprocess chips + labels into 4-class multimasks, then seeded random train/validation split via hot_fair_utilities.split_training_2_validation." + }, + "fair:hyperparameters_spec": [ + { + "key": "backbone", + "type": "str", + "default": "efficientnetb0", + "description": "Encoder backbone for U-Net segmentation model" + }, + { + "key": "epochs", + "type": "int", + "default": 20, + "min": 1, + "max": 20, + "description": "Number of training epochs" + }, + { + "key": "batch_size", + "type": "int", + "default": 8, + "min": 1, + "max": 16, + "description": "Samples per training batch" + }, + { + "key": "learning_rate", + "type": "float", + "default": 0.0003, + "min": 1e-7, + "max": 1.0, + "description": "Optimizer learning rate" + }, + { + "key": "val_ratio", + "type": "float", + "default": 0.15, + "min": 0.05, + "max": 0.5, + "description": "Fraction of data held out for validation" + }, + { + "key": "split_seed", + "type": "int", + "default": 42, + "min": 0, + "max": 2147483647, + "description": "Random seed for reproducible train/val split" + }, + { + "key": "boundary_width", + "type": "int", + "default": 3, + "min": 1, + "max": 10, + "description": "Width of building boundary class in multimasks" + }, + { + "key": "contact_spacing", + "type": "int", + "default": 8, + "min": 1, + "max": 20, + "description": "Spacing for contact points between neighbouring buildings" + }, + { + "key": "early_stopping_patience", + "type": "int", + "default": 10, + "min": 1, + "max": 50, + "description": "Epochs to wait before stopping if validation metric does not improve" + }, + { + "key": "augmentation", + "type": "bool", + "default": false, + "description": "Enable data augmentation during training" + }, + { + "key": "imgsz", + "type": "int", + "default": 256, + "min": 128, + "max": 512, + "description": "Input image size (square chips)" + }, + { + "key": "num_classes", + "type": "int", + "default": 4, + "min": 2, + "max": 10, + "description": "Number of output classes (background, building, boundary, contact)" + }, + { + "key": "confidence_threshold", + "type": "float", + "default": 0.5, + "min": 0.0, + "max": 1.0, + "description": "Minimum per-pixel confidence to retain a predicted building class at inference" + }, + { + "key": "min_class_value", + "type": "int", + "default": 1, + "min": 0, + "max": 255, + "description": "Minimum class index value to retain as a foreground prediction at inference" + } + ], + "inference_time_per_256x256": { + "value": 45, + "unit": "ms" + }, + "mlm:input": [ + { + "name": "RGB GeoTIFF chips", + "bands": [ + {"name": "red"}, + {"name": "green"}, + {"name": "blue"} + ], + "input": { + "shape": [-1, 256, 256, 3], + "dim_order": ["batch", "height", "width", "bands"], + "data_type": "float32" + }, + "norm_by_channel": false, + "norm_clip": [0.0, 1.0], + "pre_processing_function": { + "format": "python", + "expression": "models.ramp.pipeline:preprocess" + } + } + ], + "mlm:output": [ + { + "name": "4-class sparse-categorical multimask", + "tasks": ["semantic-segmentation"], + "result": { + "shape": [-1, 256, 256, 4], + "dim_order": ["batch", "height", "width", "classes"], + "data_type": "float32" + }, + "classification:classes": [ + { + "name": "background", + "value": 0, + "description": "Non-building pixels." + }, + { + "name": "building", + "value": 1, + "description": "Building interior pixels." + }, + { + "name": "boundary", + "value": 2, + "description": "Building boundary pixels (helps separate adjacent buildings)." + }, + { + "name": "contact", + "value": 3, + "description": "Close-contact points between neighbouring buildings." + } + ], + "post_processing_function": { + "format": "python", + "expression": "models.ramp.pipeline:postprocess" + }, + "bands": [] + } + ], + "mlm:hyperparameters": { + "training.backbone": "efficientnetb0", + "training.epochs": 20, + "training.batch_size": 8, + "training.learning_rate": 0.0003, + "training.val_ratio": 0.15, + "training.split_seed": 42, + "training.boundary_width": 3, + "training.contact_spacing": 8, + "training.early_stopping_patience": 10, + "training.augmentation": false, + "training.imgsz": 256, + "training.num_classes": 4, + "inference.confidence_threshold": 0.5, + "inference.min_class_value": 1 + }, + "licenses": { + "model": "Apache-2.0", + "data": "ODbL-1.0", + "code": "Apache-2.0" + } + }, + "assets": { + "checkpoint": { + "href": "https://huggingface.co/hotosm/ramp/resolve/74daea54694f2e4924f1222520c614c7f5c029fe/v1-baseline.zip", + "title": "RAMP EfficientNetB0 + U-Net baseline weights (zipped TF SavedModel)", + "type": "application/zip; framework=tensorflow", + "roles": [ + "mlm:model", + "mlm:weights" + ], + "mlm:artifact_type": "tf.keras.Model.save", + "raster:bands": [ + {"name": "red"}, + {"name": "green"}, + {"name": "blue"} + ] + }, + "model": { + "href": "https://huggingface.co/hotosm/ramp/resolve/83c77a7e5feb3af62e3604d7bb96c6c6e9ff1a96/ramp-v1.onnx", + "title": "Portable ONNX inference model", + "type": "application/octet-stream; framework=onnx", + "roles": [ + "mlm:model", + "mlm:compiled" + ], + "mlm:artifact_type": "onnx", + "raster:bands": [ + {"name": "red"}, + {"name": "green"}, + {"name": "blue"} + ] + }, + "source-code": { + "href": "https://github.com/hotosm/fAIr-models/tree/main/models/ramp", + "type": "text/html", + "title": "RAMP model pack source", + "roles": [ + "code" + ], + "mlm:entrypoint": "models.ramp.pipeline:training_pipeline" + }, + "mlm:training": { + "href": "ghcr.io/hotosm/fair-models/ramp:v1", + "type": "application/vnd.oci.image.index.v1+json", + "title": "Training Docker Image", + "roles": [ + "mlm:training-runtime" + ] + }, + "mlm:inference": { + "href": "ghcr.io/hotosm/fair-models/ramp:v1", + "type": "application/vnd.oci.image.index.v1+json", + "title": "Inference Docker Image", + "roles": [ + "mlm:inference-runtime" + ] + }, + "readme": { + "href": "https://raw.githubusercontent.com/hotosm/fAIr-models/refs/heads/main/models/ramp/README.md", + "type": "text/markdown", + "roles": [ + "metadata" + ], + "title": "Model README" + } + }, + "links": [ + { + "rel": "license", + "href": "https://www.apache.org/licenses/LICENSE-2.0", + "type": "text/html", + "title": "Apache License 2.0" + }, + { + "rel": "cite-as", + "href": "https://github.com/radiantearth/ramp-code", + "type": "text/html", + "title": "RAMP — Replicable AI for Microplanning" + } + ] +} diff --git a/models/ramp/tests/conftest.py b/models/ramp/tests/conftest.py new file mode 100644 index 0000000..4086e38 --- /dev/null +++ b/models/ramp/tests/conftest.py @@ -0,0 +1,129 @@ +"""Per-model test fixtures: toy OAM-tiled RAMP dataset (chips + labels). + +``create_toy_data`` is importable by ``models/test_integration.py`` to materialize a +minimal 4-class semantic-segmentation dataset. Chip TIFFs follow the OAM naming +convention ``OAM-{x}-{y}-{z}.tif`` so ``hot_fair_utilities.preprocessing.clip_labels`` +can recover the tile bounds. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import mercantile +import numpy as np +import pytest +import rasterio +from rasterio.crs import CRS +from rasterio.transform import from_bounds + +CHIP_COUNT = 6 +CHIP_SIZE = 128 +_OAM_ZOOM = 18 +_TILE_XY = [(100, 200), (101, 200), (102, 200), (100, 201), (101, 201), (102, 201)] + +_bounds = mercantile.bounds(mercantile.Tile(x=_TILE_XY[0][0], y=_TILE_XY[0][1], z=_OAM_ZOOM)) +_WEST, _SOUTH, _EAST, _NORTH = _bounds.west, _bounds.south, _bounds.east, _bounds.north +_GEOMETRY = { + "type": "Polygon", + "coordinates": [[[_WEST, _SOUTH], [_EAST, _SOUTH], [_EAST, _NORTH], [_WEST, _NORTH], [_WEST, _SOUTH]]], +} +_BBOX = [_WEST, _SOUTH, _EAST, _NORTH] + + +def create_toy_data(root: Path) -> dict[str, Path]: + chips_dir = root / "chips" + chips_dir.mkdir() + + labels_dir = root / "labels" + labels_dir.mkdir() + + features: list[dict[str, Any]] = [] + for i in range(CHIP_COUNT): + tx, ty = _TILE_XY[i] + tile = mercantile.Tile(x=tx, y=ty, z=_OAM_ZOOM) + b = mercantile.bounds(tile) + transform = from_bounds(b.west, b.south, b.east, b.north, CHIP_SIZE, CHIP_SIZE) + chip_name = f"OAM-{tx}-{ty}-{_OAM_ZOOM}.tif" + with rasterio.open( + chips_dir / chip_name, + "w", + driver="GTiff", + width=CHIP_SIZE, + height=CHIP_SIZE, + count=3, + dtype="uint8", + crs=CRS.from_epsg(4326), + transform=transform, + ) as dst: + dst.write(np.random.randint(0, 255, (3, CHIP_SIZE, CHIP_SIZE), dtype=np.uint8)) + + # Tiny square polygon well inside each chip footprint. + w, h = b.east - b.west, b.north - b.south + cx = b.west + w * 0.35 + cy = b.south + h * 0.35 + s = min(w, h) * 0.25 + features.append( + { + "type": "Feature", + "properties": {"building": 1}, + "geometry": { + "type": "Polygon", + "coordinates": [[[cx, cy], [cx + s, cy], [cx + s, cy + s], [cx, cy + s], [cx, cy]]], + }, + } + ) + + labels_geojson = labels_dir / "labels.geojson" + labels_geojson.write_text(json.dumps({"type": "FeatureCollection", "features": features})) + + stac_path = root / "dataset-stac-item.json" + stac_path.write_text(json.dumps(_build_dataset_stac_item(chips_dir, labels_geojson), indent=2)) + # `models/ramp/pipeline.py` expects `dataset_labels` to be a single GeoJSON file. + return {"chips": chips_dir, "labels": labels_geojson, "dataset_stac_item": stac_path} + + +@pytest.fixture(scope="session") +def generate_toy_dataset(tmp_path_factory: pytest.TempPathFactory) -> dict[str, Path]: + return create_toy_data(tmp_path_factory.mktemp("toy_ramp")) + + +def _build_dataset_stac_item(chips_dir: Path, labels_geojson: Path) -> dict[str, Any]: + return { + "type": "Feature", + "stac_version": "1.1.0", + "stac_extensions": [ + "https://stac-extensions.github.io/label/v1.0.1/schema.json", + ], + "id": "toy-ramp", + "geometry": _GEOMETRY, + "bbox": _BBOX, + "properties": { + "datetime": "2024-01-01T00:00:00Z", + "description": "Toy RAMP semantic-segmentation dataset", + "label:type": "vector", + "label:tasks": ["segmentation"], + "label:classes": [{"name": "building", "classes": ["building"]}], + "label:description": "Toy polygon labels", + "keywords": ["building"], + "providers": [ + { + "name": "HOTOSM", + "roles": ["producer"], + "url": "https://www.hotosm.org", + "description": "Humanitarian OpenStreetMap Team", + } + ], + "fair:user_id": "test", + "version": "1", + "deprecated": False, + "license": "CC-BY-4.0", + }, + "assets": { + "chips": {"href": str(chips_dir), "type": "image/tiff", "roles": ["data"]}, + "labels": {"href": str(labels_geojson), "type": "application/geo+json", "roles": ["labels"]}, + }, + "links": [], + } diff --git a/models/ramp/tests/test_steps.py b/models/ramp/tests/test_steps.py new file mode 100644 index 0000000..4d210cf --- /dev/null +++ b/models/ramp/tests/test_steps.py @@ -0,0 +1,181 @@ +"""Step tests for the RAMP pipeline. + +Each test calls ``step.entrypoint(...)`` directly. Heavy RAMP/TF operations +(``train_ramp_model``, SavedModel loading, tf2onnx conversion) are patched so +tests run quickly and do not require GPU resources. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from pathlib import Path +from typing import Any +from unittest.mock import patch + + +@contextmanager +def _noop_mlflow_ctx(*_args: Any, **_kwargs: Any): + yield + + +class _FakeKerasModel: + def __init__(self, payload: bytes): + self._payload = payload + + def save(self, path: str) -> None: + Path(path).write_bytes(self._payload) + + +def test_split_dataset(toy_chips: Path, toy_labels: Path, base_hyperparameters: dict[str, Any]) -> None: + from models.ramp.pipeline import split_dataset + + hyperparameters = dict(base_hyperparameters) + hyperparameters.update( + { + "epochs": 1, + "batch_size": 1, + "val_fraction": 0.25, + "split_seed": 42, + "boundary_width": 1, + "contact_spacing": 2, + } + ) + + with patch("models.ramp.pipeline.log_metadata"): + result = split_dataset.entrypoint( + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + hyperparameters=hyperparameters, + ) + + assert result["strategy"] == "random" + assert result["train_count"] > 0 + assert result["val_count"] > 0 + assert "_ramp_train_dir" in result + assert "_preprocessed_dir" in result + ramp_dir = Path(result["_ramp_train_dir"]) + assert ramp_dir.exists() + assert (ramp_dir / "chips").is_dir() + assert (ramp_dir / "val-chips").is_dir() + + +def test_train_model(toy_chips: Path, toy_labels: Path, base_hyperparameters: dict[str, Any], tmp_path: Path) -> None: + import tensorflow as tf + + from models.ramp.pipeline import train_model + + ramp_train_dir = tmp_path / "ramp_training_work" + (ramp_train_dir / "chips").mkdir(parents=True) + (ramp_train_dir / "val-chips").mkdir(parents=True) + + split_info = { + "_work_dir": str(tmp_path), + "_preprocessed_dir": str(tmp_path / "preprocessed"), + "_ramp_train_dir": str(ramp_train_dir), + "strategy": "random", + "val_ratio": 0.25, + "seed": 42, + "train_count": 3, + "val_count": 1, + "description": "test split", + } + hyperparameters = dict(base_hyperparameters) + hyperparameters.update({"epochs": 1, "batch_size": 1}) + + expected = b"fake-keras-bytes" + fake_model_path = tmp_path / "best_model.keras" + fake_model_path.write_bytes(b"stub") + + with ( + patch("models.ramp.pipeline.mlflow_training_context", _noop_mlflow_ctx, create=True), + patch("models.ramp.pipeline.train_ramp_model", return_value=fake_model_path), + patch.object(tf.keras.models, "load_model", return_value=_FakeKerasModel(expected)), + patch("models.ramp.pipeline.log_metadata"), + ): + model_bytes = train_model.entrypoint( + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + base_model_weights="https://example.com/baseline.zip", + hyperparameters=hyperparameters, + split_info=split_info, + num_classes=4, + ) + + assert isinstance(model_bytes, bytes) + assert model_bytes == expected + + +def test_evaluate_model( + toy_chips: Path, toy_labels: Path, base_hyperparameters: dict[str, Any], tmp_path: Path +) -> None: + from models.ramp.pipeline import evaluate_model + + ramp_train_dir = tmp_path / "ramp_training_work" + (ramp_train_dir / "chips").mkdir(parents=True) + (ramp_train_dir / "val-chips").mkdir(parents=True) + (ramp_train_dir / "val-multimasks").mkdir(parents=True) + + split_info = { + "_work_dir": str(tmp_path), + "_preprocessed_dir": str(tmp_path / "preprocessed"), + "_ramp_train_dir": str(ramp_train_dir), + "strategy": "random", + "val_ratio": 0.25, + "seed": 42, + "train_count": 3, + "val_count": 1, + "description": "test split", + } + + fake_saved_model_dir = tmp_path / "saved_model" + fake_saved_model_dir.mkdir() + (fake_saved_model_dir / "saved_model.pb").write_bytes(b"\x08\x01") + + with ( + patch("models.ramp.pipeline.mlflow_training_context", _noop_mlflow_ctx, create=True), + patch("models.ramp.pipeline._restore_checkpoint", return_value=fake_saved_model_dir), + patch("models.ramp.pipeline.log_metadata"), + ): + metrics = evaluate_model.entrypoint( + trained_model=b"fake", + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + hyperparameters=base_hyperparameters, + split_info=split_info, + ) + + expected = {"fair:accuracy", "fair:mean_iou", "fair:precision", "fair:recall"} + assert set(metrics.keys()) == expected + for key in expected: + assert isinstance(metrics[key], float) + + +def test_export_onnx(tmp_path: Path) -> None: + import onnx + import tensorflow as tf + from onnx import TensorProto, helper + + from models.ramp.pipeline import export_onnx + + # Build a toy ONNX model and capture its bytes. + x = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 1]) + y = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1]) + node = helper.make_node("Identity", ["input"], ["output"]) + graph = helper.make_graph([node], "toy", [x], [y]) + toy_model = helper.make_model(graph, producer_name="test") + toy_bytes = toy_model.SerializeToString() + + fake_keras_path = tmp_path / "model.keras" + fake_keras_path.write_bytes(b"stub") + + with ( + patch("models.ramp.pipeline._restore_checkpoint", return_value=fake_keras_path), + patch.object(tf.keras.models, "load_model", return_value=_FakeKerasModel(b"ignored")), + patch("tf2onnx.convert.from_keras", side_effect=lambda _m, opset, output_path: Path(output_path).write_bytes(toy_bytes)), + patch("models.ramp.pipeline.log_metadata"), + ): + exported = export_onnx.entrypoint(trained_model=b"fake") + + assert isinstance(exported, bytes) + loaded = onnx.load_from_string(exported) + onnx.checker.check_model(loaded)