ftw-prue-ccby

PRUE U-Nets with EfficientNet encoders from the FTW v3.1 release, trained only on the CC-BY and CC0 countries of the Fields of The World dataset.

Files

Links are pinned to commit ffe4452.

The .pth files hold smp.Unet(encoder_name, in_channels=8, classes=3) state dicts. The ONNX files are single files with no external data, and accept any batch size and any height and width that are multiples of 32.

Input and output

The input is two Sentinel-2 L2A scenes of the same area from different points in the growing season, stacked as B04, B03, B02, B08 of the first scene followed by the same four bands of the second (8 bands, the layout ftw inference download writes), with reflectance divided by 3000. The output is per-pixel logits for three classes: background, field, and boundary.

Use with ONNX Runtime

import numpy as np
import onnxruntime as ort
import rasterio

url = "https://huggingface.co/fieldsoftheworld/ftw-prue-ccby/resolve/ffe4452768e1054b1d7c2ab39b84b4170f90c7cf/prue_efnetb3_ccby-b35d59a6.onnx"
# download url to prue_efnetb3_ccby-b35d59a6.onnx, e.g. with curl -LO
session = ort.InferenceSession("prue_efnetb3_ccby-b35d59a6.onnx", providers=["CPUExecutionProvider"])

with rasterio.open("stacked.tif") as src:  # 8 bands, height and width multiples of 32
    image = src.read().astype(np.float32)[None] / 3000.0

logits = session.run(None, {"image": image})[0]  # (1, 3, height, width)
classes = logits.argmax(axis=1)[0]  # 0 background, 1 field, 2 boundary

Metrics

Computed on the FTW Full dataset with ftw model test -p3 -t3 --temporal_options stacked --model path/to/model.ckpt --countries full_data.

The CC-BY-NC checkpoints from the paper, for reference:

How the files were made

The .pth files come from the v3.1 Lightning checkpoints. The suffix is the first 8 hex digits of the file's sha256.

import hashlib
import os

import segmentation_models_pytorch as smp
import torch

urls = {
    "efficientnet-b3": "https://github.com/fieldsoftheworld/ftw-baselines/releases/download/v3.1/prue_efnetb3_ccby_checkpoint.ckpt",
    "efficientnet-b5": "https://github.com/fieldsoftheworld/ftw-baselines/releases/download/v3.1/prue_efnetb5_ccby_checkpoint.ckpt",
    "efficientnet-b7": "https://github.com/fieldsoftheworld/ftw-baselines/releases/download/v3.1/prue_efnetb7_ccby_checkpoint.ckpt",
}
for encoder_name, url in urls.items():
    state_dict = torch.hub.load_state_dict_from_url(url, weights_only=True, map_location="cpu")["state_dict"]
    state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
    model = smp.Unet(encoder_name=encoder_name, in_channels=8, classes=3, encoder_weights=None)
    model.load_state_dict(state_dict)
    filename = os.path.basename(url).replace("_checkpoint.ckpt", ".pth")
    torch.save(model.state_dict(), filename)
    sha256 = hashlib.sha256(open(filename, "rb").read()).hexdigest()[:8]
    os.rename(filename, filename.replace(".pth", f"-{sha256}.pth"))

The ONNX files are exported from those state dicts with opset 18:

h, w = torch.export.Dim("h", min=2), torch.export.Dim("w", min=2)
torch.onnx.export(
    model.eval(),
    (torch.zeros(1, 8, 256, 256),),
    pth_filename.replace(".pth", ".onnx"),
    input_names=["image"],
    output_names=["logits"],
    dynamic_shapes={"x": {0: torch.export.Dim("batch"), 2: h * 32, 3: w * 32}},
    opset_version=18,
    dynamo=True,
    external_data=False,
)

On a 352x544 tile, ONNX Runtime and PyTorch predict the same class for every pixel, with logits within 3e-4.