stac-model 0.5.1

A PydanticV2 validation and serialization libary for the STAC ML Model Extension

stac-model

A PydanticV2 and PySTAC validation and serialization library for the STAC ML Model Extension

Installation

pip install -U stac-model

or install with uv:

uv add stac-model

Then you can run

stac-model --help

Creating example metadata JSON for a STAC Item

stac-model

This will make this example item for an example model.

Validating Model Metadata

An alternative use of stac_model is to validate config files containing model metadata using the MLModelProperties schema.

Given a YAML or JSON file with the structure in examples/torch/mlm-metadata.yaml, the model metadata can be validated as follows:

import yaml
from stac_model.schema import MLModelProperties

with open("examples/mlm-metadata.yaml", "r", encoding="utf-8") as f:
    metadata = yaml.safe_load(f)

MLModelProperties.model_validate(metadata["properties"])  

Exporting and Packaging PyTorch Models, Transforms, and Model Metadata

As of PyTorch 2.8, and stac_model 0.4.0, you can now export and package PyTorch models, transforms, and model metadata using functions in stac_model.torch.export. Below is an example of exporting a U-Net model pretrained on the Fields of The World (FTW) dataset for field boundary segmentation in Sentinel-2 satellite imagery using the TorchGeo library.

📝 Note: To customize the metadata for your model you can use this example as a template.

import torch
import torchvision.transforms.v2 as T
from torchgeo.models import Unet_Weights, unet
from stac_model.torch.export import save

weights = Unet_Weights.SENTINEL2_3CLASS_FTW
transforms = torch.nn.Sequential(
  T.Resize((256, 256)),
  T.Normalize(mean=[0.0], std=[3000.0])
)
model = unet(weights=weights)

save(
    output_file="ftw.pt2",
    model=model,  # Must be an nn.Module
    transforms=transforms,  # Must be an nn.Module
    metadata="metadata.yaml",  # Can be a metadata yaml or stac_model.schema.MLModelProperties object
    input_shape=[-1, 8, -1, -1],  # -1 indicates a dynamic shaped dimension
    device="cpu",
    dtype=torch.float32,
    aoti_compile_and_package=False,  # True for AOTInductor compile otherwise use torch.export
)

The model, transforms, and metadata can then be loaded into an environment with only torch and stac_model as required dependencies like below:

import yaml
from torch.export.pt2_archive._package import load_pt2

pt2 = load_pt2(archive_path)
metadata = yaml.safe_load(pt2.extra_files["mlm-metadata"])

# If exported with aoti_compile_and_package=True
model = pt2.aoti_runners["model"]
transforms = pt2.aoti_runners["transforms"]

# If exported with aoti_compile_and_package=False
model = pt2.exported_programs["model"].module()
transforms = pt2.exported_programs["transforms"].module()

# Inference
batch = ...  # An input batch tensor
outputs = model(transforms(batch))

Creating a STAC Item from a PyTorch Model

You can generate a valid STAC Item using the Machine Learning Model (MLM) Extension.

The example below demonstrates creating a STAC Item from a U-Net model pretrained on the Fields of The World (FTW) dataset for field boundary segmentation in Sentinel-2 satellite imagery, using the TorchGeo library

from stac_model.examples import unet_mlm
from stac_model.torch import MLModelExtension
from torchgeo.models import unet, Unet_Weights

# Use default TorchGeo UNet weights
weights = Unet_Weights.SENTINEL2_2CLASS_NC_FTW
model = unet(weights=weights)

# Create an ItemMLModelExtension using the MLM extension
item_ext = MLModelExtension.from_torch(
    model,
    weights=weights,
    item_id="pytorch_geo_unet"
)

For a more complete example including STAC Item properties, geometry, and datetime ranges, see unet_mlm() in stac_model/examples.py.

📈 Releases

You can see the list of available releases on the GitHub Releases page.

📄 License

This project is licenced under the terms of the Apache Software License 2.0 licence. See LICENSE for more details.

💗 Credits