Update README.md

README.md

@@ -1,3 +1,46 @@
----
-license: mit
----
+---
+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"))
+```