torchgeo/seco-eco

Model Weights extracted below:

import os
import hashlib

import torch
import timm


# download weights
url = "https://huggingface.co/eplekh/secoeco/resolve/221f37e38589df20e4982c3449b94e9a034f69e6/seco-eco_e100.ckpt"
ckpt = torch.hub.load_state_dict_from_url(url, map_location="cpu", progress=True)
arch, image_size, bands = ckpt["hyper_parameters"]["arch"], ckpt["hyper_parameters"]["in_size"], ckpt["hyper_parameters"]["bands"]
print(arch, image_size, bands)

# Bands correspond to B9 from https://github.com/PlekhanovaElena/ssl4eco/blob/7445e048035f7ae31c0eb45e1ed8426c9989fe56/pretraining/pretrain_seco_3heads.py#L220
bands = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'NDVI']

# map weights to timm and torchvision compatible
layer_mapping = {
    "0" : "conv1",
    "1" : "bn1",
    "4" : "layer1",
    "5" : "layer2",
    "6" : "layer3",
    "7" : "layer4",
}
state_dict = {k.replace("encoder_q.", ""): v for k, v in ckpt["state_dict"].items() if k.startswith("encoder_q.")}
state_dict = {k.replace(k.split(".")[0], layer_mapping[k.split(".")[0]]): v for k, v in state_dict.items()}

model = timm.create_model("resnet50", pretrained=False, in_chans=len(bands), num_classes=0)
model.load_state_dict(state_dict, strict=True)

# save and compute hash
filename = "resnet50_sentinel2_all_seco_eco.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"))