1from pathlib import Path
2import torch
3import torchvision.transforms.v2 as T
4from stac_model.torch.export import save
5import segmentation_models_pytorch as smp
6
7
8path = "FTW-Release-Full-3-class-unet-efficientnetb5-weight0.75-3xlonger.ckpt"
9ckpt = torch.load(path, map_location="cpu", weights_only=False)
10hparams = ckpt["hyper_parameters"]
11state_dict = {k.replace("model.", ""): v for k, v in ckpt["state_dict"].items()}
12del state_dict["criterion.weight"]
13model = smp.Unet(
14 encoder_name=hparams["backbone"],
15 encoder_weights=None,
16 in_channels=hparams["in_channels"],
17 classes=hparams["num_classes"],
18)
19model.load_state_dict(state_dict, strict=True)
20
21transforms = torch.nn.Sequential(
22 torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
23 T.Normalize(mean=[0.0], std=[3000.0])
24)
25
26save(
27 output_file=Path("model.pt2"),
28 input_shape=[-1, hparams["in_channels"], -1, -1],
29 model=model,
30 transforms=transforms,
31 metadata=None,
32 device="cpu",
33 dtype=torch.float32,
34 aoti_compile_and_package=False
35)
