| --- |
| license: mit |
| --- |
| |
| Rehosted from https://huggingface.co/saadlahrichi/WSTSPlus |
|
|
| Converted using the following code: |
|
|
| ```python |
| import os |
| import segmentation_models_pytorch as smp |
| import torch |
| import hashlib |
| |
| urls = { |
| "t1-all": "https://hf.co/saadlahrichi/WSTSPlus/resolve/main/trained_model_weights/Res18Unet_T1/All/fold6_testAP0.577.pth", |
| "t1-multi": "https://hf.co/saadlahrichi/WSTSPlus/resolve/main/trained_model_weights/Res18Unet_T1/Multi/fold6_testAP0.585.pth", |
| "t1-veg": "https://hf.co/saadlahrichi/WSTSPlus/resolve/main/trained_model_weights/Res18Unet_T1/Veg/fold2_testAP0.567.pth", |
| "t5-all": "https://hf.co/saadlahrichi/WSTSPlus/resolve/main/trained_model_weights/Res18Unet_T5/All/fold2_testAP0.594.pth", |
| "t5-multi": "https://hf.co/saadlahrichi/WSTSPlus/resolve/main/trained_model_weights/Res18Unet_T5/Multi/fold2_testAP0.597.pth", |
| "t5-veg": "https://hf.co/saadlahrichi/WSTSPlus/resolve/main/trained_model_weights/Res18Unet_T5/Veg/fold2_testAP0.590.pth", |
| } |
| for prefix, url in urls.items(): |
| state_dict = torch.hub.load_state_dict_from_url(url, weights_only=True, map_location="cpu") |
| state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()} |
| in_channels = state_dict["encoder.conv1.weight"].shape[1] |
| print(f"{prefix}: in_channels={in_channels}") |
| model = smp.Unet(encoder_name="resnet18", in_channels=in_channels, classes=1, encoder_weights=None) |
| model.load_state_dict(state_dict, strict=True) |
| filename = f"unet-resnet18-{prefix}.pth" |
| torch.save(model.state_dict(), filename) |
| md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()[:8] |
| os.rename(filename, filename.replace(".pth", f"-{md5}.pth")) |
| ``` |
