Update README.md

README.md

@@ -22,7 +22,7 @@ The pretrained model weights are hosted externally and can be downloaded here:
 
 ➡️ **https://staging.cortex.thetavision.nl/dataset-provider/listing/2/**
 
-Download the file (e.g., `vit_b_gastronet5m.pth`) and place it locally or on your server.
+Download the file (e.g., `dinov2.pth`) and place it locally or on your device.
 
 ---
 
@@ -34,24 +34,22 @@ import torch
 import timm
 
 # Initialize ViT‑B backbone (no classifier head)
-model = timm.create_model("vit_base_patch16_224", pretrained=False, num_classes=0)
+
+model = timm.create_model("timm/vit_base_patch14_dinov2.lvd142m",
+                          pretrained=False,
+                          num_classes=0,
+                          img_size=336,
+                          )
 
 # Update this path to where you downloaded the checkpoint
-ckpt_path = "./weights/vit_b_gastronet5m.pth"
+ckpt_path = "dinov2.pth"
 state = torch.load(ckpt_path, map_location="cpu")
-
-# Handle state dict variations
-if "model" in state:
-    state_dict = state["model"]
-elif "state_dict" in state:
-    state_dict = state["state_dict"]
-else:
-    state_dict = state
+state_dict = state['teacher']
 
 # Remove 'module.' prefix if present
-clean_state = {k.replace("module.", ""): v for k, v in state_dict.items()}
-model.load_state_dict(clean_state, strict=False)
-
+clean_state = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+msg = model.load_state_dict(clean_state, strict=False)
+print(msg)
 model.eval()
 ```