Upload README.md with huggingface_hub

README.md

@@ -74,13 +74,18 @@ weighted avg      0.969     0.967     0.967        30
 ```python
 import timm
 import torch
+from huggingface_hub import hf_hub_download
 from PIL import Image
 from safetensors.torch import load_file
 from torchvision import transforms
 
-# Load model
+# Download and load model from Hub
+weights_path = hf_hub_download(
+    repo_id="davanstrien/nls-index-card-classifier",
+    filename="classifier.safetensors"
+)
 model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=2)
-model.load_state_dict(load_file('classifier.safetensors'))
+model.load_state_dict(load_file(weights_path))
 model.eval()
 
 # Preprocess
@@ -96,10 +101,12 @@ input_tensor = transform(image).unsqueeze(0)
 
 with torch.no_grad():
     output = model(input_tensor)
+    probs = torch.softmax(output, dim=1)
     pred = output.argmax(1).item()
+    confidence = probs[0, pred].item()
 
 classes = ['other', 'index_card']
-print(f"Prediction: {classes[pred]}")
+print(f"Prediction: {classes[pred]} ({confidence:.1%})")
 ```
 
 ## Training