Views
No views yet
1import os
2import hashlib
3
4import torch
5import segmentation_models_pytorch as smp
6
7
8url = "https://huggingface.co/restor/tcd-unet-r50/resolve/main/model.pt"
9state_dict = torch.hub.load_state_dict_from_url(url, weights_only=False, map_location="cpu")
10model = smp.Unet(
11 encoder_name="resnet50",
12 encoder_weights=None,
13 in_channels=3,
14 classes=2,
15)
16model.load_state_dict(state_dict, strict=True)
17filename = "unet_resnet50_oam_rgb_tcd.pth"
18torch.save(model.state_dict(), filename)
19md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()[:8]
20os.rename(filename, filename.replace(".pth", f"-{md5}.pth"))