Update README.md

README.md

@@ -60,35 +60,100 @@ pip install torch torchvision transformers scikit-learn pillow joblib numpy hugg
 ### Single Image Inference
 
 ```python
+
 import json
 import joblib
 from pathlib import Path
 from PIL import Image
 import torch
 import numpy as np
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfApi
 from transformers import AutoImageProcessor, ViTModel
+import warnings
 
-# Configuration
-REPO_ID = "Meeteshn/vit_fruit_ripeness_classifier"  
+# ----------------- CONFIG -----------------
+REPO_ID = "Meeteshn/vit_fruit_ripeness_classifier"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load model components
-processor = AutoImageProcessor.from_pretrained(REPO_ID, subfolder="vit_fruit_ripeness_updated/processor", use_fast=True)
-backbone  = ViTModel.from_pretrained(REPO_ID, subfolder="vit_fruit_ripeness_updated/vit_backbone")
-backbone.to(DEVICE)
-backbone.eval()
-
-# Load sklearn artifacts
-scaler_path = hf_hub_download(REPO_ID, "scaler.joblib")
-clf_path = hf_hub_download(REPO_ID, "logistic_model.joblib")
-metadata_path = hf_hub_download(REPO_ID, "metadata.json")
+NESTED_FOLDER = "vit_fruit_ripeness_updated"  # your repo uses this nested folder
+TOP_K = 5
+# ------------------------------------------
+
+def hf_download_try(repo_id: str, filename: str, nested_folder: str = NESTED_FOLDER):
+    """
+    Try to download `filename` from repo root, then from nested_folder/filename.
+    Returns local path to downloaded file or raises an informative error.
+    """
+    candidates = [filename, f"{nested_folder}/{filename}"]
+    last_exc = None
+    for f in candidates:
+        try:
+            print(f"Trying to download '{f}' from '{repo_id}'...")
+            path = hf_hub_download(repo_id=repo_id, filename=f)
+            print("Downloaded:", path)
+            return path
+        except Exception as e:
+            print(f"Not found at '{f}': {e}")
+            last_exc = e
+    raise RuntimeError(f"Could not download '{filename}' from repo '{repo_id}'. Last error: {last_exc}")
+
+def load_processor_and_backbone(repo_id: str, nested_folder: str = NESTED_FOLDER, device: str = DEVICE):
+    """
+    Try several likely subfolder locations for processor/backbone.
+    Returns (processor, backbone).
+    """
+    # candidate subfolders for processor
+    proc_candidates = [
+        "processor",
+        f"{nested_folder}/processor",
+        "",  # no subfolder (root)
+    ]
+    last_exc = None
+    for sub in proc_candidates:
+        try:
+            if sub == "":
+                print(f"Trying AutoImageProcessor.from_pretrained('{repo_id}')")
+                processor = AutoImageProcessor.from_pretrained(repo_id, use_fast=True)
+            else:
+                print(f"Trying AutoImageProcessor.from_pretrained('{repo_id}', subfolder='{sub}')")
+                processor = AutoImageProcessor.from_pretrained(repo_id, subfolder=sub, use_fast=True)
+            # now try backbone with matching guessed subfolder
+            backbone_sub = sub.replace("processor", "vit_backbone") if sub and "processor" in sub else ("vit_backbone" if sub == "" else f"{nested_folder}/vit_backbone")
+            try:
+                print(f"Trying ViTModel.from_pretrained('{repo_id}', subfolder='{backbone_sub}')")
+                backbone = ViTModel.from_pretrained(repo_id, subfolder=backbone_sub)
+            except Exception as e_backbone:
+                # final fallback: try root vit_backbone
+                print(f"Backbone attempt failed for sub='{backbone_sub}': {e_backbone}. Trying root 'vit_backbone'.")
+                backbone = ViTModel.from_pretrained(repo_id, subfolder="vit_backbone")
+            backbone.to(device)
+            backbone.eval()
+            print(f"Loaded processor/backbone from subfolder='{sub or 'root'}'")
+            return processor, backbone
+        except Exception as e:
+            print(f"Processor load failed for sub='{sub}': {e}")
+            last_exc = e
+    # ultimate fallback: official ViT from hub
+    warnings.warn("Could not load processor/backbone from repo; falling back to official 'google/vit-base-patch16-224'.")
+    processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224", use_fast=True)
+    backbone = ViTModel.from_pretrained("google/vit-base-patch16-224")
+    backbone.to(device)
+    backbone.eval()
+    return processor, backbone
+
+# ----------------- Load assets (robust) -----------------
+processor, backbone = load_processor_and_backbone(REPO_ID, nested_folder=NESTED_FOLDER, device=DEVICE)
+
+# Download sklearn artifacts (try root then nested)
+scaler_path = hf_download_try(REPO_ID, "scaler.joblib", nested_folder=NESTED_FOLDER)
+clf_path    = hf_download_try(REPO_ID, "logistic_model.joblib", nested_folder=NESTED_FOLDER)
+metadata_path = hf_download_try(REPO_ID, "metadata.json", nested_folder=NESTED_FOLDER)
 
 scaler = joblib.load(scaler_path)
 clf = joblib.load(clf_path)
 metadata = json.loads(Path(metadata_path).read_text(encoding="utf-8"))
 classes = metadata["classes"]
 
+# ----------------- Prediction function -----------------
 def predict(image_path: str):
     """Predict ripeness condition for a single image."""
     img = Image.open(image_path).convert("RGB")
@@ -103,20 +168,27 @@ def predict(image_path: str):
         feat = pooled.cpu().numpy()
 
     feat_scaled = scaler.transform(feat)
-    probs = clf.predict_proba(feat_scaled)[0]
+    # get probabilities (works for sklearn logistic / classifiers with predict_proba)
+    if hasattr(clf, "predict_proba"):
+        probs = clf.predict_proba(feat_scaled)[0]
+    else:
+        # fallback for classifiers without predict_proba
+        dec = clf.decision_function(feat_scaled)[0]
+        exp = np.exp(dec - np.max(dec))
+        probs = exp / exp.sum()
+
     idx = int(np.argmax(probs))
-    
-    return classes[idx], float(probs[idx]), {
-        classes[i]: float(probs[i]) for i in range(len(classes))
-    }
+    return classes[idx], float(probs[idx]), {classes[i]: float(probs[i]) for i in range(len(classes))}
 
-# Example usage
+# ----------------- Example usage -----------------
 if __name__ == "__main__":
-    label, prob, all_probs = predict("my_apple.jpg")
+    sample_image = "my_apple.jpg"  # change as needed
+    label, prob, all_probs = predict(sample_image)
     print(f"Prediction: {label} ({prob*100:.2f}%)")
-    print("\nTop 5 probabilities:")
-    for cls, p in sorted(all_probs.items(), key=lambda x: -x[1])[:5]:
+    print("\nTop probabilities:")
+    for cls, p in sorted(all_probs.items(), key=lambda x: -x[1])[:TOP_K]:
         print(f"  {cls}: {p*100:.2f}%")
+
 ```
 
 ### Batch Prediction