Update app.py

app.py

@@ -9,7 +9,8 @@ import matplotlib.pyplot as plt
 import numpy as np
 from sklearn.preprocessing import LabelEncoder
 from huggingface_hub import hf_hub_download
-from transformers import pipeline
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import torch
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import re
@@ -31,14 +32,28 @@ def load_image(image_file):
 
 def classify_image(image):
     try:
-        # Create a pipeline for image classification
-        classifier = pipeline('image-classification', model="dima806/car_models_image_detection", device=-1)  # Use -1 for CPU, or 0 for GPU if available
-        
-        # Classify the image
-        results = classifier(image)
+        # Load the model and feature extractor
+        model_name = "dima806/car_models_image_detection"
+        feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+        model = AutoModelForImageClassification.from_pretrained(model_name)
+
+        # Preprocess the image
+        inputs = feature_extractor(images=image, return_tensors="pt")
+
+        # Perform inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # Get the predicted class
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
         
-        # Return top 5 predictions
-        return results[:5]
+        # Get the class label and score
+        predicted_class_label = model.config.id2label[predicted_class_idx]
+        score = torch.nn.functional.softmax(logits, dim=-1)[0, predicted_class_idx].item()
+
+        # Return the top prediction
+        return [{'label': predicted_class_label, 'score': score}]
     
     except Exception as e:
         st.error(f"Classification error: {e}")