Serving an ONNX Land-Cover Model with ONNX Runtime

Serve a land-cover ONNX model with onnxruntime InferenceSession: read a raster in windows, run tiled inference, and write a georeferenced GeoTIFF with rasterio.

TL;DR Answer

Load the .onnx file with onnxruntime.InferenceSession, read the source raster in windows with rasterio, reshape each window to the (N, bands) float32 layout the model expects, call session.run, reshape the class predictions back to the window grid, and write them into a GeoTIFF that copies the source CRS, transform, and nodata. This turns an exported model into a running service that classifies a full scene tile by tile. Getting the model into portable ONNX form is covered by ONNX Export for Geospatial Model Inference; this page is the concrete serving how-to.

import onnxruntime as ort

session = ort.InferenceSession(
    "land_cover.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
input_meta = session.get_inputs()[0]
print(input_meta.name, input_meta.type, input_meta.shape)
# "float_input tensor(float) ['N', 6]"

Part of: ONNX Export for Geospatial Model Inference

Why Serving Differs from Export

Exporting a model answers “can this run outside Python?” Serving answers “how do I feed a 20,000 x 20,000 pixel Sentinel-2 scene through it without exhausting memory, losing the georeference, or silently classifying fill pixels?” Those are different problems, and the second one is where most land-cover pipelines break.

A pixel-wise classifier trained with scikit-learn or XGBoost expects a 2-D table of shape (n_pixels, n_bands). A raster on disk is a 3-D array of shape (bands, rows, cols). The serving layer has to transpose and flatten every tile into the table layout, run it, and fold the flat prediction vector back onto the row/col grid — all while carrying the CRS and affine transform through untouched so the output overlays the input exactly.

Three failures recur in production:

  • Silent dtype mismatch. The model expects tensor(float) (float32) but the raster reads as uint16 or float64. ONNX Runtime raises a type error, or worse, an implicit cast changes the feature scaling the model was calibrated on.
  • Georeference loss. The prediction array is written with a default transform, so the map is numerically correct but spatially adrift by hundreds of metres.
  • Nodata classified as land. Fill pixels at scene edges get fed to the model and come back as “water” or “cropland”, polluting area statistics and any downstream model-drift-detection-for-geospatial-inference baseline.

The OnnxTileInferencer class below closes all three gaps.

Core Principles for a Robust Inference Service

  • Read the contract from the session, not from memory. session.get_inputs() reports the exact input name, element type, and shape. Never hard-code the feed key — the ONNX exporter may have named it float_input, input, or X.
  • Cast explicitly to float32. Rasters rarely store float32. Convert every window with .astype(np.float32) before session.run, matching the tensor(float) type the classifier was exported with.
  • Window, never load the whole scene. Iterate dataset.block_windows() or a fixed tiling so peak memory stays bounded regardless of scene size.
  • Mask before you predict. Compute a validity mask from the nodata value, run only valid pixels through the model, and scatter results back. This saves compute and keeps fill pixels out of the classes.
  • Copy the profile, override only what changes. Start from src.profile, then set count=1, dtype to the class dtype, and the output nodata. CRS and transform pass through unchanged.
  • Pick providers as a fallback chain. ["CUDAExecutionProvider", "CPUExecutionProvider"] uses the GPU when present and degrades to CPU cleanly on a machine without CUDA.

Production-Ready Code

The OnnxTileInferencer wraps the session, reads the input contract once, and exposes a single run_scene method that streams a source raster to a classified GeoTIFF. It validates the band count against the model’s declared feature dimension so a mismatched raster fails immediately rather than deep inside session.run.

from __future__ import annotations

import logging
from pathlib import Path

import numpy as np
import onnxruntime as ort
import rasterio
from rasterio.windows import Window

logger = logging.getLogger(__name__)


class OnnxTileInferencer:
    """Serve a pixel-wise land-cover ONNX model over a raster, tile by tile.

    Parameters
    ----------
    model_path : str or Path
        Path to the exported .onnx classifier. It must accept a
        (N, bands) float32 tensor and return per-pixel class labels.
    providers : list[str] | None
        ONNX Runtime execution providers in priority order. Defaults to
        CUDA with a CPU fallback.
    output_nodata : int, default=255
        Class value written to masked / invalid pixels.
    """

    def __init__(
        self,
        model_path: str | Path,
        providers: list[str] | None = None,
        output_nodata: int = 255,
    ) -> None:
        if providers is None:
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self.session = ort.InferenceSession(str(model_path), providers=providers)
        self.output_nodata = output_nodata

        meta = self.session.get_inputs()[0]
        self.input_name = meta.name
        self.output_name = self.session.get_outputs()[0].name
        # Trailing static dimension is the feature (band) count.
        self.n_features = meta.shape[-1] if isinstance(meta.shape[-1], int) else None

        logger.info(
            "Loaded %s | provider=%s | input=%s%s -> output=%s",
            model_path, self.session.get_providers()[0],
            self.input_name, meta.shape, self.output_name,
        )

    def _predict_window(self, block: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Classify one window. block is (bands, h, w); mask is (h, w) of valid px."""
        bands, height, width = block.shape
        # (bands, h, w) -> (h*w, bands), the table layout the model expects.
        table = block.reshape(bands, -1).T.astype(np.float32)
        flat_mask = mask.reshape(-1)

        out = np.full(height * width, self.output_nodata, dtype=np.uint8)
        valid = table[flat_mask]
        if valid.shape[0] > 0:
            preds = self.session.run(
                [self.output_name], {self.input_name: valid}
            )[0]
            out[flat_mask] = np.asarray(preds, dtype=np.uint8).reshape(-1)
        return out.reshape(height, width)

    def run_scene(self, src_path: str | Path, dst_path: str | Path) -> None:
        """Stream src_path through the model and write a classified GeoTIFF."""
        with rasterio.open(src_path) as src:
            if self.n_features is not None and src.count != self.n_features:
                raise ValueError(
                    f"Model expects {self.n_features} bands but "
                    f"{src_path} has {src.count}."
                )

            profile = src.profile.copy()
            profile.update(
                count=1, dtype="uint8", nodata=self.output_nodata,
                compress="deflate",
            )
            src_nodata = src.nodata

            with rasterio.open(dst_path, "w", **profile) as dst:
                for _, window in src.block_windows(1):
                    block = src.read(window=window)  # (bands, h, w)
                    if src_nodata is not None:
                        mask = np.all(block != src_nodata, axis=0)
                    else:
                        mask = np.ones(block.shape[1:], dtype=bool)

                    classified = self._predict_window(block, mask)
                    dst.write(classified, 1, window=window)

        logger.info("Wrote classified scene to %s", dst_path)

Step-by-Step Walkthrough

1. Inspect the model contract

Before touching a raster, confirm what the model actually wants. The input name and element type drive every later step.

import onnxruntime as ort

session = ort.InferenceSession(
    "land_cover.onnx", providers=["CPUExecutionProvider"]
)
inp = session.get_inputs()[0]
print("name  :", inp.name)     # e.g. float_input
print("type  :", inp.type)     # tensor(float) -> cast to np.float32
print("shape :", inp.shape)    # ['N', 6] -> dynamic rows, 6 bands
print("active:", session.get_providers()[0])

A dynamic first dimension (reported as a string like 'N' or None) is what lets you feed a variable number of pixels per window. The trailing 6 is the band count your raster must supply.

2. Run a single window to validate the shape plumbing

import numpy as np
import rasterio
from rasterio.windows import Window

inferencer = OnnxTileInferencer("land_cover.onnx")

with rasterio.open("sentinel2_scene.tif") as src:
    win = Window(0, 0, 256, 256)
    block = src.read(window=win)          # (6, 256, 256)
    mask = np.ones(block.shape[1:], bool)
    classified = inferencer._predict_window(block, mask)

print(classified.shape)                   # (256, 256)
print(np.unique(classified))              # e.g. [1 2 3 4 5]

If classified.shape matches the window and the class values fall inside your legend, the transpose/flatten/reshape round trip is correct.

3. Run the full scene

inferencer = OnnxTileInferencer(
    "land_cover.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
    output_nodata=255,
)
inferencer.run_scene("sentinel2_scene.tif", "land_cover_map.tif")

The service reads the scene one internal tile at a time, so a multi-gigabyte input never lands in memory whole. If the source is not internally tiled, block_windows yields one strip per row; retile the source to 256 or 512 pixel blocks first for better throughput.

4. Batch several windows per session call

session.run has fixed per-call overhead. On the GPU especially, feeding one small window at a time starves the device. Accumulate valid rows from several windows and run them in one batch to raise utilisation.

def run_batched(inferencer, src_path, dst_path, batch_windows=8):
    import rasterio, numpy as np
    with rasterio.open(src_path) as src:
        profile = src.profile.copy()
        profile.update(count=1, dtype="uint8",
                       nodata=inferencer.output_nodata, compress="deflate")
        windows = [w for _, w in src.block_windows(1)]
        with rasterio.open(dst_path, "w", **profile) as dst:
            for i in range(0, len(windows), batch_windows):
                for w in windows[i:i + batch_windows]:
                    block = src.read(window=w)
                    m = (np.all(block != src.nodata, axis=0)
                         if src.nodata is not None
                         else np.ones(block.shape[1:], bool))
                    dst.write(inferencer._predict_window(block, m), 1, window=w)

For a pixel classifier the effective batch is already the flattened pixel count of a window, so the win here is mostly reduced Python-loop and I/O overhead; for a patch-based CNN, true window batching along the sample axis matters far more.

How a Tile Moves Through the Service

The diagram traces one window from disk to the classified GeoTIFF, showing where the reshape, mask, and georeference steps sit.

Tiled ONNX inference data flow A source GeoTIFF window of shape bands by height by width is reshaped into an N by bands float32 table, masked for nodata, run through onnxruntime InferenceSession, reshaped back to the tile grid, and written into an output GeoTIFF that copies the source CRS, transform, and nodata. Source GeoTIFF window read (bands, h, w) Reshape + cast (N, bands) float32 + mask session.run CPU / CUDA class labels Reshape back to (h, w) scatter mask Output GeoTIFF (1 band, uint8) profile copied: CRS + transform + nodata preserved src.profile CRS & affine transform carried forward

Verification

A classified GeoTIFF that opens without error can still be wrong: shifted, mis-projected, or full of impossible class codes. Check three invariants against the source before shipping the map.

import numpy as np
import rasterio

VALID_CLASSES = {1, 2, 3, 4, 5, 6}   # your land-cover legend
OUTPUT_NODATA = 255

with rasterio.open("sentinel2_scene.tif") as src, \
     rasterio.open("land_cover_map.tif") as dst:

    # 1. Spatial footprint matches exactly.
    assert dst.shape == src.shape, "row/col grid changed"
    assert dst.crs == src.crs, "CRS not preserved"
    assert dst.transform.almost_equals(src.transform), "transform drifted"

    # 2. Only legal class codes (plus nodata) were written.
    data = dst.read(1)
    present = set(np.unique(data)) - {OUTPUT_NODATA}
    assert present <= VALID_CLASSES, f"unexpected classes: {present - VALID_CLASSES}"

    # 3. Fill pixels stayed masked.
    if src.nodata is not None:
        src_mask = np.all(src.read() == src.nodata, axis=0)
        assert np.all(data[src_mask] == OUTPUT_NODATA), "nodata leaked into classes"

print("Verification passed: footprint, CRS, transform, classes, and nodata all consistent.")

The transform and CRS assertions catch georeference loss; the class-set check catches dtype or reshape corruption; the nodata check confirms fill pixels never reached the model. The same reshape-back-to-grid logic underpins any raster classifier, whether the labels come from ONNX Runtime or straight from a tree model, as shown in Mapping XGBoost Predictions Back to a Raster Grid. Once the map verifies, package the session and its providers into a reproducible image following Containerizing Geospatial Inference Pipelines so the CUDA and GDAL dependencies stay pinned.

FAQ

How do I find the input name and dtype my ONNX model expects?

Call session.get_inputs() on the InferenceSession. Each entry exposes .name, .type (for example tensor(float)), and .shape. Use the reported name as the key in the feed dictionary passed to session.run, and cast your NumPy array to the matching dtype — almost always float32 for a land-cover classifier. A dynamic dimension appears as a string or None in .shape, which is why you feed a concrete (N, bands) array at run time rather than a fixed batch.

How do I run inference on the GPU instead of the CPU?

Pass providers to InferenceSession, for example ["CUDAExecutionProvider", "CPUExecutionProvider"]. ONNX Runtime tries each provider in order and falls back to CPU when CUDA is unavailable, so the list doubles as a graceful degradation path. Confirm the active provider with session.get_providers() after construction. The CUDA provider needs a matching onnxruntime-gpu build and compatible CUDA and cuDNN libraries, which is easiest to guarantee inside a container.

How do I keep nodata pixels out of the prediction?

Build a boolean mask from the source nodata value before running the model, run inference only on valid pixels, then scatter the class predictions back into a full array whose masked pixels are set to the output nodata value. Writing that nodata value into the destination GeoTIFF profile means colour ramps and area statistics treat masked pixels as absent rather than as a real land-cover class.

Part of: ONNX Export for Geospatial Model Inference — Geospatial MLOps and Model Deployment