cattle-breed-classifier / example_usage.py

first commit

abca75a verified about 1 year ago

5.86 kB

"""
Hugging Face Model Hub Integration Example
==========================================
This script demonstrates how to use the cattle breed classification model
from Hugging Face Model Hub.
"""

import onnxruntime as ort
import numpy as np
import json
from PIL import Image
from torchvision import transforms
import requests
from huggingface_hub import hf_hub_download
import os

class CattleBreedClassifier:
    def __init__(self, model_name="your-username/cattle-breed-classifier"):
        """
        Initialize the classifier by downloading model files from Hugging Face
        
        Args:
            model_name: HuggingFace model repository name
        """
        self.model_name = model_name
        self.session = None
        self.prototypes = None
        self.metadata = None
        
        # Download and load model files
        self._download_model_files()
        self._load_model()
        self._load_prototypes()
        
    def _download_model_files(self):
        """Download model files from Hugging Face Hub"""
        print("📥 Downloading model files from Hugging Face...")
        
        # Download ONNX model
        self.model_path = hf_hub_download(
            repo_id=self.model_name,
            filename="model.onnx"
        )
        
        # Download prototypes
        self.prototypes_path = hf_hub_download(
            repo_id=self.model_name,
            filename="prototypes.json"
        )
        
        # Download metadata
        self.metadata_path = hf_hub_download(
            repo_id=self.model_name,
            filename="metadata.json"
        )
        
        print("✅ Model files downloaded successfully!")
        
    def _load_model(self):
        """Load the ONNX model"""
        self.session = ort.InferenceSession(self.model_path)
        print("✅ ONNX model loaded")
        
    def _load_prototypes(self):
        """Load breed prototypes"""
        with open(self.prototypes_path, 'r') as f:
            self.prototypes = json.load(f)
        
        with open(self.metadata_path, 'r') as f:
            self.metadata = json.load(f)
            
        print(f"✅ Loaded prototypes for {len(self.prototypes['prototypes'])} breeds")
        
    def preprocess_image(self, image_input):
        """
        Preprocess image for model inference
        
        Args:
            image_input: PIL Image, numpy array, or file path
            
        Returns:
            numpy.ndarray: Preprocessed image tensor
        """
        # Handle different input types
        if isinstance(image_input, str):
            image = Image.open(image_input).convert('RGB')
        elif isinstance(image_input, np.ndarray):
            image = Image.fromarray(image_input).convert('RGB')
        else:
            image = image_input.convert('RGB')
            
        # Apply preprocessing
        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]
            )
        ])
        
        tensor = transform(image).unsqueeze(0)
        return tensor.numpy()
    
    def predict(self, image_input, return_all_scores=False):
        """
        Predict cattle/buffalo breed from image
        
        Args:
            image_input: Image input (PIL Image, numpy array, or file path)
            return_all_scores: Whether to return scores for all breeds
            
        Returns:
            dict: Prediction results
        """
        # Preprocess image
        input_data = self.preprocess_image(image_input)
        
        # Run inference
        features = self.session.run(None, {'input': input_data})[0][0]
        
        # Calculate similarities with all breed prototypes
        similarities = {}
        for breed, prototype in self.prototypes['prototypes'].items():
            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]
        
        # Determine animal type
        buffalo_breeds = ['Bhadawari', 'Jaffarbadi', 'Mehsana', 'Murrah', 'Surti']
        animal_type = 'Buffalo' if predicted_breed in buffalo_breeds else 'Cattle'
        
        result = {
            'predicted_breed': predicted_breed,
            'confidence': confidence,
            'animal_type': animal_type
        }
        
        if return_all_scores:
            result['all_scores'] = similarities
            
        return result

# Example usage
def main():
    # Initialize classifier (will download model from Hugging Face)
    classifier = CattleBreedClassifier("your-username/cattle-breed-classifier")
    
    # Example 1: Predict from local image
    image_path = "path/to/your/image.jpg"
    if os.path.exists(image_path):
        result = classifier.predict(image_path, return_all_scores=True)
        
        print(f"\n🐄 Prediction Results:")
        print(f"Animal Type: {result['animal_type']}")
        print(f"Predicted Breed: {result['predicted_breed']}")
        print(f"Confidence: {result['confidence']:.4f}")
        
        print(f"\n📊 All Breed Scores:")
        for breed, score in sorted(result['all_scores'].items(), 
                                 key=lambda x: x[1], reverse=True):
            print(f"  {breed}: {score:.4f}")
    
    # Example 2: Predict from PIL Image
    from PIL import Image
    image = Image.open(image_path)
    result = classifier.predict(image)
    print(f"\nDirect PIL prediction: {result['predicted_breed']} ({result['confidence']:.4f})")

