1import os
2import hashlib
3
4import torch
5import segmentation_models_pytorch as smp
6
7
8url = "https://github.com/microsoft/ai4g-flood/raw/refs/heads/main/models/ai4g_sar_model.ckpt"
9state_dict = torch.hub.load_state_dict_from_url(url, weights_only=False, map_location="cpu")["state_dict"]
10state_dict = {k.replace("model.model.", ""): v for k, v in state_dict.items() if "model.model." in k}
11model = smp.Unet(
12 encoder_name="mobilenet_v2",
13 encoder_weights=None,
14 in_channels=2,
15 classes=2,
16)
17model.load_state_dict(state_dict, strict=True)
18
19filename = "unet_mobilenetv2_sentinel1_ai4g_flood.pth"
20torch.save(model.state_dict(), filename)
21md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()[:8]
22os.rename(filename, filename.replace(".pth", f"-{md5}.pth"))