|
|
--- |
|
|
license: mit |
|
|
library_name: torchgeo |
|
|
--- |
|
|
|
|
|
|
|
|
Model Weights extracted below: |
|
|
|
|
|
```python |
|
|
import os |
|
|
import hashlib |
|
|
|
|
|
import torch |
|
|
import timm |
|
|
|
|
|
|
|
|
# download weights |
|
|
url = "https://huggingface.co/eplekh/secoeco/resolve/main/ablation_B12_weights.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 = ['B1', 'B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B9', 'B11', 'B12'] |
|
|
|
|
|
# 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")) |
|
|
``` |
