Update README.md

README.md

@@ -23,3 +23,42 @@ merge_mode: weighted_average
 
 # Model Card
 
+Exported using the following code:
+
+```python
+from pathlib import Path
+import torch
+import torchvision.transforms.v2 as T
+from stac_model.torch.export import save
+import segmentation_models_pytorch as smp
+
+
+path = "FTW-Release-Full-3-class-unet-efficientnetb5-weight0.75-3xlonger.ckpt"
+ckpt = torch.load(path, map_location="cpu", weights_only=False)
+hparams = ckpt["hyper_parameters"]
+state_dict = {k.replace("model.", ""): v for k, v in ckpt["state_dict"].items()}
+del state_dict["criterion.weight"]
+model = smp.Unet(
+    encoder_name=hparams["backbone"],
+    encoder_weights=None,
+    in_channels=hparams["in_channels"],
+    classes=hparams["num_classes"],
+)
+model.load_state_dict(state_dict, strict=True)
+
+transforms = torch.nn.Sequential(
+    torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+    T.Normalize(mean=[0.0], std=[3000.0])
+)
+
+save(
+    output_file=Path("model.pt2"),
+    input_shape=[-1, hparams["in_channels"], -1, -1],
+    model=model,
+    transforms=transforms,
+    metadata=None,
+    device="cpu",
+    dtype=torch.float32,
+    aoti_compile_and_package=False
+)
+```