Update README.md

README.md

@@ -108,19 +108,47 @@ This repository hosts a fine-tuned `vit_large_patch16_384` classifier
 ## Usage
 
 ```python
-import timm
-import torch
+import json, torch, timm
+from huggingface_hub import hf_hub_download
+from timm.data.transforms_factory import create_transform
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from PIL import Image
 
-model = timm.create_model(
-    "vit_large_patch16_384",
-    num_classes=92,
-    pretrained=False,
-)
-state_dict = torch.hub.load_state_dict_from_url(
-    "https://huggingface.co/rexologue/vit_large_384_for_trees/resolve/main/pytorch_model.bin",
-    map_location="cpu",
-    file_name="rexologue--vit_large_384_for_trees.bin",
+REPO = "rexologue/vit_large_384_for_trees"
+MODEL_NAME = "vit_large_patch16_384"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 1) labels
+labels_path = hf_hub_download(REPO, filename="labels.json")
+with open(labels_path, "r", encoding="utf-8") as f:
+    raw = json.load(f)
+labels = [raw[str(i)] for i in range(len(raw))] if isinstance(raw, dict) else list(raw)
+
+# 2) weights
+ckpt_path = hf_hub_download(REPO, filename="pytorch_model.bin")
+state = torch.load(ckpt_path, map_location="cpu")
+if any(k.startswith("module.") for k in state):  # DDP fix
+    state = {k.replace("module.", "", 1): v for k, v in state.items()}
+
+# 3) model
+model = timm.create_model(MODEL_NAME, num_classes=len(labels), pretrained=False)
+model.load_state_dict(state, strict=True)
+model.to(DEVICE).eval()
+
+# 4) preprocessing (ViT-L/16 @ 384 w/ ImageNet mean/std + bicubic)
+transform = create_transform(
+    input_size=(3, 384, 384),
+    interpolation="bicubic",
+    mean=IMAGENET_DEFAULT_MEAN,
+    std=IMAGENET_DEFAULT_STD,
 )
-model.load_state_dict(state_dict)
-model.eval()
+
+# 5) run
+img = Image.open("your_image.jpg").convert("RGB")
+x = transform(img).unsqueeze(0).to(DEVICE)
+with torch.no_grad():
+    logits = model(x)
+probs = torch.softmax(logits, dim=1)[0].cpu()
+topk = probs.topk(k=min(5, len(labels)))
+print([(labels[i], float(probs[i])) for i in topk.indices])
 ```