isaaccorley/wstsplus

Rehosted from https://huggingface.co/saadlahrichi/WSTSPlus

Converted using the following code:

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"))