if __name__ == "__main__":
    main()

	"""
	Hugging Face Model Hub Integration Example
	==========================================

	This script demonstrates how to use the cattle breed classification model
	from Hugging Face Model Hub.
	"""

	import onnxruntime as ort
	import numpy as np
	import json
	from PIL import Image
	from torchvision import transforms
	import requests
	from huggingface_hub import hf_hub_download
	import os

	class CattleBreedClassifier:
	def __init__(self, model_name="your-username/cattle-breed-classifier"):
	"""
	Initialize the classifier by downloading model files from Hugging Face

	Args:
	model_name: HuggingFace model repository name
	"""
	self.model_name = model_name
	self.session = None
	self.prototypes = None
	self.metadata = None

	# Download and load model files
	self._download_model_files()
	self._load_model()
	self._load_prototypes()

	def _download_model_files(self):
	"""Download model files from Hugging Face Hub"""
	print("📥 Downloading model files from Hugging Face...")

	# Download ONNX model
	self.model_path = hf_hub_download(
	repo_id=self.model_name,
	filename="model.onnx"
	)

	# Download prototypes
	self.prototypes_path = hf_hub_download(
	repo_id=self.model_name,
	filename="prototypes.json"
	)

	# Download metadata
	self.metadata_path = hf_hub_download(
	repo_id=self.model_name,
	filename="metadata.json"
	)

	print("✅ Model files downloaded successfully!")

	def _load_model(self):
	"""Load the ONNX model"""
	self.session = ort.InferenceSession(self.model_path)
	print("✅ ONNX model loaded")

	def _load_prototypes(self):
	"""Load breed prototypes"""
	with open(self.prototypes_path, 'r') as f:
	self.prototypes = json.load(f)

	with open(self.metadata_path, 'r') as f:
	self.metadata = json.load(f)

	print(f"✅ Loaded prototypes for {len(self.prototypes['prototypes'])} breeds")

	def preprocess_image(self, image_input):
	"""
	Preprocess image for model inference

	Args:
	image_input: PIL Image, numpy array, or file path

	Returns:
	numpy.ndarray: Preprocessed image tensor
	"""
	# Handle different input types
	if isinstance(image_input, str):
	image = Image.open(image_input).convert('RGB')
	elif isinstance(image_input, np.ndarray):
	image = Image.fromarray(image_input).convert('RGB')
	else:
	image = image_input.convert('RGB')

	# Apply preprocessing
	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]
	)
	])

	tensor = transform(image).unsqueeze(0)
	return tensor.numpy()

	def predict(self, image_input, return_all_scores=False):
	"""
	Predict cattle/buffalo breed from image

	Args:
	image_input: Image input (PIL Image, numpy array, or file path)
	return_all_scores: Whether to return scores for all breeds

	Returns:
	dict: Prediction results
	"""
	# Preprocess image
	input_data = self.preprocess_image(image_input)

	# Run inference
	features = self.session.run(None, {'input': input_data})[0][0]

	# Calculate similarities with all breed prototypes
	similarities = {}
	for breed, prototype in self.prototypes['prototypes'].items():
	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]

	# Determine animal type
	buffalo_breeds = ['Bhadawari', 'Jaffarbadi', 'Mehsana', 'Murrah', 'Surti']
	animal_type = 'Buffalo' if predicted_breed in buffalo_breeds else 'Cattle'

	result = {
	'predicted_breed': predicted_breed,
	'confidence': confidence,
	'animal_type': animal_type
	}

	if return_all_scores:
	result['all_scores'] = similarities

	return result

	# Example usage
	def main():
	# Initialize classifier (will download model from Hugging Face)
	classifier = CattleBreedClassifier("your-username/cattle-breed-classifier")

	# Example 1: Predict from local image
	image_path = "path/to/your/image.jpg"
	if os.path.exists(image_path):
	result = classifier.predict(image_path, return_all_scores=True)

	print(f"\n🐄 Prediction Results:")
	print(f"Animal Type: {result['animal_type']}")
	print(f"Predicted Breed: {result['predicted_breed']}")
	print(f"Confidence: {result['confidence']:.4f}")

	print(f"\n📊 All Breed Scores:")
	for breed, score in sorted(result['all_scores'].items(),
	key=lambda x: x[1], reverse=True):
	print(f" {breed}: {score:.4f}")

	# Example 2: Predict from PIL Image
	from PIL import Image
	image = Image.open(image_path)
	result = classifier.predict(image)
	print(f"\nDirect PIL prediction: {result['predicted_breed']} ({result['confidence']:.4f})")

	if __name__ == "__main__":
	main()