wherobots/ftw-v1-pt2

Model exported using the following code:

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 = "3_Class_FULL_FTW_Pretrained.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
)