ONNX Model Support for Geospatial Inference¶

This notebook demonstrates how to use ONNX Runtime with GeoAI for efficient geospatial model inference. The ONNX module mirrors the AutoGeoModel API, allowing you to export Hugging Face models to ONNX format and run inference on GeoTIFF data using ONNX Runtime.

Why ONNX?¶

• Faster inference: ONNX Runtime provides optimized execution across different hardware
• Edge deployment: Run models without PyTorch or GPU drivers
• Cross-platform: Deploy on Windows, Linux, macOS, and mobile devices
• Smaller footprint: No need for full PyTorch installation at inference time

In [ ]:

# %pip install -U geoai-py[onnx] transformers

# %pip install -U geoai-py[onnx] transformers

In [ ]:

from geoai import download_file
from geoai.onnx import (
    export_to_onnx,
    ONNXGeoModel,
    onnx_semantic_segmentation,
    onnx_image_classification,
)

from geoai import download_file from geoai.onnx import ( export_to_onnx, ONNXGeoModel, onnx_semantic_segmentation, onnx_image_classification, )

Download Sample Data¶

Download sample aerial imagery for the examples.

In [ ]:

image_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/aerial.tif"
image_path = download_file(image_url, "aerial.tif")

image_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/aerial.tif" image_path = download_file(image_url, "aerial.tif")

1. Export a Hugging Face Model to ONNX¶

The export_to_onnx() function converts a pretrained Hugging Face model to ONNX format. It automatically handles model loading, tracing, and saving metadata (label mappings, input sizes, task type) to a sidecar JSON file.

Export a Semantic Segmentation Model¶

In [ ]:

seg_onnx_path = export_to_onnx(
    "nvidia/segformer-b0-finetuned-ade-512-512",
    "segformer.onnx",
    task="semantic-segmentation",
)
print(f"Model exported to: {seg_onnx_path}")

seg_onnx_path = export_to_onnx( "nvidia/segformer-b0-finetuned-ade-512-512", "segformer.onnx", task="semantic-segmentation", ) print(f"Model exported to: {seg_onnx_path}")

Export an Image Classification Model¶

In [ ]:

cls_onnx_path = export_to_onnx(
    "google/vit-base-patch16-224",
    "vit_classifier.onnx",
    task="image-classification",
    input_height=224,
    input_width=224,
)
print(f"Model exported to: {cls_onnx_path}")

cls_onnx_path = export_to_onnx( "google/vit-base-patch16-224", "vit_classifier.onnx", task="image-classification", input_height=224, input_width=224, ) print(f"Model exported to: {cls_onnx_path}")

2. Semantic Segmentation with ONNX¶

Using ONNXGeoModel Directly¶

The ONNXGeoModel class provides full control over inference. It supports tiled processing for large GeoTIFF files and preserves georeferencing.

In [ ]:

model = ONNXGeoModel("segformer.onnx", task="semantic-segmentation")

result = model.predict(
    image_path,
    output_path="segmentation_onnx.tif",
)

mask = result.get("mask", result.get("output"))
print(f"Segmentation mask shape: {mask.shape}")
print(f"Unique classes: {len(set(mask.flatten()))}")
print(f"Output saved to: segmentation_onnx.tif")

model = ONNXGeoModel("segformer.onnx", task="semantic-segmentation") result = model.predict( image_path, output_path="segmentation_onnx.tif", ) mask = result.get("mask", result.get("output")) print(f"Segmentation mask shape: {mask.shape}") print(f"Unique classes: {len(set(mask.flatten()))}") print(f"Output saved to: segmentation_onnx.tif")

Vectorize Segmentation Results¶

Save segmentation results as both raster (GeoTIFF) and vector (GeoJSON) outputs.

In [ ]:

result = model.predict(
    image_path,
    output_path="segmentation_onnx.tif",
    output_vector_path="segmentation_onnx.geojson",
    min_object_area=50,
    simplify_tolerance=1.0,
)

if "geodataframe" in result:
    print("Vectorized segmentation saved to segmentation_onnx.geojson")
    print(f"Number of polygons: {len(result['geodataframe'])}")
    print(result["geodataframe"].head())

result = model.predict( image_path, output_path="segmentation_onnx.tif", output_vector_path="segmentation_onnx.geojson", min_object_area=50, simplify_tolerance=1.0, ) if "geodataframe" in result: print("Vectorized segmentation saved to segmentation_onnx.geojson") print(f"Number of polygons: {len(result['geodataframe'])}") print(result["geodataframe"].head())

