Fixing linting and adding eo bands to pass test

Author: Nazim-crim

examples/collection.json

@@ -75,6 +75,10 @@
     {
       "href": "item_multi_io.json",
       "rel": "item"
+    },
+    {
+      "href": "torchgeo.json",
+      "rel": "item"
     }
   ]
 }

examples/torchgeo.json

@@ -113,6 +113,32 @@
       "type": "application/octet-stream; application=pytorch",
       "title": "U-Net (efficientnet-b3) weights trained on FTW dataset",
       "description": "A U-Net segmentation model with efficientnet-b3 encoder trained on FTW imagery with 2-class labels. Weights are non-commercial.",
+      "eo:bands": [
+        {
+          "name": "B4"
+        },
+        {
+          "name": "B3"
+        },
+        {
+          "name": "B2"
+        },
+        {
+          "name": "B8A"
+        },
+        {
+          "name": "B4"
+        },
+        {
+          "name": "B3"
+        },
+        {
+          "name": "B2"
+        },
+        {
+          "name": "B8A"
+        }
+      ],
       "roles": [
         "mlm:model",
         "mlm:weights",
@@ -147,6 +173,7 @@
     58.21
   ],
   "stac_extensions": [
-    "https://crim-ca.github.io/mlm-extension/v1.3.0/schema.json"
+    "https://crim-ca.github.io/mlm-extension/v1.3.0/schema.json",
+    "https://stac-extensions.github.io/eo/v1.1.0/schema.json"
   ]
 }

stac_model/examples.py

@@ -5,14 +5,13 @@
 from dateutil.parser import parse as parse_dt
 from pystac.extensions.eo import Band, EOExtension
 from pystac.extensions.file import FileExtension
+from torchgeo.models import Unet_Weights, unet
 
 from stac_model.base import ProcessingExpression
 from stac_model.input import InputStructure, MLMStatistic, ModelInput
 from stac_model.output import MLMClassification, ModelOutput, ModelResult
 from stac_model.schema import ItemMLModelExtension, MLModelExtension, MLModelProperties
 
-from torchgeo.models import unet, Unet_Weights
-
 
 def eurosat_resnet() -> ItemMLModelExtension:
     input_struct = InputStructure(
@@ -242,7 +241,9 @@ def eurosat_resnet() -> ItemMLModelExtension:
 
 def unet_mlm() -> ItemMLModelExtension:
     """
-    Example of a UNet model using pytorchgeo SENTINEL2_2CLASS_NC_FTW default weights with Machine Learning Model Extension metadata.
+    Example of a UNet model using PyTorchGeo SENTINEL2_2CLASS_NC_FTW default weights.
+
+    Returns an ItemMLModelExtension with Machine Learning Model Extension metadata.
     """
     weights = Unet_Weights.SENTINEL2_2CLASS_NC_FTW
     model = unet(weights=weights)

stac_model/schema.py

@@ -16,6 +16,7 @@
 )
 
 import pystac
+import torch.nn as nn
 from pydantic import ConfigDict, Field
 from pydantic.fields import FieldInfo
 from pystac.extensions.base import (
@@ -29,8 +30,6 @@
 from stac_model.output import ModelOutput
 from stac_model.runtime import Runtime
 
-import torch.nn as nn
-
 T = TypeVar(
     "T",
     pystac.Collection,

stac_model/utils.py

@@ -1,11 +1,13 @@
+from typing import Optional, cast
+
+import torch.nn as nn
 from pystac import Asset, Item, Link
+from pystac.extensions.eo import Band, EOExtension
+
 from stac_model.input import InputStructure, ModelInput
 from stac_model.output import MLMClassification, ModelOutput, ModelResult
 from stac_model.schema import ItemMLModelExtension, MLModelExtension, MLModelProperties
 
-import torch.nn as nn
-from typing import Optional
-
 
 def get_input_channels(state_dict: dict) -> int:
     """
@@ -131,7 +133,10 @@ def from_torch(
 
     # Model weights asset
     assets["model"] = Asset(
-        title=f"{meta.get('model', 'Model')} ({meta.get('encoder', '')}) weights trained on {meta.get('dataset', 'dataset')} dataset",
+        title=(
+            f"{meta.get('model', 'Model')} ({meta.get('encoder', '')}) weights "
+            f"trained on {meta.get('dataset', 'dataset')} dataset"
+        ),
         description=(
             f"A {meta.get('model', 'Model')} segmentation model with {meta.get('encoder', '')} encoder "
             f"trained on {meta.get('dataset', 'dataset')} imagery with {meta.get('num_classes', '?')}-class labels. "
@@ -194,4 +199,17 @@ def from_torch(
     ext = MLModelExtension.ext(item, add_if_missing=True)
     ext.apply(mlm_props)
 
-    return ItemMLModelExtension(item)
+    eo_model_asset = cast(
+        EOExtension[Asset],
+        EOExtension.ext(assets["model"], add_if_missing=True),
+    )
+    eo_bands = []
+    for name in bands:
+        band = Band({})
+        band.apply(name=name)
+        eo_bands.append(band)
+    eo_model_asset.apply(bands=eo_bands)
+
+    item_mlm = MLModelExtension.ext(item, add_if_missing=True)
+    item_mlm.apply(mlm_props.model_dump(by_alias=True, exclude_unset=True, exclude_defaults=True))
+    return item_mlm