Metadata-Version: 2.1
Name: stac-model
Version: 0.6.0
Summary: A PydanticV2 validation and serialization libary for the STAC ML Model Extension
Keywords: STAC,SpatioTemporal Asset Catalog,Machine Learning Model,Artificial Intelligence
Author-Email: Ryan Avery <ryan@wherobots.com>, Francis Charette-Migneault <francis.charette-migneault@crim.ca>
License: Apache Software License 2.0
Classifier: Development Status :: 4 - Beta
Classifier: Operating System :: OS Independent
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Framework :: Pydantic
Classifier: Framework :: Pydantic :: 2
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Information Technology
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: File Formats :: JSON :: JSON Schema
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: GIS
Classifier: Topic :: Scientific/Engineering :: Image Processing
Classifier: Topic :: Scientific/Engineering :: Image Recognition
Project-URL: homepage, https://github.com/stac-extensions/mlm/blob/main/README_STAC_MODEL.md
Project-URL: repository, https://github.com/crim-ca/mlm-extension
Requires-Python: <4.0,>=3.11
Requires-Dist: typer<1.0.0,>=0.9.0
Requires-Dist: rich<16.0.0,>=13.7.0
Requires-Dist: pydantic<3.0.0,>=2.6.3
Requires-Dist: pydantic-core<3,>=2
Requires-Dist: pystac<2.0.0,>=1.9.0
Requires-Dist: shapely<3,>=2
Requires-Dist: jsonschema<5.0.0,>=4.21.1
Requires-Dist: pip>=25.0.0
Provides-Extra: torch
Requires-Dist: torch>=2.8.0; extra == "torch"
Requires-Dist: torchgeo>=0.8.1; extra == "torch"
Requires-Dist: torchvision<1,>=0.21; extra == "torch"
Provides-Extra: torch-cu126
Requires-Dist: pytorch-triton; extra == "torch-cu126"
Requires-Dist: torch>=2.8.0; extra == "torch-cu126"
Requires-Dist: torchgeo>=0.8.1; extra == "torch-cu126"
Requires-Dist: torchvision<1,>=0.21; extra == "torch-cu126"
Description-Content-Type: text/markdown

# stac-model

<!--lint disable no-html -->

<div align="center">

[![Python support][bp1]][bp2]
[![PyPI Release][bp3]][bp2]
[![Repository][bscm1]][bp4]
[![Releases][bscm2]][bp5]

[![Contributions Welcome][bp8]][bp9]

[![uv][bp11]][bp12]
[![Pre-commit][bp15]][bp16]
[![Semantic versions][blic3]][bp5]
[![Pipelines][bscm6]][bscm7]

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

</div>

## Python Version Support

> [!NOTE]
> `stac-model` package versions and MLM specification versions are independently managed.
>
> - **MLM Extension Specification**: Versioned as `vX.Y.Z` (e.g., `v1.5.1`)
> - **stac-model Package**: Versioned as `stac-model-vX.Y.Z` (e.g., `stac-model-v0.5.2`)

**Python Version Compatibility:**

| `stac-model` Version | Python Version |
|----------------------| -------------- |
| `>=0.6.0`            | 3.11+          |
| `<0.6.0`             | 3.10+          |

## Installation

```shell
pip install -U stac-model
```

or install with uv:

```shell
uv add stac-model
```

Then you can run

```shell
stac-model --help
```

## Creating example metadata JSON for a STAC Item

```shell
stac-model
```

This will make [this example item](./examples/item_basic.json) 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](./examples/torch/mlm-metadata.yaml), the model metadata can be validated as follows:

```python
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](https://fieldsofthe.world/) for
field boundary segmentation in Sentinel-2 satellite imagery using the [TorchGeo](https://github.com/microsoft/torchgeo) library.

> 📝 **Note:** To customize the metadata for your model you can use this [example](./tests/torch/metadata.yaml) as a template.

```python
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:

```python
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](https://fieldsofthe.world/) for field boundary segmentation 
in Sentinel-2 satellite imagery, using the [TorchGeo](https://github.com/microsoft/torchgeo) library

```python
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`](stac_model/examples.py).

## 📈 Releases

You can see the list of available releases on the [GitHub Releases][github-releases] page.

For release procedure, see [Building and Releasing](CONTRIBUTING.md#building-and-releasing).

## 📄 License

[![License][blic1]][blic2]

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

## 💗 Credits

[![Python project templated from galactipy.][bp6]][bp7]

<!-- Anchors -->

[bp1]: https://img.shields.io/pypi/pyversions/stac-model?style=for-the-badge

[bp2]: https://pypi.org/project/stac-model/

[bp3]: https://img.shields.io/pypi/v/stac-model?style=for-the-badge&logo=pypi&color=3775a9

[bp4]: https://github.com/stac-extensions/mlm

[bp5]: https://github.com/stac-extensions/mlm/releases

[bp6]: https://img.shields.io/badge/made%20with-galactipy%20%F0%9F%8C%8C-179287?style=for-the-badge&labelColor=193A3E

[bp7]: https://kutt.it/7fYqQl

[bp8]: https://img.shields.io/static/v1.svg?label=Contributions&message=Welcome&color=0059b3&style=for-the-badge

[bp9]: https://github.com/stac-extensions/mlm/blob/main/CONTRIBUTING.md

[bp11]: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json&style=for-the-badge

[bp12]: https://docs.astral.sh/uv/

[bp15]: https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white&style=for-the-badge

[bp16]: https://github.com/stac-extensions/mlm/blob/main/.pre-commit-config.yaml

[blic1]: https://img.shields.io/github/license/stac-extensions/mlm?style=for-the-badge

[blic2]: https://github.com/stac-extensions/mlm/blob/main/LICENSE

[blic3]: https://img.shields.io/badge/%F0%9F%93%A6-semantic%20versions-4053D6?style=for-the-badge

[github-releases]: https://github.com/stac-extensions/mlm/releases

[bscm1]: https://img.shields.io/badge/GitHub-100000?style=for-the-badge&logo=github&logoColor=white

[bscm2]: https://img.shields.io/github/v/release/stac-extensions/mlm?filter=stac-model-v*&style=for-the-badge&logo=semantic-release&color=347d39

[bscm6]: https://img.shields.io/github/actions/workflow/status/stac-extensions/mlm/publish.yaml?style=for-the-badge&logo=github

[bscm7]: https://github.com/stac-extensions/mlm/blob/main/.github/workflows/publish.yaml

[hub1]: https://docs.github.com/en/code-security/dependabot/dependabot-version-updates/configuring-dependabot-version-updates#enabling-dependabot-version-updates

[hub2]: https://github.com/marketplace/actions/close-stale-issues

[hub6]: https://docs.github.com/en/code-security/dependabot

[hub8]: https://github.com/stac-extensions/mlm/blob/main/.github/release-drafter.yml

[hub9]: https://github.com/stac-extensions/mlm/blob/main/.github/.stale.yml