Using the Convenience Function¶

The onnx_semantic_segmentation() function provides a simpler interface.

In [ ]:

result = onnx_semantic_segmentation(
    image_path,
    output_path="segmentation_onnx2.tif",
    model_path="segformer.onnx",
    output_vector_path="segmentation_onnx2.geojson",
    min_object_area=50,
)

print(f"Segmentation complete")
if "geodataframe" in result:
    print(f"Number of polygons: {len(result['geodataframe'])}")

result = onnx_semantic_segmentation( image_path, output_path="segmentation_onnx2.tif", model_path="segformer.onnx", output_vector_path="segmentation_onnx2.geojson", min_object_area=50, ) print(f"Segmentation complete") if "geodataframe" in result: print(f"Number of polygons: {len(result['geodataframe'])}")

3. Image Classification with ONNX¶

In [ ]:

import numpy as np

cls_result = onnx_image_classification(
    image_path,
    model_path="vit_classifier.onnx",
)

print(f"Predicted class index: {cls_result.get('class')}")
if cls_result.get("label"):
    print(f"Predicted label: {cls_result['label']}")

if "probabilities" in cls_result and cls_result["probabilities"] is not None:
    probs = cls_result["probabilities"]
    top_indices = np.argsort(probs)[-5:][::-1]
    print("\nTop 5 predictions:")
    for idx in top_indices:
        label = (
            cls_result.get("label", f"Class {idx}")
            if idx == cls_result.get("class")
            else f"Class {idx}"
        )
        print(f"  {label}: {probs[idx]:.4f}")

import numpy as np cls_result = onnx_image_classification( image_path, model_path="vit_classifier.onnx", ) print(f"Predicted class index: {cls_result.get('class')}") if cls_result.get("label"): print(f"Predicted label: {cls_result['label']}") if "probabilities" in cls_result and cls_result["probabilities"] is not None: probs = cls_result["probabilities"] top_indices = np.argsort(probs)[-5:][::-1] print("\nTop 5 predictions:") for idx in top_indices: label = ( cls_result.get("label", f"Class {idx}") if idx == cls_result.get("class") else f"Class {idx}" ) print(f" {label}: {probs[idx]:.4f}")

Using ONNXGeoModel for Classification¶

In [ ]:

cls_model = ONNXGeoModel("vit_classifier.onnx", task="image-classification")

cls_result = cls_model.predict(image_path)

print(f"Predicted class: {cls_result.get('class')}")
print(f"Label: {cls_result.get('label', 'N/A')}")

cls_model = ONNXGeoModel("vit_classifier.onnx", task="image-classification") cls_result = cls_model.predict(image_path) print(f"Predicted class: {cls_result.get('class')}") print(f"Label: {cls_result.get('label', 'N/A')}")

4. Performance Notes¶

ONNX Runtime provides several advantages over PyTorch for inference:

Feature	PyTorch	ONNX Runtime
Installation size	Large (~2 GB+)	Small (~50 MB)
GPU requirement	Often needed	CPU-optimized
Edge deployment	Difficult	Straightforward
Cross-platform	Python-centric	C++, C#, Java, JS, etc.
Inference speed	Baseline	Often 1.5-3x faster

ONNX models are ideal for:

• Production deployment: Smaller, faster, no training framework needed
• Edge devices: Run on IoT devices, drones, and embedded systems
• Web applications: Use ONNX.js or ONNX Runtime Web
• Batch processing: Faster throughput for large-scale geospatial workflows

Cleanup¶

In [ ]:

import os

files_to_remove = [
    "aerial.tif",
    "segformer.onnx",
    "segformer.onnx.json",
    "vit_classifier.onnx",
    "vit_classifier.onnx.json",
    "segmentation_onnx.tif",
    "segmentation_onnx.geojson",
    "segmentation_onnx2.tif",
    "segmentation_onnx2.geojson",
]

for f in files_to_remove:
    if os.path.exists(f):
        os.remove(f)
        print(f"Removed {f}")

import os files_to_remove = [ "aerial.tif", "segformer.onnx", "segformer.onnx.json", "vit_classifier.onnx", "vit_classifier.onnx.json", "segmentation_onnx.tif", "segmentation_onnx.geojson", "segmentation_onnx2.tif", "segmentation_onnx2.geojson", ] for f in files_to_remove: if os.path.exists(f): os.remove(f) print(f"Removed {f}")