Update README.md

README.md

@@ -22,25 +22,50 @@ license: cc-by-nc-2.0
 ## 🖼️ Quick Usage
 
 ```python
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-from PIL import Image
+## 🔎 Inference Example
+
+```python
 import torch
+from PIL import Image
+from torchvision import models, transforms
+from huggingface_hub import hf_hub_download
 
-# Load model & feature extractor
-extractor = AutoFeatureExtractor.from_pretrained("your-username/RistoNet")
-model = AutoModelForImageClassification.from_pretrained("your-username/RistoNet")
+# ------------------------------
+# 1. Load model
+# ------------------------------
+MODEL_REPO = "Orkidee/RistoNet"
+MODEL_FILE = "ristonet.pth"
 
-# Load an image
-image = Image.open("my_dish.jpg")
+model = models.efficientnet_b0(weights=None)  # no pretrained weights
+num_features = model.classifier[1].in_features
+model.classifier[1] = torch.nn.Linear(num_features, 2)  # 2 classes
 
-# Preprocess
-inputs = extractor(images=image, return_tensors="pt")
+# Download weights from Hub and load
+model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+model.load_state_dict(torch.load(model_path, map_location="cpu"))
+model.eval()
+
+# ------------------------------
+# 2. Define preprocessing
+# ------------------------------
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
+
+# ------------------------------
+# 3. Run inference on an image
+# ------------------------------
+image = Image.open("my_food.jpg").convert("RGB")
+input_tensor = transform(image).unsqueeze(0)  # add batch dim
 
-# Predict
 with torch.no_grad():
-    outputs = model(**inputs)
-    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    outputs = model(input_tensor)
+    probs = torch.nn.functional.softmax(outputs, dim=1)
+    predicted_class = probs.argmax().item()
+
+print("Predicted class:", predicted_class)
+print("Probabilities:", probs.numpy())
 
-# Top class
-predicted_class = predictions.argmax().item()
-print(f"Predicted class: {predicted_class}")