Update README.md

README.md

@@ -46,58 +46,55 @@ This model classifies cattle and buffalo breeds using computer vision. It's base
 ### Using ONNX Runtime
 
 ```python
+from huggingface_hub import hf_hub_download
 import onnxruntime as ort
 import numpy as np
-from PIL import Image
 import json
-
-# Load model
-session = ort.InferenceSession('model.onnx')
-
-# Load breed prototypes
-with open('prototypes.json', 'r') as f:
-    prototypes = json.load(f)
-
-# Preprocess image
-def preprocess_image(image_path):
-    from torchvision import transforms
-    
+from PIL import Image
+from torchvision import transforms
+
+def setup_model():
+    print("📥 Downloading model from Hugging Face...")
+    model_path = hf_hub_download("vishnuamar/cattle-breed-classifier", "model.onnx")
+    prototypes_path = hf_hub_download("vishnuamar/cattle-breed-classifier", "prototypes.json")
+    session = ort.InferenceSession(model_path)
+    with open(prototypes_path, 'r') as f:
+        prototypes = json.load(f)
+    print("✅ Model ready!")
+    return session, prototypes
+
+def predict_breed(session, prototypes, image_path):
+    image = Image.open(image_path).convert('RGB')
     transform = transforms.Compose([
         transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.ToTensor(),
-        transforms.Normalize(
-            mean=[0.485, 0.456, 0.406],
-            std=[0.229, 0.224, 0.225]
-        )
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
-    
-    image = Image.open(image_path).convert('RGB')
-    tensor = transform(image).unsqueeze(0)
-    return tensor.numpy()
-
-# Predict breed
-def predict_breed(image_path):
-    # Get features
-    input_data = preprocess_image(image_path)
+    input_data = transform(image).unsqueeze(0).numpy()
     features = session.run(None, {'input': input_data})[0][0]
-    
-    # Calculate similarities
     similarities = {}
     for breed, prototype in prototypes['prototypes'].items():
-        similarity = np.dot(features, prototype)
+        similarity = np.dot(features, np.array(prototype))
         similarities[breed] = float(similarity)
-    
-    # Get top prediction
     predicted_breed = max(similarities, key=similarities.get)
     confidence = similarities[predicted_breed]
-    
-    return predicted_breed, confidence, similarities
-
-# Example usage
-breed, confidence, all_scores = predict_breed('path/to/image.jpg')
-print(f"Predicted breed: {breed}")
-print(f"Confidence: {confidence:.4f}")
+    buffalo_breeds = ['Bhadawari', 'Jaffarbadi', 'Mehsana', 'Murrah', 'Surti']
+    animal_type = 'Buffalo' if predicted_breed in buffalo_breeds else 'Cattle'
+    return {
+        'breed': predicted_breed,
+        'confidence': confidence,
+        'animal_type': animal_type,
+        'all_scores': similarities
+    }
+
+if __name__ == "__main__":
+    session, prototypes = setup_model()
+    image_path = "path/to/your/image.jpg"  # Change this to your test image
+    result = predict_breed(session, prototypes, image_path)
+    print(f"\n🐄 Animal: {result['animal_type']}")
+    print(f"Breed: {result['breed']}")
+    print(f"Confidence: {result['confidence']:.4f}")
 ```
 
 ### Integration with Mobile Apps